vmi-core 0.7.0

use isr_macros::Field;
use zerocopy::{FromBytes, Immutable, IntoBytes};

use super::session::VmiSession;
use crate::{
    AccessContext, AddressContext, Architecture, Pa, Registers as _, Va, VcpuId, VmiCore, VmiError,
    VmiRead, VmiWrite,
    driver::VmiSetRegisters,
    os::{NoOS, VmiOs},
};

/// A VMI state.
///
/// The state combines access to a [`VmiSession`] with [`Architecture::Registers`]
/// to provide unified access to VMI operations in the context of a specific
/// virtual machine state.
pub struct VmiState<'a, Os>
where
    Os: VmiOs,
{
    /// The VMI session.
    session: VmiSession<'a, Os>,

    /// The CPU registers associated with the current VM state.
    registers: &'a <Os::Architecture as Architecture>::Registers,
}

impl<Os> Clone for VmiState<'_, Os>
where
    Os: VmiOs,
{
    fn clone(&self) -> Self {
        *self
    }
}

impl<Os> Copy for VmiState<'_, Os> where Os: VmiOs {}

impl<'a, Os> std::ops::Deref for VmiState<'a, Os>
where
    Os: VmiOs,
{
    type Target = VmiSession<'a, Os>;

    fn deref(&self) -> &Self::Target {
        &self.session
    }
}

impl<'a, Os> VmiState<'a, Os>
where
    Os: VmiOs,
{
    /// Creates a new VMI state.
    pub fn new(
        session: &'a VmiSession<'a, Os>,
        registers: &'a <Os::Architecture as Architecture>::Registers,
    ) -> Self {
        Self {
            session: *session,
            registers,
        }
    }

    /// Creates a new VMI state with the specified registers.
    pub fn with_registers(
        &'a self,
        registers: &'a <Os::Architecture as Architecture>::Registers,
    ) -> Self {
        Self {
            session: self.session,
            registers,
        }
    }

    /// Creates a new VMI state without an OS-specific implementation.
    pub fn without_os(&self) -> VmiState<'a, NoOS<Os::Driver>> {
        VmiState {
            session: self.session.without_os(),
            registers: self.registers,
        }
    }

    // Note that `core()` and `underlying_os()` are delegated to the `VmiSession`.

    /// Returns the VMI session.
    pub fn session(&self) -> &VmiSession<'a, Os> {
        &self.session
    }

    /// Returns the CPU registers associated with the current event.
    pub fn registers(&self) -> &'a <Os::Architecture as Architecture>::Registers {
        self.registers
    }

    /// Returns a wrapper providing access to OS-specific operations.
    pub fn os(&self) -> VmiOsState<'a, Os> {
        VmiOsState(*self)
    }

    /// Creates an address context for a given virtual address.
    pub fn access_context(&self, address: Va) -> AccessContext {
        self.registers().access_context(address)
    }

    /// Creates an address context for a given virtual address.
    pub fn address_context(&self, address: Va) -> AddressContext {
        self.registers().address_context(address)
    }

    /// Returns the physical address of the root of the current page table
    /// hierarchy for a given virtual address.
    pub fn translation_root(&self, va: Va) -> Pa {
        self.registers().translation_root(va)
    }
}

///////////////////////////////////////////////////////////////////////////////
// VmiRead
///////////////////////////////////////////////////////////////////////////////

impl<'a, Os> VmiState<'a, Os>
where
    Os: VmiOs,
    Os::Driver: VmiRead,
{
    /// Returns the return address from the current stack frame.
    pub fn return_address(&self) -> Result<Va, VmiError> {
        self.registers().return_address(self.core())
    }

    /// Translates a virtual address to a physical address.
    pub fn translate_address(&self, va: Va) -> Result<Pa, VmiError> {
        self.core().translate_address(self.address_context(va))
    }

    // region: Read

    /// Reads memory from the virtual machine.
    pub fn read(&self, address: Va, buffer: &mut [u8]) -> Result<(), VmiError> {
        self.read_in(self.access_context(address), buffer)
    }

    /// Reads a single byte from the virtual machine.
    pub fn read_u8(&self, address: Va) -> Result<u8, VmiError> {
        self.read_u8_in(self.access_context(address))
    }

    /// Reads a 16-bit unsigned integer from the virtual machine.
    pub fn read_u16(&self, address: Va) -> Result<u16, VmiError> {
        self.read_u16_in(self.access_context(address))
    }

    /// Reads a 32-bit unsigned integer from the virtual machine.
    pub fn read_u32(&self, address: Va) -> Result<u32, VmiError> {
        self.read_u32_in(self.access_context(address))
    }

    /// Reads a 64-bit unsigned integer from the virtual machine.
    pub fn read_u64(&self, address: Va) -> Result<u64, VmiError> {
        self.read_u64_in(self.access_context(address))
    }

    /// Reads an unsigned integer of the specified size from the virtual machine.
    ///
    /// This method reads an unsigned integer of the specified size (in bytes)
    /// from the virtual machine. Note that the size must be 1, 2, 4, or 8.
    ///
    /// The result is returned as a [`u64`] to accommodate the widest possible
    /// integer size.
    pub fn read_uint(&self, address: Va, size: usize) -> Result<u64, VmiError> {
        self.read_uint_in(self.access_context(address), size)
    }

    /// Reads a field of a structure from the virtual machine.
    ///
    /// This method reads a field from the virtual machine. The field is
    /// defined by the provided [`Field`] structure, which specifies the
    /// offset and size of the field within the memory region.
    ///
    /// The result is returned as a [`u64`] to accommodate the widest possible
    /// integer size.
    pub fn read_field(&self, base_address: Va, field: &Field) -> Result<u64, VmiError> {
        self.read_field_in(self.access_context(base_address), field)
    }

    /// Reads an address-sized unsigned integer from the virtual machine.
    pub fn read_address(&self, address: Va) -> Result<u64, VmiError> {
        self.read_address_in(self.access_context(address))
    }

    /// Reads an address-sized unsigned integer from the virtual machine.
    pub fn read_address_native(&self, address: Va) -> Result<u64, VmiError> {
        self.read_address_native_in(self.access_context(address))
    }

    /// Reads a 32-bit address from the virtual machine.
    pub fn read_address32(&self, address: Va) -> Result<u64, VmiError> {
        self.read_address32_in(self.access_context(address))
    }

    /// Reads a 64-bit address from the virtual machine.
    pub fn read_address64(&self, address: Va) -> Result<u64, VmiError> {
        self.read_address64_in(self.access_context(address))
    }

    /// Reads a virtual address from the virtual machine.
    pub fn read_va(&self, address: Va) -> Result<Va, VmiError> {
        self.read_va_in(self.access_context(address))
    }

    /// Reads a virtual address from the virtual machine.
    pub fn read_va_native(&self, address: Va) -> Result<Va, VmiError> {
        self.read_va_native_in(self.access_context(address))
    }

    /// Reads a 32-bit virtual address from the virtual machine.
    pub fn read_va32(&self, address: Va) -> Result<Va, VmiError> {
        self.read_va32_in(self.access_context(address))
    }

    /// Reads a 64-bit virtual address from the virtual machine.
    pub fn read_va64(&self, address: Va) -> Result<Va, VmiError> {
        self.read_va64_in(self.access_context(address))
    }

    /// Reads a null-terminated string of bytes from the virtual machine with a
    /// specified limit.
    pub fn read_string_bytes_limited(
        &self,
        address: Va,
        limit: usize,
    ) -> Result<Vec<u8>, VmiError> {
        self.read_string_bytes_limited_in(self.access_context(address), limit)
    }

    /// Reads a null-terminated string of bytes from the virtual machine.
    pub fn read_string_bytes(&self, address: Va) -> Result<Vec<u8>, VmiError> {
        self.read_string_bytes_in(self.access_context(address))
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine
    /// with a specified limit.
    pub fn read_string_utf16_bytes_limited(
        &self,
        address: Va,
        limit: usize,
    ) -> Result<Vec<u16>, VmiError> {
        self.read_string_utf16_bytes_limited_in(self.access_context(address), limit)
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine.
    pub fn read_string_utf16_bytes(&self, address: Va) -> Result<Vec<u16>, VmiError> {
        self.read_string_utf16_bytes_in(self.access_context(address))
    }

    /// Reads a null-terminated string from the virtual machine with a specified
    /// limit.
    pub fn read_string_limited(&self, address: Va, limit: usize) -> Result<String, VmiError> {
        self.read_string_limited_in(self.access_context(address), limit)
    }

    /// Reads a null-terminated string from the virtual machine.
    pub fn read_string(&self, address: Va) -> Result<String, VmiError> {
        self.read_string_in(self.access_context(address))
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine
    /// with a specified limit.
    pub fn read_string_utf16_limited(&self, address: Va, limit: usize) -> Result<String, VmiError> {
        self.read_string_utf16_limited_in(self.access_context(address), limit)
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine.
    pub fn read_string_utf16(&self, address: Va) -> Result<String, VmiError> {
        self.read_string_utf16_in(self.access_context(address))
    }

    /// Reads a struct from the virtual machine.
    pub fn read_struct<T>(&self, address: Va) -> Result<T, VmiError>
    where
        T: IntoBytes + FromBytes,
    {
        self.read_struct_in(self.access_context(address))
    }

    // endregion: Read

    // region: Read in

    /// Reads memory from the virtual machine.
    pub fn read_in(
        &self,
        ctx: impl Into<AccessContext>,
        buffer: &mut [u8],
    ) -> Result<(), VmiError> {
        self.core().read(ctx, buffer)
    }

    /// Reads a single byte from the virtual machine.
    pub fn read_u8_in(&self, ctx: impl Into<AccessContext>) -> Result<u8, VmiError> {
        self.core().read_u8(ctx)
    }

    /// Reads a 16-bit unsigned integer from the virtual machine.
    pub fn read_u16_in(&self, ctx: impl Into<AccessContext>) -> Result<u16, VmiError> {
        self.core().read_u16(ctx)
    }

    /// Reads a 32-bit unsigned integer from the virtual machine.
    pub fn read_u32_in(&self, ctx: impl Into<AccessContext>) -> Result<u32, VmiError> {
        self.core().read_u32(ctx)
    }

    /// Reads a 64-bit unsigned integer from the virtual machine.
    pub fn read_u64_in(&self, ctx: impl Into<AccessContext>) -> Result<u64, VmiError> {
        self.core().read_u64(ctx)
    }

    /// Reads an unsigned integer of the specified size from the virtual machine.
    ///
    /// This method reads an unsigned integer of the specified size (in bytes)
    /// from the virtual machine. Note that the size must be 1, 2, 4, or 8.
    ///
    /// The result is returned as a [`u64`] to accommodate the widest possible
    /// integer size.
    pub fn read_uint_in(
        &self,
        ctx: impl Into<AccessContext>,
        size: usize,
    ) -> Result<u64, VmiError> {
        self.core().read_uint(ctx, size)
    }

    /// Reads a field of a structure from the virtual machine.
    ///
    /// This method reads a field from the virtual machine. The field is
    /// defined by the provided [`Field`] structure, which specifies the
    /// offset and size of the field within the memory region.
    ///
    /// The result is returned as a [`u64`] to accommodate the widest possible
    /// integer size.
    pub fn read_field_in(
        &self,
        ctx: impl Into<AccessContext>,
        field: &Field,
    ) -> Result<u64, VmiError> {
        self.core().read_field(ctx, field)
    }

    /// Reads an address-sized unsigned integer from the virtual machine.
    pub fn read_address_in(&self, ctx: impl Into<AccessContext>) -> Result<u64, VmiError> {
        self.core()
            .read_address(ctx, self.registers().effective_address_width())
    }

    /// Reads an address-sized unsigned integer from the virtual machine.
    pub fn read_address_native_in(&self, ctx: impl Into<AccessContext>) -> Result<u64, VmiError> {
        self.core()
            .read_address(ctx, self.registers().address_width())
    }

    /// Reads a 32-bit address from the virtual machine.
    pub fn read_address32_in(&self, ctx: impl Into<AccessContext>) -> Result<u64, VmiError> {
        self.core().read_address32(ctx)
    }

    /// Reads a 64-bit address from the virtual machine.
    pub fn read_address64_in(&self, ctx: impl Into<AccessContext>) -> Result<u64, VmiError> {
        self.core().read_address64(ctx)
    }

    /// Reads a virtual address from the virtual machine.
    pub fn read_va_in(&self, ctx: impl Into<AccessContext>) -> Result<Va, VmiError> {
        self.core()
            .read_va(ctx, self.registers().effective_address_width())
    }

    /// Reads a virtual address from the virtual machine.
    pub fn read_va_native_in(&self, ctx: impl Into<AccessContext>) -> Result<Va, VmiError> {
        self.core().read_va(ctx, self.registers().address_width())
    }

    /// Reads a 32-bit virtual address from the virtual machine.
    pub fn read_va32_in(&self, ctx: impl Into<AccessContext>) -> Result<Va, VmiError> {
        self.core().read_va32(ctx)
    }

    /// Reads a 64-bit virtual address from the virtual machine.
    pub fn read_va64_in(&self, ctx: impl Into<AccessContext>) -> Result<Va, VmiError> {
        self.core().read_va64(ctx)
    }

    /// Reads a null-terminated string of bytes from the virtual machine with a
    /// specified limit.
    pub fn read_string_bytes_limited_in(
        &self,
        ctx: impl Into<AccessContext>,
        limit: usize,
    ) -> Result<Vec<u8>, VmiError> {
        self.core().read_string_bytes_limited(ctx, limit)
    }

    /// Reads a null-terminated string of bytes from the virtual machine.
    pub fn read_string_bytes_in(&self, ctx: impl Into<AccessContext>) -> Result<Vec<u8>, VmiError> {
        self.core().read_string_bytes(ctx)
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine
    /// with a specified limit.
    pub fn read_string_utf16_bytes_limited_in(
        &self,
        ctx: impl Into<AccessContext>,
        limit: usize,
    ) -> Result<Vec<u16>, VmiError> {
        self.core().read_string_utf16_bytes_limited(ctx, limit)
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine.
    pub fn read_string_utf16_bytes_in(
        &self,
        ctx: impl Into<AccessContext>,
    ) -> Result<Vec<u16>, VmiError> {
        self.core().read_string_utf16_bytes(ctx)
    }

    /// Reads a null-terminated string from the virtual machine with a specified
    /// limit.
    pub fn read_string_limited_in(
        &self,
        ctx: impl Into<AccessContext>,
        limit: usize,
    ) -> Result<String, VmiError> {
        self.core().read_string_limited(ctx, limit)
    }

    /// Reads a null-terminated string from the virtual machine.
    pub fn read_string_in(&self, ctx: impl Into<AccessContext>) -> Result<String, VmiError> {
        self.core().read_string(ctx)
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine
    /// with a specified limit.
    pub fn read_string_utf16_limited_in(
        &self,
        ctx: impl Into<AccessContext>,
        limit: usize,
    ) -> Result<String, VmiError> {
        self.core().read_string_utf16_limited(ctx, limit)
    }

    /// Reads a null-terminated wide string (UTF-16) from the virtual machine.
    pub fn read_string_utf16_in(&self, ctx: impl Into<AccessContext>) -> Result<String, VmiError> {
        self.core().read_string_utf16(ctx)
    }

    /// Reads a struct from the virtual machine.
    pub fn read_struct_in<T>(&self, ctx: impl Into<AccessContext>) -> Result<T, VmiError>
    where
        T: IntoBytes + FromBytes,
    {
        self.core().read_struct(ctx)
    }

    // endregion: Read in
}

///////////////////////////////////////////////////////////////////////////////
// VmiRead + VmiWrite
///////////////////////////////////////////////////////////////////////////////

impl<'a, Os> VmiState<'a, Os>
where
    Os: VmiOs,
    Os::Driver: VmiRead + VmiWrite,
{
    /// Writes memory to the virtual machine.
    pub fn write(&self, address: Va, buffer: &[u8]) -> Result<(), VmiError> {
        self.write_in(self.access_context(address), buffer)
    }

    /// Writes memory to the virtual machine.
    pub fn write_in(&self, ctx: impl Into<AccessContext>, buffer: &[u8]) -> Result<(), VmiError> {
        self.core().write(ctx, buffer)
    }

    /// Writes a single byte to the virtual machine.
    pub fn write_u8(&self, address: Va, value: u8) -> Result<(), VmiError> {
        self.core().write_u8(self.access_context(address), value)
    }

    /// Writes a 16-bit unsigned integer to the virtual machine.
    pub fn write_u16(&self, address: Va, value: u16) -> Result<(), VmiError> {
        self.core().write_u16(self.access_context(address), value)
    }

    /// Writes a 32-bit unsigned integer to the virtual machine.
    pub fn write_u32(&self, address: Va, value: u32) -> Result<(), VmiError> {
        self.core().write_u32(self.access_context(address), value)
    }

    /// Writes a 64-bit unsigned integer to the virtual machine.
    pub fn write_u64(&self, address: Va, value: u64) -> Result<(), VmiError> {
        self.core().write_u64(self.access_context(address), value)
    }

    /// Writes a struct to the virtual machine.
    pub fn write_struct<T>(&self, address: Va, value: T) -> Result<(), VmiError>
    where
        T: FromBytes + IntoBytes + Immutable,
    {
        self.core()
            .write_struct(self.access_context(address), value)
    }
}

///////////////////////////////////////////////////////////////////////////////
// VmiSetRegisters
///////////////////////////////////////////////////////////////////////////////

impl<'a, Os> VmiState<'a, Os>
where
    Os: VmiOs,
    Os::Driver: VmiSetRegisters,
{
    /// Sets the registers of a virtual CPU.
    pub fn set_registers(
        &self,
        vcpu: VcpuId,
        registers: <Os::Architecture as Architecture>::Registers,
    ) -> Result<(), VmiError> {
        self.core().set_registers(vcpu, registers)
    }
}

/// Wrapper providing access to OS-specific operations.
pub struct VmiOsState<'a, Os>(VmiState<'a, Os>)
where
    Os: VmiOs;

impl<'a, Os> VmiOsState<'a, Os>
where
    Os: VmiOs,
{
    /// Returns the VMI core.
    pub fn core(&self) -> &'a VmiCore<Os::Driver> {
        self.0.core()
    }

    /// Returns the underlying OS-specific implementation.
    pub fn underlying_os(&self) -> &'a Os {
        self.0.underlying_os()
    }

    /// Returns the VMI session.
    pub fn session(&self) -> &VmiSession<'a, Os> {
        self.0.session()
    }

    /// Returns the VMI state.
    pub fn state(&self) -> VmiState<'a, Os> {
        self.0
    }

    /// Returns the CPU registers associated with the current event.
    pub fn registers(&self) -> &<Os::Architecture as Architecture>::Registers {
        self.0.registers()
    }

    /// Retrieves a specific function argument according to the calling
    /// convention of the operating system.
    pub fn function_argument_for_registers(
        &self,
        registers: &<Os::Architecture as Architecture>::Registers,
        index: u64,
    ) -> Result<u64, VmiError> {
        Os::function_argument(self.0.with_registers(registers), index)
    }

    /// Retrieves the return value of a function.
    pub fn function_return_value_for_registers(
        &self,
        registers: &<Os::Architecture as Architecture>::Registers,
    ) -> Result<u64, VmiError> {
        Os::function_return_value(self.0.with_registers(registers))
    }
}