dyncvoke-spoof 0.1.1

//! Walk UNWIND_INFO records to derive stack sizes and RBP offsets. Three
//! walkers (rbp_offset, stack_frame, ignoring_set_fpreg) verified against
//! live ntdll, kernelbase, and kernel32 prologues.

use core::ffi::c_void;

use crate::types::{
    ImageRuntimeFunction, Registers, UnwindCode, UnwindInfo, UnwindOpCode::*,
    UnwindOpCode, UNW_FLAG_CHAININFO, UNW_FLAG_EHANDLER,
};

/// Determine RBP push location and total stack size. Returns None if the
/// frame layout isn't spoof-compatible (RSP saved as nonvol, RBP pushed
/// twice, etc.).
pub unsafe fn rbp_offset(
    module: *mut c_void,
    runtime: &ImageRuntimeFunction,
) -> Option<(u32, u32)> {
    let unwind_info = (module as usize + runtime.UnwindData as usize) as *mut UnwindInfo;
    let unwind_code = (unwind_info as *mut u8).add(4) as *mut UnwindCode;
    let flag = (*unwind_info).VersionFlags.Flags();

    let mut i = 0usize;
    let mut total_stack = 0u32;
    let mut rbp_pushed = false;
    let mut stack_offset = 0u32;

    while i < (*unwind_info).CountOfCodes as usize {
        let unwind_code = unwind_code.add(i);
        let op_info = (*unwind_code).Anonymous.OpInfo() as usize;
        let unwind_op = (*unwind_code).Anonymous.UnwindOp();

        match UnwindOpCode::try_from(unwind_op) {
            Ok(UWOP_PUSH_NONVOL) => {
                if Registers::Rsp == op_info {
                    return None;
                }
                if Registers::Rbp == op_info {
                    if rbp_pushed {
                        return None;
                    }
                    rbp_pushed = true;
                    stack_offset = total_stack;
                }
                total_stack += 8;
                i += 1;
            }
            Ok(UWOP_ALLOC_LARGE) => {
                if (*unwind_code).Anonymous.OpInfo() == 0 {
                    let frame_offset = ((*unwind_code.add(1)).FrameOffset as i32) * 8;
                    total_stack += frame_offset as u32;
                    i += 2;
                } else {
                    let frame_offset = *(unwind_code.add(1) as *mut i32);
                    total_stack += frame_offset as u32;
                    i += 3;
                }
            }
            Ok(UWOP_ALLOC_SMALL) => {
                total_stack += ((op_info + 1) * 8) as u32;
                i += 1;
            }
            Ok(UWOP_SAVE_NONVOL) => {
                if Registers::Rsp == op_info {
                    return None;
                }
                if Registers::Rbp == op_info {
                    if rbp_pushed {
                        return None;
                    }
                    let offset = (*unwind_code.add(1)).FrameOffset * 8;
                    stack_offset = total_stack + offset as u32;
                    rbp_pushed = true;
                }
                i += 2;
            }
            Ok(UWOP_SAVE_NONVOL_BIG) => {
                if Registers::Rsp == op_info {
                    return None;
                }
                if Registers::Rbp == op_info {
                    if rbp_pushed {
                        return None;
                    }
                    let offset = *(unwind_code.add(1) as *mut u32);
                    stack_offset = total_stack + offset;
                    rbp_pushed = true;
                }
                i += 3;
            }
            Ok(UWOP_SET_FPREG) => return None,
            Ok(UWOP_SAVE_XMM128) => i += 2,
            Ok(UWOP_SAVE_XMM128BIG) => i += 3,
            Ok(UWOP_EPILOG) | Ok(UWOP_SPARE_CODE) => i += 1,
            Ok(UWOP_PUSH_MACH_FRAME) => {
                total_stack += if op_info == 0 { 0x40 } else { 0x48 };
                i += 1;
            }
            _ => return None,
        }
    }

    if (flag & UNW_FLAG_CHAININFO) != 0 {
        let count = (*unwind_info).CountOfCodes as usize;
        let index = if count & 1 == 1 { count + 1 } else { count };
        let runtime = unwind_code.add(index) as *const ImageRuntimeFunction;
        if let Some((_, child_total)) = rbp_offset(module, &*runtime) {
            total_stack += child_total;
        } else {
            return None;
        }
    }

    Some((stack_offset, total_stack))
}

/// Compute total stack size, rejecting any frame that uses RBP/RSP as a
/// non-volatile save target. Used when picking the first decoy prologue.
pub unsafe fn stack_frame(
    module: *mut c_void,
    runtime: &ImageRuntimeFunction,
) -> Option<(bool, u32)> {
    let unwind_info = (module as usize + runtime.UnwindData as usize) as *mut UnwindInfo;
    let unwind_code = (unwind_info as *mut u8).add(4) as *mut UnwindCode;
    let flag = (*unwind_info).VersionFlags.Flags();

    let mut i = 0usize;
    let mut set_fpreg_hit = false;
    let mut total_stack = 0i32;

    while i < (*unwind_info).CountOfCodes as usize {
        let unwind_code = unwind_code.add(i);
        let op_info = (*unwind_code).Anonymous.OpInfo() as usize;
        let unwind_op = (*unwind_code).Anonymous.UnwindOp();

        match UnwindOpCode::try_from(unwind_op) {
            Ok(UWOP_PUSH_NONVOL) => {
                if Registers::Rsp == op_info && !set_fpreg_hit {
                    return None;
                }
                total_stack += 8;
                i += 1;
            }
            Ok(UWOP_ALLOC_SMALL) => {
                total_stack += ((op_info + 1) * 8) as i32;
                i += 1;
            }
            Ok(UWOP_ALLOC_LARGE) => {
                if (*unwind_code).Anonymous.OpInfo() == 0 {
                    let frame_offset = ((*unwind_code.add(1)).FrameOffset as i32) * 8;
                    total_stack += frame_offset;
                    i += 2;
                } else {
                    let frame_offset = *(unwind_code.add(1) as *mut i32);
                    total_stack += frame_offset;
                    i += 3;
                }
            }
            Ok(UWOP_SAVE_NONVOL) => {
                if Registers::Rsp == op_info || Registers::Rbp == op_info {
                    return None;
                }
                i += 2;
            }
            Ok(UWOP_SAVE_NONVOL_BIG) => {
                if Registers::Rsp == op_info || Registers::Rbp == op_info {
                    return None;
                }
                i += 3;
            }
            Ok(UWOP_SAVE_XMM128) => i += 2,
            Ok(UWOP_SAVE_XMM128BIG) => i += 3,
            Ok(UWOP_SET_FPREG) => {
                if (flag & UNW_FLAG_EHANDLER) != 0 && (flag & UNW_FLAG_CHAININFO) != 0 {
                    return None;
                }
                if (*unwind_info).FrameInfo.FrameRegister() != Registers::Rbp as u8 {
                    return None;
                }
                set_fpreg_hit = true;
                let offset = ((*unwind_info).FrameInfo.FrameOffset() as i32) << 4;
                total_stack -= offset;
                i += 1;
            }
            Ok(UWOP_EPILOG) | Ok(UWOP_SPARE_CODE) => i += 1,
            Ok(UWOP_PUSH_MACH_FRAME) => {
                total_stack += if op_info == 0 { 0x40 } else { 0x48 };
                i += 1;
            }
            _ => return None,
        }
    }

    if (flag & UNW_FLAG_CHAININFO) != 0 {
        let count = (*unwind_info).CountOfCodes as usize;
        let index = if count & 1 == 1 { count + 1 } else { count };
        let runtime = unwind_code.add(index) as *const ImageRuntimeFunction;
        if let Some((chained_fpreg_hit, chained_stack)) = stack_frame(module, &*runtime) {
            total_stack += chained_stack as i32;
            set_fpreg_hit |= chained_fpreg_hit;
        } else {
            return None;
        }
    }

    Some((set_fpreg_hit, total_stack as u32))
}

/// Total stack size, treating set_fpreg as a no-op. Used for sizing the
/// frames we splice in (RtlUserThreadStart, BaseThreadInitThunk, gadget
/// frames).
pub unsafe fn ignoring_set_fpreg(
    module: *mut c_void,
    runtime: &ImageRuntimeFunction,
) -> Option<u32> {
    let unwind_info = (module as usize + runtime.UnwindData as usize) as *mut UnwindInfo;
    let unwind_code = (unwind_info as *mut u8).add(4) as *mut UnwindCode;
    let flag = (*unwind_info).VersionFlags.Flags();

    let mut i = 0usize;
    let mut total_stack = 0u32;

    while i < (*unwind_info).CountOfCodes as usize {
        let unwind_code = unwind_code.add(i);
        let op_info = (*unwind_code).Anonymous.OpInfo() as usize;
        let unwind_op = (*unwind_code).Anonymous.UnwindOp();

        match UnwindOpCode::try_from(unwind_op) {
            Ok(UWOP_PUSH_NONVOL) => {
                if Registers::Rsp == op_info {
                    return None;
                }
                total_stack += 8;
                i += 1;
            }
            Ok(UWOP_ALLOC_SMALL) => {
                total_stack += ((op_info + 1) * 8) as u32;
                i += 1;
            }
            Ok(UWOP_ALLOC_LARGE) => {
                if (*unwind_code).Anonymous.OpInfo() == 0 {
                    let frame_offset = ((*unwind_code.add(1)).FrameOffset as i32) * 8;
                    total_stack += frame_offset as u32;
                    i += 2;
                } else {
                    let frame_offset = *(unwind_code.add(1) as *mut i32);
                    total_stack += frame_offset as u32;
                    i += 3;
                }
            }
            Ok(UWOP_SAVE_NONVOL) => {
                if Registers::Rsp == op_info {
                    return None;
                }
                i += 2;
            }
            Ok(UWOP_SAVE_NONVOL_BIG) => {
                if Registers::Rsp == op_info {
                    return None;
                }
                i += 3;
            }
            Ok(UWOP_SAVE_XMM128) => i += 2,
            Ok(UWOP_SAVE_XMM128BIG) => i += 3,
            Ok(UWOP_SET_FPREG) => i += 1,
            Ok(UWOP_EPILOG) | Ok(UWOP_SPARE_CODE) => i += 1,
            Ok(UWOP_PUSH_MACH_FRAME) => {
                total_stack += if op_info == 0 { 0x40 } else { 0x48 };
                i += 1;
            }
            _ => return None,
        }
    }

    if (flag & UNW_FLAG_CHAININFO) != 0 {
        let count = (*unwind_info).CountOfCodes as usize;
        let index = if count & 1 == 1 { count + 1 } else { count };
        let runtime = unwind_code.add(index) as *const ImageRuntimeFunction;
        if let Some(chained_stack) = ignoring_set_fpreg(module, &*runtime) {
            total_stack += chained_stack;
        } else {
            return None;
        }
    }

    Some(total_stack)
}