vmi-os-windows 0.7.0

use once_cell::unsync::OnceCell;
use vmi_core::{Registers as _, Va, VmiError, VmiState, VmiVa, driver::VmiRead};

use super::WindowsHandleTableEntry;
use crate::{ArchAdapter, HandleTableEntryIterator, OffsetsExt, WindowsOs, offset};

/// A Windows handle table.
///
/// A handle table in Windows tracks handles to kernel objects
/// for a specific process, allowing access control and management.
///
/// # Implementation Details
///
/// Corresponds to `_HANDLE_TABLE`.
pub struct WindowsHandleTable<'a, Driver>
where
    Driver: VmiRead,
    Driver::Architecture: ArchAdapter<Driver>,
{
    /// The VMI state.
    vmi: VmiState<'a, WindowsOs<Driver>>,

    /// Address of the handle table.
    va: Va,

    /// Corresponds to `_HANDLE_TABLE.TableCode`.
    table_code: OnceCell<u64>,

    /// Corresponds to `_HANDLE_TABLE.NextHandleNeedingPool`.
    next_handle_needing_pool: OnceCell<u64>,
}

impl<Driver> VmiVa for WindowsHandleTable<'_, Driver>
where
    Driver: VmiRead,
    Driver::Architecture: ArchAdapter<Driver>,
{
    fn va(&self) -> Va {
        self.va
    }
}

impl<'a, Driver> WindowsHandleTable<'a, Driver>
where
    Driver: VmiRead,
    Driver::Architecture: ArchAdapter<Driver>,
{
    /// Creates a new Windows module object.
    pub fn new(vmi: VmiState<'a, WindowsOs<Driver>>, va: Va) -> Self {
        Self {
            vmi,
            va,
            table_code: OnceCell::new(),
            next_handle_needing_pool: OnceCell::new(),
        }
    }

    /// Returns the table code of the handle table.
    ///
    /// # Notes
    ///
    /// This value is cached after the first read.
    ///
    /// # Implementation Details
    ///
    /// Corresponds to `_HANDLE_TABLE.TableCode`.
    pub fn table_code(&self) -> Result<u64, VmiError> {
        self.table_code
            .get_or_try_init(|| {
                let HANDLE_TABLE = offset!(self.vmi, _HANDLE_TABLE);

                self.vmi.read_field(self.va, &HANDLE_TABLE.TableCode)
            })
            .copied()
    }

    /// Returns the next handle needing pool.
    ///
    /// This value tracks the next handle slot that requires additional pool
    /// allocation.
    ///
    /// # Notes
    ///
    /// This value is cached after the first read.
    ///
    /// # Implementation Details
    ///
    /// Corresponds to `_HANDLE_TABLE.NextHandleNeedingPool`.
    pub fn next_handle_needing_pool(&self) -> Result<u64, VmiError> {
        self.next_handle_needing_pool
            .get_or_try_init(|| {
                let HANDLE_TABLE = offset!(self.vmi, _HANDLE_TABLE);

                self.vmi
                    .read_field(self.va, &HANDLE_TABLE.NextHandleNeedingPool)
            })
            .copied()
    }

    /// Iterates over all handle table entries.
    ///
    /// Returns an iterator over all handle table entries that have a valid
    /// object pointer. The iterator yields a tuple containing the handle
    /// value and the handle table entry.
    ///
    /// # Implementation Details
    ///
    /// The functionality is similar to the Windows kernel's internal
    /// `ExpSnapShotHandleTables()` function.
    pub fn iter(
        &self,
    ) -> Result<
        impl Iterator<Item = Result<(u64, WindowsHandleTableEntry<'a, Driver>), VmiError>>
        + use<'a, Driver>,
        VmiError,
    > {
        Ok(HandleTableEntryIterator::new(
            self.vmi,
            self.table_code()?,
            self.next_handle_needing_pool()?,
        ))
    }

    /// Performs a lookup in the handle table to find the address of a handle
    /// table entry.
    ///
    /// Implements the multi-level handle table lookup algorithm used by
    /// Windows. Returns the virtual address of the handle table entry.
    ///
    /// # Implementation Details
    ///
    /// The functionality is similar to the Windows kernel's internal
    /// `ExpLookupHandleTableEntry()` function.
    pub fn lookup(
        &self,
        handle: u64,
    ) -> Result<Option<WindowsHandleTableEntry<'a, Driver>>, VmiError> {
        lookup_handle_entry(
            self.vmi,
            self.table_code()?,
            self.next_handle_needing_pool()?,
            handle,
        )
    }
}

/// Looks up a handle table entry given snapshots of `_HANDLE_TABLE.TableCode`
/// and `_HANDLE_TABLE.NextHandleNeedingPool`. Factored out so both
/// [`WindowsHandleTable::lookup`] and [`HandleTableEntryIterator`] can share
/// the multi-level lookup algorithm without the iterator having to borrow
/// from [`WindowsHandleTable`].
pub(crate) fn lookup_handle_entry<'a, Driver>(
    vmi: VmiState<'a, WindowsOs<Driver>>,
    table_code: u64,
    next_handle_needing_pool: u64,
    handle: u64,
) -> Result<Option<WindowsHandleTableEntry<'a, Driver>>, VmiError>
where
    Driver: VmiRead,
    Driver::Architecture: ArchAdapter<Driver>,
{
    const PAGE_SIZE: u64 = 4096;
    const LEVEL_CODE_MASK: u64 = 3;
    const HANDLE_VALUE_INC: u64 = 4;

    let sizeof_handle_table_entry = match vmi.underlying_os().offsets.ext() {
        Some(OffsetsExt::V1(offsets)) => offsets._HANDLE_TABLE_ENTRY.len() as u64,
        Some(OffsetsExt::V2(offsets)) => offsets._HANDLE_TABLE_ENTRY.len() as u64,
        None => unimplemented!(),
    };

    let address_width = vmi.registers().address_width() as u64;
    let lowlevel_count = PAGE_SIZE / sizeof_handle_table_entry; // TABLE_PAGE_SIZE / sizeof(HANDLE_TABLE_ENTRY)
    let midlevel_count = PAGE_SIZE / address_width; // PAGE_SIZE / sizeof(PHANDLE_TABLE_ENTRY)

    // The 2 least significant bits of a handle are available to the
    // application and are ignored by the system.
    let mut index = handle & !0b11;

    // See if this can be a valid handle given the table levels.
    if index >= next_handle_needing_pool {
        return Ok(None);
    }

    let level = table_code & LEVEL_CODE_MASK;
    let table = Va(table_code - level);

    let entry = match level {
        0 => table + index * (sizeof_handle_table_entry / HANDLE_VALUE_INC),

        1 => {
            let table2 = table;
            let i = index % (lowlevel_count * HANDLE_VALUE_INC);

            index -= i;
            let j = index / (lowlevel_count * HANDLE_VALUE_INC);

            let table1 = vmi.read_va_native(table2 + j * address_width)?;

            table1 + i * (sizeof_handle_table_entry / HANDLE_VALUE_INC)
        }

        2 => {
            let table3 = table;
            let i = index % (lowlevel_count * HANDLE_VALUE_INC);

            index -= i;
            let mut k = index / (lowlevel_count * HANDLE_VALUE_INC);

            let j = k % midlevel_count;
            k -= j;
            k /= midlevel_count;

            let table2 = vmi.read_va_native(table3 + k * address_width)?;
            let table1 = vmi.read_va_native(table2 + j * address_width)?;

            table1 + i * (sizeof_handle_table_entry / HANDLE_VALUE_INC)
        }

        _ => unreachable!(),
    };

    Ok(Some(WindowsHandleTableEntry::new(vmi, entry)))
}