Skip to content

Commit

Permalink
fix(os-windows): enumerate trees in order
Browse files Browse the repository at this point in the history
  • Loading branch information
wbenny committed Nov 13, 2024
1 parent 5bb796f commit 30f4875
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 18 deletions.
48 changes: 31 additions & 17 deletions crates/vmi-os-windows/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -610,19 +610,15 @@ where

let left = Va(balanced_node.read(MMADDRESS_NODE.LeftChild)?);
if !left.is_null() {
if !callback(left) {
return Ok(());
}

self.enumerate_tree_node_v1(vmi, registers, left, callback, offsets)?;
}

if !callback(node) {
return Ok(());
}

let right = Va(balanced_node.read(MMADDRESS_NODE.RightChild)?);
if !right.is_null() {
if !callback(right) {
return Ok(());
}

self.enumerate_tree_node_v1(vmi, registers, right, callback, offsets)?;
}

Expand All @@ -648,19 +644,15 @@ where

let left = Va(balanced_node.read(RTL_BALANCED_NODE.Left)?);
if !left.is_null() {
if !callback(left) {
return Ok(());
}

self.enumerate_tree_node_v2(vmi, registers, left, callback, offsets)?;
}

if !callback(node) {
return Ok(());
}

let right = Va(balanced_node.read(RTL_BALANCED_NODE.Right)?);
if !right.is_null() {
if !callback(right) {
return Ok(());
}

self.enumerate_tree_node_v2(vmi, registers, right, callback, offsets)?;
}

Expand All @@ -675,6 +667,29 @@ where
mut callback: impl FnMut(Va) -> bool,
offsets: &v1::Offsets,
) -> Result<(), VmiError> {
let MM_AVL_TABLE = &offsets._MM_AVL_TABLE;
let MMADDRESS_NODE = &offsets._MMADDRESS_NODE;

// NumberGenericTableElements is a ULONG_PTR, which is the same size
// as a pointer.
let count = vmi.read_va(
registers.address_context(root + MM_AVL_TABLE.NumberGenericTableElements.offset),
registers.address_width(),
)?;

let count = MM_AVL_TABLE.NumberGenericTableElements.value_from(count.0);
if count == 0 {
return Ok(());
}

// Table->BalancedRoot.RightChild
let root = vmi.read_va(
registers.address_context(
root + MM_AVL_TABLE.BalancedRoot.offset + MMADDRESS_NODE.RightChild.offset,
),
registers.address_width(),
)?;

self.enumerate_tree_node_v1(vmi, registers, root, &mut callback, offsets)
}

Expand All @@ -686,7 +701,6 @@ where
mut callback: impl FnMut(Va) -> bool,
offsets: &v2::Offsets,
) -> Result<(), VmiError> {
callback(root);
self.enumerate_tree_node_v2(vmi, registers, root, &mut callback, offsets)
}

Expand Down
3 changes: 2 additions & 1 deletion crates/vmi-os-windows/src/offsets/v1.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use isr_macros::{offsets, Field};
use isr_macros::{offsets, Bitfield, Field};

offsets! {
/// Windows 7 kernel offsets used by the [`WindowsOs`] implementation.
Expand Down Expand Up @@ -28,6 +28,7 @@ offsets! {

struct _MM_AVL_TABLE {
BalancedRoot: Field, // _MMADDRESS_NODE
NumberGenericTableElements: Bitfield, // ULONG_PTR
NodeHint: Field, // PVOID
}

Expand Down

0 comments on commit 30f4875

Please sign in to comment.