memdump-kmdf.rs

#![no_std]
#![no_main]
#![allow(non_camel_case_types, non_snake_case, dead_code)]

#[cfg(not(test))]
extern crate wdk_panic;

#[cfg(not(test))]
use wdk_alloc::WdkAllocator;
#[cfg(not(test))]
#[global_allocator]
static GLOBAL_ALLOCATOR: WdkAllocator = WdkAllocator;

extern crate alloc;

use alloc::vec::Vec;
use core::mem::{size_of, zeroed};
use core::ptr;

use wdk::println;

type NTSTATUS = i32;
type ULONG = u32;
type SIZE_T = usize;
type HANDLE = *mut core::ffi::c_void;
type ACCESS_MASK = u32;

const STATUS_SUCCESS: NTSTATUS = 0;
const STATUS_INSUFFICIENT_RESOURCES: NTSTATUS = -1073741801; // 0xC000009A
const STATUS_UNSUCCESSFUL: NTSTATUS = -1073741823; // 0xC0000001

#[repr(C)]
#[derive(Copy, Clone)]
pub struct PHYSICAL_ADDRESS {
    pub QuadPart: i64,
}

#[repr(C)]
pub struct PHYSICAL_MEMORY_RANGE {
    pub BaseAddress: PHYSICAL_ADDRESS,
    pub NumberOfBytes: i64,
}

#[repr(C)]
pub union MM_COPY_ADDRESS {
    pub VirtualAddress: *mut core::ffi::c_void,
    pub PhysicalAddress: PHYSICAL_ADDRESS,
}

#[repr(C)]
pub struct UNICODE_STRING {
    pub Length: u16,
    pub MaximumLength: u16,
    pub Buffer: *mut u16,
}

#[repr(C)]
pub struct OBJECT_ATTRIBUTES {
    pub Length: ULONG,
    pub RootDirectory: HANDLE,
    pub ObjectName: *mut UNICODE_STRING,
    pub Attributes: ULONG,
    pub SecurityDescriptor: *mut core::ffi::c_void,
    pub SecurityQualityOfService: *mut core::ffi::c_void,
}

#[repr(C)]
pub struct IO_STATUS_BLOCK {
    pub Pointer: *mut core::ffi::c_void,
    pub Information: usize,
}

#[repr(C)]
pub struct DRIVER_OBJECT {
    pub DriverStart: *mut core::ffi::c_void,
    _pad: [u8; 0x200],
    pub DriverUnload: Option<unsafe extern "C" fn(*mut DRIVER_OBJECT)>,
}

unsafe extern "system" {

    fn MmGetPhysicalMemoryRanges() -> *mut PHYSICAL_MEMORY_RANGE;

    fn MmCopyMemory(
        TargetAddress: *mut core::ffi::c_void,
        SourceAddress: MM_COPY_ADDRESS,
        NumberOfBytes: SIZE_T,
        Flags: ULONG,
        NumberOfBytesTransferred: *mut SIZE_T,
    ) -> NTSTATUS;

    fn ExAllocatePoolWithTag(
        pool_type: ULONG,
        number_of_bytes: SIZE_T,
        tag: u32,
    ) -> *mut core::ffi::c_void;
    fn ExFreePoolWithTag(p: *mut core::ffi::c_void, tag: u32);

    fn ZwCreateFile(
        FileHandle: *mut HANDLE,
        DesiredAccess: ACCESS_MASK,
        ObjectAttributes: *mut OBJECT_ATTRIBUTES,
        IoStatusBlock: *mut IO_STATUS_BLOCK,
        AllocationSize: *mut i64,
        FileAttributes: ULONG,
        ShareAccess: ULONG,
        CreateDisposition: ULONG,
        CreateOptions: ULONG,
        EaBuffer: *mut core::ffi::c_void,
        EaLength: ULONG,
    ) -> NTSTATUS;

    fn ZwWriteFile(
        FileHandle: HANDLE,
        Event: HANDLE,
        ApcRoutine: *mut core::ffi::c_void,
        ApcContext: *mut core::ffi::c_void,
        IoStatusBlock: *mut IO_STATUS_BLOCK,
        Buffer: *const core::ffi::c_void,
        Length: ULONG,
        ByteOffset: *mut i64,
        Key: *mut ULONG,
    ) -> NTSTATUS;

    fn ZwClose(Handle: HANDLE);
}

const MM_COPY_MEMORY_PHYSICAL: ULONG = 0x1;
const POOL_TYPE_NON_PAGED: ULONG = 0;
const TAG: u32 = 0x504D3848; // 'H8MP' (8MB Physical)

const OBJ_KERNEL_HANDLE: ULONG = 0x00000200;
const OBJ_CASE_INSENSITIVE: ULONG = 0x00000040;

const FILE_ATTRIBUTE_NORMAL: ULONG = 0x00000080;
const FILE_OVERWRITE_IF: ULONG = 0x00000005;
const FILE_SYNCHRONOUS_IO_NONALERT: ULONG = 0x00000020;
const FILE_NON_DIRECTORY_FILE: ULONG = 0x00000040;

const GENERIC_WRITE: ACCESS_MASK = 0x40000000;
const FILE_SHARE_READ: ULONG = 0x00000001;
const FILE_SHARE_WRITE: ULONG = 0x00000002;

const CHUNK_SIZE: usize = 8 * 1024 * 1024; // 8MB

// L"\\??\\C:\\Windows\\OriginalMemoryDump.dmp"
static PATH_W: [u16; 38] = [
    0x005C, 0x003F, 0x003F, 0x005C, // \??\
    0x0043, 0x003A, 0x005C, // C:\
    0x0057, 0x0069, 0x006E, 0x0064, 0x006F, 0x0077, 0x0073, 0x005C, // Windows\
    0x004F, 0x0072, 0x0069, 0x0067, 0x0069, 0x006E, 0x0061, 0x006C, // Original
    0x004D, 0x0065, 0x006D, 0x006F, 0x0072, 0x0079, // Memory
    0x0044, 0x0075, 0x006D, 0x0070, // Dump
    0x002E, 0x0064, 0x006D, 0x0070, // .dmp
    0x0000,
];

#[inline(always)]
unsafe fn initialize_object_attributes(
    oa: *mut OBJECT_ATTRIBUTES,
    name: *mut UNICODE_STRING,
    attributes: ULONG,
    root: HANDLE,
    sd: *mut core::ffi::c_void,
) {
    unsafe {
        (*oa).Length = size_of::<OBJECT_ATTRIBUTES>() as ULONG;
        (*oa).RootDirectory = root;
        (*oa).Attributes = attributes;
        (*oa).ObjectName = name;
        (*oa).SecurityDescriptor = sd;
        (*oa).SecurityQualityOfService = ptr::null_mut();
    }
}

unsafe fn make_unicode_const(buf: *const u16) -> UNICODE_STRING {
    let mut len = 0usize;
    unsafe {
        while *buf.add(len) != 0 {
            len += 1;
        }
    }
    UNICODE_STRING {
        Length: (len * 2) as u16,
        MaximumLength: ((len + 1) * 2) as u16,
        Buffer: buf as *mut u16,
    }
}

fn percent(done: u64, total: u64) -> u64 {
    if total == 0 {
        100
    } else {
        (done.saturating_mul(100)).saturating_div(total)
    }
}

unsafe fn enumerate_physical_ranges() -> Result<Vec<(u64, u64)>, NTSTATUS> {
    let p = unsafe { MmGetPhysicalMemoryRanges() };
    if p.is_null() {
        return Err(STATUS_UNSUCCESSFUL);
    }
    let mut out: Vec<(u64, u64)> = Vec::new();
    let mut i = 0usize;
    loop {
        let r = unsafe { p.add(i) };
        let base = unsafe { (*r).BaseAddress.QuadPart as u64 };
        let size = unsafe { (*r).NumberOfBytes as u64 };
        if base == 0 && size == 0 {
            break;
        }
        out.push((base, size));
        i += 1;
    }
    Ok(out)
}

unsafe fn copy_phys_chunk(phys: u64, buf: *mut u8, bytes: usize) -> Result<usize, NTSTATUS> {
    let src = MM_COPY_ADDRESS {
        PhysicalAddress: PHYSICAL_ADDRESS {
            QuadPart: phys as i64,
        },
    };
    let mut transferred: SIZE_T = 0;
    let st = unsafe {
        MmCopyMemory(
            buf as *mut _,
            src,
            bytes,
            MM_COPY_MEMORY_PHYSICAL,
            &mut transferred as *mut SIZE_T,
        )
    };
    if st != STATUS_SUCCESS {
        return Err(st);
    }
    Ok(transferred)
}

unsafe fn write_all_sync(
    handle: HANDLE,
    mut ptr_buf: *const u8,
    mut len: usize,
    mut offset: i64,
) -> Result<(), NTSTATUS> {
    while len > 0 {
        let mut ios: IO_STATUS_BLOCK = unsafe { zeroed() };
        let chunk = if len > u32::MAX as usize {
            u32::MAX as usize
        } else {
            len
        };
        let st = unsafe {
            ZwWriteFile(
                handle,
                ptr::null_mut(),
                ptr::null_mut(),
                ptr::null_mut(),
                &mut ios as *mut IO_STATUS_BLOCK,
                ptr_buf as *const _,
                chunk as u32,
                &mut offset as *mut i64,
                ptr::null_mut(),
            )
        };
        if st != STATUS_SUCCESS {
            return Err(st);
        }
        let written = ios.Information;
        if written == 0 {
            return Err(STATUS_UNSUCCESSFUL);
        }
        unsafe {
            ptr_buf = ptr_buf.add(written);
        }
        len -= written;
        offset += written as i64;
    }
    Ok(())
}

unsafe fn dump_all_physical_memory_8mb(path: *const u16) -> Result<(), NTSTATUS> {
    let mut us = unsafe { make_unicode_const(path) };
    let mut oa: OBJECT_ATTRIBUTES = unsafe { zeroed() };
    unsafe {
        initialize_object_attributes(
            &mut oa as *mut _,
            &mut us as *mut UNICODE_STRING,
            OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
            ptr::null_mut(),
            ptr::null_mut(),
        );
    }

    let mut ios: IO_STATUS_BLOCK = unsafe { zeroed() };
    let mut h: HANDLE = ptr::null_mut();
    let st = unsafe {
        ZwCreateFile(
            &mut h as *mut HANDLE,
            GENERIC_WRITE,
            &mut oa as *mut OBJECT_ATTRIBUTES,
            &mut ios as *mut IO_STATUS_BLOCK,
            ptr::null_mut(),
            FILE_ATTRIBUTE_NORMAL,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            FILE_OVERWRITE_IF,
            FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT,
            ptr::null_mut(),
            0,
        )
    };
    if st != STATUS_SUCCESS {
        println!("dump: ZwCreateFile failed: 0x{:08X}", st as u32);
        return Err(st);
    }
    println!("dump: opened file ok");

    let buf = unsafe { ExAllocatePoolWithTag(POOL_TYPE_NON_PAGED, CHUNK_SIZE, TAG) };
    if buf.is_null() {
        unsafe {
            ZwClose(h);
        }
        println!("dump: buffer allocation failed");
        return Err(STATUS_INSUFFICIENT_RESOURCES);
    }
    println!("dump: allocated {} bytes buffer", CHUNK_SIZE);

    let ranges = unsafe { enumerate_physical_ranges()? };
    let total_bytes: u64 = ranges.iter().map(|(_, sz)| *sz).sum();
    println!(
        "dump: total physical bytes reported = {} (0x{:X})",
        total_bytes, total_bytes
    );

    let mut global_done: u64 = 0;
    let mut file_offset: i64 = 0;

    for (base, size) in ranges {
        if size == 0 {
            continue;
        }
        let mut off: u64 = 0;
        while off < size {
            let remain = (size - off) as usize;
            let this = if remain > CHUNK_SIZE {
                CHUNK_SIZE
            } else {
                remain
            };

            match unsafe { copy_phys_chunk(base + off, buf as *mut u8, this) } {
                Ok(transferred) => {
                    if transferred == 0 {
                        println!("dump: zero bytes transferred at phys=0x{:X}", base + off);
                        off = off.saturating_add(0x1000);
                        global_done = global_done.saturating_add(0x1000);
                        continue;
                    }
                    if let Err(wst) =
                        unsafe { write_all_sync(h, buf as *const u8, transferred, file_offset) }
                    {
                        println!("dump: ZwWriteFile failed: 0x{:08X}", wst as u32);
                        unsafe {
                            ExFreePoolWithTag(buf, TAG);
                            ZwClose(h);
                        }
                        return Err(wst);
                    }

                    off = off.saturating_add(transferred as u64);
                    file_offset += transferred as i64;
                    global_done = global_done.saturating_add(transferred as u64);

                    let p = percent(global_done, total_bytes);
                    if (global_done % (1u64 << 30)) < transferred as u64 {
                        println!(
                            "dump: progress {}% ({} / {} bytes)",
                            p, global_done, total_bytes
                        );
                    }
                }
                Err(cst) => {
                    println!(
                        "dump: MmCopyMemory failed at phys=0x{:X}, status=0x{:08X} -> skip 4KiB",
                        base + off,
                        cst as u32
                    );
                    off = off.saturating_add(0x1000);
                    global_done = global_done.saturating_add(0x1000);
                }
            }
        }
    }

    unsafe {
        ExFreePoolWithTag(buf, TAG);
        ZwClose(h);
    }
    println!("dump: completed, total written = {} bytes", global_done);
    Ok(())
}

#[unsafe(no_mangle)]
pub unsafe extern "system" fn DriverEntry(
    driver: *mut DRIVER_OBJECT,
    _path: *mut UNICODE_STRING,
) -> NTSTATUS {
    unsafe {
        (*driver).DriverUnload = Some(driver_unload);
    }

    println!("DriverEntry: start memory dump (8MB chunks)");
    let st = unsafe { dump_all_physical_memory_8mb(PATH_W.as_ptr()) };
    if let Err(e) = st {
        println!("DriverEntry: dump failed: 0x{:08X}", e as u32);
        return e;
    }
    println!("DriverEntry: dump done");
    STATUS_SUCCESS
}

unsafe extern "C" fn driver_unload(_driver: *mut DRIVER_OBJECT) {
    println!("DriverUnload");
}