PEB

I thought walking the PEB to find Kernel32.dll and then lookup the functions manually might save some output bytes, but it didn't - even when I had possible UB code. It can probably be optimised more, but as it stands it's 784 bytes on x64 and 688 bytes on x86. Once adding the linker flag /MERGE:.rdata=.text, we save 40 bytes (744 and 688, respectively). The non PEB version is currently at 655 and 576, and I'm not convinced we can reduce it by another 100 bytes.

This might change if we were looking up more functions (or using LoadLibrary to get more imports), but for our two functions from Kernel32, it seems like a useless endeavour.

I'm leaving this here in-case anyone wants to walk the PEB in the future. I have no idea how this works on other architectures (eg aarch64).

#![no_std]
#![no_main]

#[repr(C)]
struct ListEntry {
    flink: *mut ListEntry,
    blink: *mut ListEntry,
}

#[repr(C)]
struct UnicodeString {
    length: u16,
    maximum_length: u16,
    buffer: *mut u16,
}

#[repr(C)]
struct LdrDataTableEntry {
    in_load_order_links: ListEntry,
    in_memory_order_links: ListEntry,
    in_initialization_order_links: ListEntry,
    dll_base: *mut u8,
    entry_point: *mut u8,
    size_of_image: u32,
    full_dll_name: UnicodeString,
    base_dll_name: UnicodeString,
    flags: u32,
    load_count: u16,
    tls_index: u16,
    hash_links: ListEntry,
    time_date_stamp: u32,
}

#[repr(C)]
struct PebLdrData {
    length: u32,
    initialized: u8,
    ss_handle: *mut u8,
    in_load_order_module_list: ListEntry,
    in_memory_order_module_list: ListEntry,
    in_initialization_order_module_list: ListEntry,
}

#[repr(C)]
struct Peb {
    inherited_address_space: u8,
    read_image_file_exec_options: u8,
    being_debugged: u8,
    bit_field: u8,
    mutant: *mut u8,
    image_base_address: *mut u8,
    ldr: *mut PebLdrData,
    process_parameters: *mut u8,
    sub_system_data: *mut u8,
    process_heap: *mut u8,
    fast_peb_lock: *mut u8,
    atl_thunk_s_list_ptr: *mut u8,
    ifeo_key: *mut u8,
    cross_process_flags: u32,
    kernel_callback_table: *mut u8,
    system_reserved: u32,
    atl_thunk_s_list_ptr32: u32,
    api_set_map: *mut u8,
}

#[repr(C)]
struct ImageDosHeader {
    e_magic: u16,
    e_cblp: u16,
    e_cp: u16,
    e_crlc: u16,
    e_cparhdr: u16,
    e_minalloc: u16,
    e_maxalloc: u16,
    e_ss: u16,
    e_sp: u16,
    e_csum: u16,
    e_ip: u16,
    e_cs: u16,
    e_lfarlc: u16,
    e_ovno: u16,
    e_res: [u16; 4],
    e_oemid: u16,
    e_oeminfo: u16,
    e_res2: [u16; 10],
    e_lfanew: i32,
}

#[repr(C)]
struct ImageFileHeader {
    machine: u16,
    number_of_sections: u16,
    time_date_stamp: u32,
    pointer_to_symbol_table: u32,
    number_of_symbols: u32,
    size_of_optional_header: u16,
    characteristics: u16,
}

#[repr(C)]
struct ImageDataDirectory {
    virtual_address: u32,
    size: u32,
}

#[repr(C)]
struct ImageOptionalHeader {
    magic: u16,
    major_linker_version: u8,
    minor_linker_version: u8,
    size_of_code: u32,
    size_of_initialized_data: u32,
    size_of_uninitialized_data: u32,
    address_of_entry_point: u32,
    base_of_code: u32,
    #[cfg(target_arch = "x86")]
    base_of_data: u32,
    image_base: usize,
    section_alignment: u32,
    file_alignment: u32,
    major_os_version: u16,
    minor_os_version: u16,
    major_image_version: u16,
    minor_image_version: u16,
    major_subsystem_version: u16,
    minor_subsystem_version: u16,
    win32_version_value: u32,
    size_of_image: u32,
    size_of_headers: u32,
    checksum: u32,
    subsystem: u16,
    dll_characteristics: u16,
    size_of_stack_reserve: usize,
    size_of_stack_commit: usize,
    size_of_heap_reserve: usize,
    size_of_heap_commit: usize,
    loader_flags: u32,
    number_of_rva_and_sizes: u32,
    data_directory: [ImageDataDirectory; 16],
}

#[repr(C)]
struct ImageNtHeaders {
    signature: u32,
    file_header: ImageFileHeader,
    optional_header: ImageOptionalHeader,
}

#[repr(C)]
struct ImageExportDirectory {
    characteristics: u32,
    time_date_stamp: u32,
    major_version: u16,
    minor_version: u16,
    name: u32,
    base: u32,
    number_of_functions: u32,
    number_of_names: u32,
    address_of_functions: u32,
    address_of_names: u32,
    address_of_name_ordinals: u32,
}

use core::{ffi::c_void, panic::PanicInfo};

const STD_OUTPUT_HANDLE: u32 = -11i32 as u32;

type ExitProcessFn = unsafe extern "system" fn(u32) -> !;
type WriteFileFn = unsafe extern "system" fn(
    hfile: *mut c_void,
    lpbuffer: *const u8,
    nnumberofbytestowrite: u32,
    lpnumberofbyteswritten: *mut u32,
    lpoverlapped: *mut c_void,
) -> i32;

#[unsafe(no_mangle)]
pub extern "C" fn mainCRTStartup() -> ! {
    let msg = b"Hello, world!\n";
    unsafe {
        // NOTE: UB if kernel32 is not loaded (or another module with the same length is loaded first - unlikely)
        let peb: *mut Peb;
        #[cfg(target_arch = "x86")]
        core::arch::asm!("mov {}, fs:[0x30]", out(reg) peb);
        #[cfg(target_arch = "x86_64")]
        core::arch::asm!("mov {}, gs:[0x60]", out(reg) peb);
        let ldr = (*peb).ldr;

        let mut entry = (*ldr).in_load_order_module_list.flink as *mut LdrDataTableEntry;

        // Skip the first entry (main executable)
        entry = (*entry).in_load_order_links.flink as *mut LdrDataTableEntry;

        let module_base;
        loop {
            let current_entry = &*entry;

            if *(&current_entry.base_dll_name).buffer.offset(12) == 0 {
                module_base = current_entry.dll_base;
                break;
            }

            entry = (*entry).in_load_order_links.flink as *mut LdrDataTableEntry;
        }

        let dos_header = module_base as *const ImageDosHeader;
        let nt_headers =
            (module_base as usize + (*dos_header).e_lfanew as usize) as *const ImageNtHeaders;
        let export_dir = (module_base as usize
            + (*nt_headers).optional_header.data_directory[0].virtual_address as usize)
            as *const ImageExportDirectory;

        let name_table =
            (module_base as usize + (*export_dir).address_of_names as usize) as *const u32;
        let func_table =
            (module_base as usize + (*export_dir).address_of_functions as usize) as *const u32;
        let ordinal_table =
            (module_base as usize + (*export_dir).address_of_name_ordinals as usize) as *const u16;

        let mut exit_process: *mut u8 = core::ptr::null_mut();
        let mut write_file: *mut u8 = core::ptr::null_mut();

        for i in 0..(*export_dir).number_of_names {
            let name_ptr =
                (module_base as usize + *name_table.offset(i as isize) as usize) as *const u8;

            if exit_process.is_null() && matches_name(name_ptr, b"ExitProcess") {
                let ordinal = *ordinal_table.offset(i as isize);
                exit_process =
                    (module_base as usize + *func_table.offset(ordinal as isize) as usize) as *mut u8;
            }

            if write_file.is_null() && matches_name(name_ptr, b"WriteFile") {
                let ordinal = *ordinal_table.offset(i as isize);
                write_file =
                    (module_base as usize + *func_table.offset(ordinal as isize) as usize) as *mut u8;
            }

            if !exit_process.is_null() && !write_file.is_null() {
                break;
            }
        }

        let handle = STD_OUTPUT_HANDLE as *mut c_void;

        let exit_process: ExitProcessFn = core::mem::transmute(exit_process);
        let write_file: WriteFileFn = core::mem::transmute(write_file);

        write_file(
            handle,
            msg.as_ptr(),
            msg.len() as u32,
            0 as *mut _,
            0 as *mut _,
        );
        exit_process(0)
    }
}

#[inline(always)]
unsafe fn matches_name(ptr: *const u8, name: &[u8]) -> bool {
    unsafe {
        for i in 0..name.len() {
            if *ptr.add(i) != name[i] {
                return false;
            }
        }
        *ptr.add(name.len()) == 0
    }
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    loop {}
}