lancelot 0.2.1

binary analysis framework for x32/x64 PE files
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

use std::collections::BTreeMap;

use anyhow::Result;
use log::debug;

use crate::{
    analysis::{cfg, dis},
    aspace::AddressSpace,
    module::Module,
    RVA, VA,
};

#[derive(Default)]
pub struct CallGraph {
    // call instruction indexes...
    /// map from function start to the addresses that call here.
    /// lookup via `call_instruction_functions` to figure out the functions that
    /// call here.
    pub calls_to:   BTreeMap<VA, Vec<VA>>,
    /// map from an instruction to the addresses that it calls (usually one).
    pub calls_from: BTreeMap<VA, Vec<VA>>,

    // instruction membership indexes...
    /// map from function start to the instructions in its CFG that call
    /// elsewhere. lookup via `calls_to` to figoure out the functions that
    /// this function calls.
    pub function_call_instructions: BTreeMap<VA, Vec<VA>>,
    /// map from instruction to starts of functions whose CFGs contain the
    /// instruction (usually one).
    pub call_instruction_functions: BTreeMap<VA, Vec<VA>>,
}

pub fn build_call_graph(module: &Module, cfgs: &BTreeMap<VA, cfg::CFG>) -> Result<CallGraph> {
    debug!("call graph");

    let mut cg: CallGraph = Default::default();
    let decoder = dis::get_disassembler(module)?;

    for (&function, cfg) in cfgs.iter() {
        debug!("call graph: {:#x}", function);

        // ensure there are at least (empty) entries for all the keys in `functions`
        cg.function_call_instructions.entry(function).or_default();
        cg.calls_to.entry(function).or_default();

        for basic_block in cfg.basic_blocks.values() {
            let buf = module
                .address_space
                .read_bytes(basic_block.address, basic_block.length as usize)?;

            for (offset, insn) in dis::linear_disassemble(&decoder, &buf) {
                if let Ok(Some(insn)) = insn {
                    if matches!(insn.mnemonic, zydis::enums::Mnemonic::CALL) {
                        let va = basic_block.address + offset as RVA;
                        for flow in cfg::get_call_insn_flow(module, va, &insn)?.iter() {
                            if let cfg::Flow::Call(target) = *flow {
                                cg.calls_from.entry(va).or_default().push(target);
                                cg.calls_to.entry(target).or_default().push(va);
                                cg.function_call_instructions.entry(function).or_default().push(va);
                                cg.call_instruction_functions.entry(va).or_default().push(function);
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(cg)
}

#[cfg(test)]
mod tests {
    use crate::{
        analysis::{call_graph, cfg::CFG, pe},
        rsrc::*,
        VA,
    };
    use anyhow::Result;
    use std::collections::BTreeMap;

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let mut cfgs: BTreeMap<VA, CFG> = Default::default();
        for &function in pe::find_function_starts(&pe)?.iter() {
            if let Ok(cfg) = crate::analysis::cfg::build_cfg(&pe.module, function) {
                cfgs.insert(function, cfg);
            }
        }

        let cg = call_graph::build_call_graph(&pe.module, &cfgs)?;

        assert_eq!(cg.calls_to[&0x180001068].len(), 2);
        assert!(cg.calls_to[&0x180001068].iter().find(|&&v| v == 0x18000F775).is_some());
        assert!(cg.calls_to[&0x180001068].iter().find(|&&v| v == 0x180060504).is_some());
        assert!(cg.calls_from[&0x180060504]
            .iter()
            .find(|&&v| v == 0x180001068)
            .is_some());
        assert!(cg.call_instruction_functions[&0x180060504]
            .iter()
            .find(|&&v| v == 0x1800602C0)
            .is_some());
        assert!(cg.function_call_instructions[&0x1800602C0]
            .iter()
            .find(|&&v| v == 0x180060504)
            .is_some());

        Ok(())
    }
}