Running Executable Machine Code at Runtime - Windows

Cargo.toml

[package]
name = "foo"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]

[dependencies.windows]
version = "0.51.1"
features = ["Win32_Foundation", "Win32_System_Memory"]

main.rs

#![forbid(unsafe_op_in_unsafe_fn)]

use std::{ffi::c_void, marker::PhantomData};

use windows::{
    core::Error,
    Win32::{
        Foundation::GetLastError,
        System::Memory::{
            VirtualAlloc, VirtualFree, VirtualProtect, MEM_COMMIT, MEM_RELEASE, PAGE_EXECUTE_READ,
            PAGE_PROTECTION_FLAGS, PAGE_READWRITE,
        },
    },
};

struct ExecuteCode<T> {
    buffer: *mut c_void,
    _phantom: PhantomData<T>,
}

unsafe impl<T> Sync for ExecuteCode<T> {}
unsafe impl<T> Send for ExecuteCode<T> {}

impl<T> ExecuteCode<T> {
    /// Const check to ensure T is correct size
    #[allow(unused)]
    const SIZE_OK: () = assert!(std::mem::size_of::<T>() == std::mem::size_of::<*mut c_void>());

    /// SAFETY:
    /// Your machine code must be safe. You should know what it does.
    /// It is machine code. It can do absolutely anything. Be careful.
    unsafe fn new(code: &[u8]) -> Result<Self, Error> {
        let buffer = unsafe { VirtualAlloc(None, code.len(), MEM_COMMIT, PAGE_READWRITE) };

        // error check
        if buffer.is_null() {
            return Err(unsafe { GetLastError().unwrap_err() });
        }

        // SAFETY:
        // The size of the buffer is the same as our code size.
        // We allocated with read and write perms
        unsafe {
            std::ptr::copy(code.as_ptr(), buffer as *mut _, code.len());
        }

        let mut old_protect = PAGE_PROTECTION_FLAGS::default();

        // Protect after writing data to prevent any writing
        //
        // SAFETY:
        // We change protection on the same size as we allocated with
        let res =
            unsafe { VirtualProtect(buffer, code.len(), PAGE_EXECUTE_READ, &mut old_protect) };

        // release memory if virtual protect failed
        if let Err(e) = res {
            unsafe {
                VirtualFree(buffer, 0, MEM_RELEASE)?;
            }

            Err(e)?;
        }

        Ok(Self {
            buffer,
            _phantom: PhantomData,
        })
    }

    /// SAFETY:
    /// T must be a fn pointer signature, and the signature MUST be correct
    unsafe fn run<F: FnOnce(T)>(&self, cb: F) {
        // SAFETY:
        // T is the correct size due to const check
        // All else, caller ensures signature is correct
        let fn_ptr: T = unsafe { std::mem::transmute_copy(&(self.buffer as *mut _ as *mut T)) };

        cb(fn_ptr);
    }
}

impl<T> Drop for ExecuteCode<T> {
    fn drop(&mut self) {
        unsafe {
            // ignore if it failed. Nothing to do if it failed
            _ = VirtualFree(self.buffer, 0, MEM_RELEASE);
        }
    }
}

fn main() {
    // add_2(a: usize) -> usize {
    //     a + 2
    // }
    #[rustfmt::skip]
    let code = &[
        0x48, 0x83, 0xC1, 0x02, // add rcx,2
        0x48, 0x89, 0xC8,       // mov rax,rcx
        0xC3,                   // ret
    ];

    let executor = unsafe { ExecuteCode::<fn(usize) -> usize>::new(code).expect("succeess") };

    unsafe {
        executor.run(|fn_| {
            let result = fn_(2);
            assert!(result == 4, "Bad error oh no");
            println!("2+2={}", result);
        })
    };
}