winstrument 1.0.0

A tiny Rust library to set undocumented Windows process instrumentation callback
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use crate::error::{Error, Result};
use crate::windows::{
    NtSetInformationProcessFn, PROCESS_INSTRUMENTATION_CALLBACK, ProcessInstrumentationCallbackInfo,
};
use defer_rs::defer;
use ntapi::ntpsapi::NtCurrentProcess;
use ntapi::winapi_local::um::winnt::NtCurrentTeb;
use std::arch::naked_asm;
use std::{ffi::c_void, sync::RwLock};
use windows::Win32::{
    Foundation::{FreeLibrary, HANDLE},
    System::{
        Diagnostics::Debug::{CONTEXT, RtlRestoreContext},
        LibraryLoader::{GetProcAddress, LoadLibraryA},
    },
};
use windows_strings::*;

pub type InstrumentationCallback = fn(&mut CONTEXT);
static INSTRUMENTATION_CALLBACK: RwLock<Option<InstrumentationCallback>> = RwLock::new(None);

#[unsafe(no_mangle)]
extern "C" fn instrumentation_callback(ctx: &mut CONTEXT) {
    let teb = unsafe { NtCurrentTeb() };

    unsafe {
        ctx.Rip = (*teb).InstrumentationCallbackPreviousPc as _;
        ctx.Rsp = (*teb).InstrumentationCallbackPreviousSp as _;
        ctx.Rcx = ctx.R10;
        ctx.R10 = ctx.Rip;

        (*teb).InstrumentationCallbackDisabled = 1;
    }

    {
        let callback = INSTRUMENTATION_CALLBACK.read().unwrap();
        callback.unwrap()(ctx);
    }

    unsafe {
        (*teb).InstrumentationCallbackDisabled = 0;

        RtlRestoreContext(ctx, None);
    }
}

#[unsafe(naked)]
#[rustfmt::skip]
pub extern "C" fn callback_stub() {
    naked_asm!(
        "cmp byte ptr gs:[0x2ec], 1",    // Check if InstrumentationCallbackDisabled is set
        "je resume",                     // Skip if so
        "",                              //
        "mov gs:[0x2e0], rsp",           // Save the stack pointer to InstrumentationCallbackPreviousSp
        "mov gs:[0x2d8], r10",           // Save the return address to InstrumentationCallbackPreviousPc
        "mov r10, rcx",                  // Save original RCX
        "sub rsp, 0x4d0",                // Alloc stack space for CONTEXT structure
        "and rsp, -0x10",                // Align stack pointer
        "mov rcx, rsp",                  // Setup RCX for RtlCaptureContext
        "call RtlCaptureContext",        // Capture the context
        "sub rsp, 0x20",                 // Shadow stack space
        "call instrumentation_callback", // Call instrumentation_callback
        "int3",                          // Should never reach here
        "",                              //
        "resume:",                       //
        "jmp r10"                        // Jump to original return address
    );
}

pub fn set_instrumentation_callback(callback: Option<InstrumentationCallback>) -> Result<()> {
    *INSTRUMENTATION_CALLBACK.write().unwrap() = callback;

    let ntdll = unsafe { LoadLibraryA(s!("ntdll.dll")) }.map_err(Error::LoadLibrary)?;
    defer! { unsafe { _ = FreeLibrary(ntdll); } }

    let nt_set_info_proc_addr = unsafe { GetProcAddress(ntdll, s!("NtSetInformationProcess")) };
    let nt_set_info_proc: NtSetInformationProcessFn =
        unsafe { std::mem::transmute(nt_set_info_proc_addr) };

    let mut info = ProcessInstrumentationCallbackInfo {
        version: 0,
        reserved: 0,
        callback: if callback.is_some() {
            callback_stub as *const c_void
        } else {
            std::ptr::null()
        },
    };

    let status = unsafe {
        nt_set_info_proc(
            HANDLE(NtCurrentProcess as _),
            PROCESS_INSTRUMENTATION_CALLBACK,
            &mut info as *mut _ as *mut c_void,
            core::mem::size_of::<ProcessInstrumentationCallbackInfo>() as u32,
        )
    };

    if status.is_err() {
        return Err(Error::NtSetInformationProcess(status));
    }

    Ok(())
}