stackwise-core 0.2.0

Artifact-first Rust stack analysis engine used by Stackwise
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use object::{Object, ObjectSection};

#[derive(Debug, Clone, Copy)]
pub struct PeUnwindRecord {
    pub begin: u64,
    pub end: u64,
    pub stack_bytes: u64,
}

pub fn parse_pe_x64_unwind<'data>(file: &object::File<'data>) -> Vec<PeUnwindRecord> {
    let Some(pdata) = file.sections().find(|section| {
        section
            .name()
            .map(|name| name.trim_end_matches('\0') == ".pdata")
            .unwrap_or(false)
    }) else {
        return Vec::new();
    };

    let Ok(pdata_bytes) = pdata.data() else {
        return Vec::new();
    };

    let image_base = file.relative_address_base();
    let sections = file
        .sections()
        .filter_map(|section| {
            let data = section.data().ok()?;
            Some(SectionRange {
                address: section.address(),
                data,
            })
        })
        .collect::<Vec<_>>();
    let mut frames = Vec::new();

    for chunk in pdata_bytes.chunks_exact(12) {
        let begin = image_base + read_u32(chunk, 0);
        let end = image_base + read_u32(chunk, 4);
        let unwind_rva = image_base + read_u32(chunk, 8);

        if let Some(unwind_bytes) = section_tail(unwind_rva, &sections) {
            let Some(bytes) = parse_unwind_info(unwind_bytes) else {
                continue;
            };
            frames.push(PeUnwindRecord {
                begin,
                end,
                stack_bytes: bytes,
            });
        }
    }

    frames
}

pub fn infer_pe_image_base<'data>(file: &object::File<'data>) -> Option<u64> {
    Some(file.relative_address_base())
}

fn section_tail<'a>(rva: u64, sections: &'a [SectionRange<'a>]) -> Option<&'a [u8]> {
    sections.iter().find_map(|section| {
        let offset = rva.checked_sub(section.address)? as usize;
        section.data.get(offset..)
    })
}

struct SectionRange<'a> {
    address: u64,
    data: &'a [u8],
}

fn parse_unwind_info(bytes: &[u8]) -> Option<u64> {
    if bytes.len() < 4 {
        return None;
    }

    let count = usize::from(bytes[2]);
    let codes_len = count.checked_mul(2)?;
    if bytes.len() < 4 + codes_len {
        return None;
    }

    let mut stack = 0u64;
    let mut index = 0usize;
    let codes = &bytes[4..4 + codes_len];

    while index < count {
        let slot = index * 2;
        let op_and_info = codes[slot + 1];
        let op = op_and_info & 0x0f;
        let info = op_and_info >> 4;

        match op {
            0 => stack += 8,
            1 if info == 0 => {
                let value = read_code_u16(codes, index + 1)?;
                stack += u64::from(value) * 8;
                index += 1;
            }
            1 => {
                let low = u32::from(read_code_u16(codes, index + 1)?);
                let high = u32::from(read_code_u16(codes, index + 2)?);
                stack += u64::from((high << 16) | low);
                index += 2;
            }
            2 => stack += u64::from(info) * 8 + 8,
            3 => {}
            4 | 8 => index += 1,
            5 | 9 => index += 2,
            10 => stack += if info == 0 { 40 } else { 48 },
            _ => {}
        }

        index += 1;
    }

    Some(stack)
}

fn read_code_u16(codes: &[u8], index: usize) -> Option<u16> {
    let offset = index.checked_mul(2)?;
    let bytes = codes.get(offset..offset + 2)?;
    Some(u16::from_le_bytes([bytes[0], bytes[1]]))
}

fn read_u32(bytes: &[u8], offset: usize) -> u64 {
    u64::from(u32::from_le_bytes([
        bytes[offset],
        bytes[offset + 1],
        bytes[offset + 2],
        bytes[offset + 3],
    ]))
}

#[cfg(test)]
mod tests {
    use super::parse_unwind_info;

    #[test]
    fn parses_small_alloc_and_push() {
        let bytes = [0x01, 0x04, 0x02, 0x00, 0x01, 0x50, 0x02, 0x32];
        assert_eq!(parse_unwind_info(&bytes), Some(40));
    }
}