Code for Episode 10

scanflow/src/disasm.rs

@@ -11,18 +11,23 @@ use std::collections::BTreeMap;
 #[derive(Default)]
 pub struct Disasm {
     map: BTreeMap<Address, Address>,
+    inverse_map: BTreeMap<Address, Vec<Address>>,
+    globals: Vec<Address>,
 }
 
 impl Disasm {
     #[allow(unused)]
     pub fn reset(&mut self) {
         self.map.clear();
+        self.inverse_map.clear();
+        self.globals.clear();
     }
 
     pub fn collect_globals(
         &mut self,
         process: &mut Win32Process<impl VirtualMemory>,
     ) -> Result<()> {
+        self.reset();
         let modules = process.module_list()?;
 
         let mut image = vec![0; size::kb(128)];
@@ -81,12 +86,24 @@ impl Disasm {
             }
         }
 
-        println!("GLOBALS {:x}", self.map.len());
+        for (&k, &v) in &self.map {
+            self.inverse_map.entry(v).or_default().push(k);
+        }
+
+        self.globals = self.inverse_map.keys().copied().collect();
 
         Ok(())
     }
 
     pub fn map(&self) -> &BTreeMap<Address, Address> {
         &self.map
     }
+
+    pub fn inverse_map(&self) -> &BTreeMap<Address, Vec<Address>> {
+        &self.inverse_map
+    }
+
+    pub fn globals(&self) -> &Vec<Address> {
+        &self.globals
+    }
 }

scanflow/src/main.rs

@@ -7,16 +7,16 @@ use memflow_win32::{Error, Result};
 
 use simplelog::{Config, LevelFilter, TermLogger, TerminalMode};
 
-use std::collections::BTreeSet;
 use std::convert::TryInto;
+use std::time::Instant;
 
-mod value_scanner;
+pub mod value_scanner;
 use value_scanner::ValueScanner;
 
-mod pointer_map;
+pub mod pointer_map;
 use pointer_map::PointerMap;
 
-mod disasm;
+pub mod disasm;
 use disasm::Disasm;
 
 #[macro_use]
@@ -97,35 +97,50 @@ fn main() -> Result<()> {
                         )?;
                     }
 
+                    let start = Instant::now();
+
                     let matches = if use_di == "y" {
                         if disasm.map().is_empty() {
                             disasm.collect_globals(&mut process)?;
                         }
-                        let set: BTreeSet<_> = disasm.map().values().copied().collect();
                         pointer_map.find_matches_addrs(
-                            lrange,
-                            urange,
+                            (lrange, urange),
                             max_depth,
-                            value_scanner.matches(),
-                            set.into_iter(),
+                            value_scanner.matches().iter().copied(),
+                            disasm.globals(),
                         )
                     } else {
-                        pointer_map.find_matches(lrange, urange, max_depth, value_scanner.matches())
+                        pointer_map.find_matches(
+                            (lrange, urange),
+                            max_depth,
+                            value_scanner.matches().iter().copied(),
+                        )
                     };
 
-                    println!("Matches found: {}", matches.len());
+                    println!(
+                        "Matches found: {} in {:.2}ms",
+                        matches.len(),
+                        start.elapsed().as_secs_f64() * 1000.0
+                    );
 
-                    for (m, offsets) in matches.into_iter().filter(|(_, v)| {
-                        if let Some(a) = filter_addr {
-                            if let Some((s, _)) = v.first() {
-                                s.as_u64() == a
+                    if matches.len() > 64 {
+                        println!("Printing first 64 matches");
+                    }
+                    for (m, offsets) in matches
+                        .into_iter()
+                        .filter(|(_, v)| {
+                            if let Some(a) = filter_addr {
+                                if let Some((s, _)) = v.first() {
+                                    s.as_u64() == a
+                                } else {
+                                    false
+                                }
                             } else {
-                                false
+                                true
                             }
-                        } else {
-                            true
-                        }
-                    }) {
+                        })
+                        .take(64)
+                    {
                         for (start, off) in offsets.into_iter() {
                             print!("{:x} + ({}) => ", start, off);
                         }

scanflow/src/pointer_map.rs

@@ -8,15 +8,21 @@ use std::ops::Bound::Included;
 #[derive(Default)]
 pub struct PointerMap {
     map: BTreeMap<Address, Address>,
+    inverse_map: BTreeMap<Address, Vec<Address>>,
+    pointers: Vec<Address>,
 }
 
 impl PointerMap {
     #[allow(unused)]
     pub fn reset(&mut self) {
-        self.map.clear()
+        self.map.clear();
+        self.inverse_map.clear();
+        self.pointers.clear();
     }
 
     pub fn create_map<T: VirtualMemory>(&mut self, mem: &mut T, size_addr: usize) -> Result<()> {
+        self.reset();
+
         let mem_map = mem.virt_page_map_range(size::mb(16), Address::null(), (1u64 << 47).into());
 
         let mut buf = vec![0; 0x1000 + size_addr - 1];
@@ -49,7 +55,11 @@ impl PointerMap {
             }
         }
 
-        println!("Map size: {:x}", self.map.len());
+        for (&k, &v) in &self.map {
+            self.inverse_map.entry(v).or_default().push(k);
+        }
+
+        self.pointers = self.map.keys().copied().collect();
 
         Ok(())
     }
@@ -58,84 +68,108 @@ impl PointerMap {
         &self.map
     }
 
+    pub fn inverse_map(&self) -> &BTreeMap<Address, Vec<Address>> {
+        &self.inverse_map
+    }
+
+    pub fn pointers(&self) -> &Vec<Address> {
+        &self.pointers
+    }
+
     fn walk_down_range(
         &self,
         addr: Address,
-        lrange: usize,
-        urange: usize,
+        (lrange, urange): (usize, usize),
         max_levels: usize,
         level: usize,
-        search_for: &[Address],
+        startpoints: &[Address],
         out: &mut Vec<(Address, Vec<(Address, isize)>)>,
-        tmp: &mut Vec<(Address, isize)>,
+        (final_addr, tmp): (Address, &mut Vec<(Address, isize)>),
     ) {
-        let min = Address::from(addr.as_u64().saturating_sub(lrange as _));
-        let max = Address::from(addr.as_u64().saturating_add(urange as _));
+        let min = Address::from(addr.as_u64().saturating_sub(urange as _));
+        let max = Address::from(addr.as_u64().saturating_add(lrange as _));
 
-        for &m in search_for.iter().filter(|&&m| m >= min && m <= max) {
-            let off = signed_diff(m, addr);
+        // Find the lower bound
+        let idx = startpoints.binary_search(&min).unwrap_or_else(|x| x);
+
+        let mut iter = startpoints
+            .iter()
+            .skip(idx)
+            .copied()
+            .take_while(|&v| v <= max);
+
+        // Pick next match
+        let mut m = iter.next();
+
+        // Go through the rest
+        for e in iter {
+            let off = signed_diff(addr, e).abs();
+            // If abs offset is smaller, overwrite
+            // < biasses more towards positive end
+            if off < signed_diff(addr, m.unwrap()).abs() {
+                m = Some(e);
+            }
+        }
+
+        // Push match if found
+        if let Some(e) = m {
+            let off = signed_diff(addr, e);
             let mut cloned = tmp.clone();
-            cloned.push((addr, off));
-            out.push((m, cloned));
+            cloned.push((e, off));
+            cloned.reverse();
+            out.push((final_addr, cloned));
         }
 
+        // Recurse downwards if possible
         if level < max_levels {
-            for (&k, &v) in self.map.range((Included(&min), Included(&max))) {
-                let off = signed_diff(k, addr);
-                tmp.push((addr, off));
-                self.walk_down_range(
-                    v,
-                    lrange,
-                    urange,
-                    max_levels,
-                    level + 1,
-                    search_for,
-                    out,
-                    tmp,
-                );
+            for (&k, vec) in self.inverse_map.range((Included(&min), Included(&max))) {
+                let off = signed_diff(addr, k);
+                tmp.push((k, off));
+                for &v in vec {
+                    self.walk_down_range(
+                        v,
+                        (lrange, urange),
+                        max_levels,
+                        level + 1,
+                        startpoints,
+                        out,
+                        (final_addr, tmp),
+                    );
+                }
                 tmp.pop();
             }
         }
     }
 
     pub fn find_matches_addrs(
         &self,
-        lrange: usize,
-        urange: usize,
+        range: (usize, usize),
         max_depth: usize,
-        search_for: &[Address],
-        addrs: impl Iterator<Item = Address>,
+        search_for: impl Iterator<Item = Address>,
+        entry_points: &[Address],
     ) -> Vec<(Address, Vec<(Address, isize)>)> {
         let mut matches = vec![];
-        for k in addrs {
+        for m in search_for {
             self.walk_down_range(
-                k,
-                lrange,
-                urange,
+                m,
+                range,
                 max_depth,
                 1,
-                search_for,
+                entry_points,
                 &mut matches,
-                &mut vec![],
+                (m, &mut vec![]),
             );
         }
         matches
     }
 
     pub fn find_matches(
         &self,
-        lrange: usize,
-        urange: usize,
+        range: (usize, usize),
         max_depth: usize,
-        search_for: &[Address],
+        search_for: impl Iterator<Item = Address>,
     ) -> Vec<(Address, Vec<(Address, isize)>)> {
-        self.find_matches_addrs(
-            lrange,
-            urange,
-            max_depth,
-            search_for,
-            self.map.keys().copied(),
-        )
+        self.find_matches_addrs(range, max_depth, search_for, &self.pointers)
     }
 }