OpenFUT/openfut-hook/src/connectex_hook.rs

use core::ffi::c_void;
/// Intercepts ConnectEx (EA/DirtySDK's preferred async connect API).
///
/// DirtySDK calls WSAIoctl(SIO_GET_EXTENSION_FUNCTION_POINTER, WSAID_CONNECTEX) once at
/// startup to get a ConnectEx function pointer, bypassing all IAT hooks.  We hook WSAIoctl
/// inline so that when it returns a ConnectEx pointer we swap it for our own wrapper.
use core::sync::atomic::{AtomicUsize, Ordering};

// Address rewriting (v4 + v6) is shared from connect_hook::redirect_if_ea, so the port
// constants and sockaddr structs no longer live here.

// SIO_GET_EXTENSION_FUNCTION_POINTER
const SIO_GET_EXT_FN: u32 = 0xC8000006;

// WSAID_CONNECTEX = {25A207B9-DDF3-4660-8EE9-76E58C74063E}
const CONNECTEX_GUID: [u8; 16] = [
    0xB9, 0x07, 0xA2, 0x25, 0xF3, 0xDD, 0x60, 0x46, 0x8E, 0xE9, 0x76, 0xE5, 0x8C, 0x74, 0x06, 0x3E,
];

// The real ConnectEx pointer, saved after WSAIoctl returns it
static REAL_CONNECTEX: AtomicUsize = AtomicUsize::new(0);

// ConnectEx function signature
type ConnectExFn = unsafe extern "system" fn(
    s: usize,
    name: *const u8,
    namelen: i32,
    send_buf: *const c_void,
    send_data_len: u32,
    bytes_sent: *mut u32,
    overlapped: *mut c_void,
) -> i32;

// WSAIoctl function address (for inline unhook/rehook)
static WSAIOCTL_ADDR: AtomicUsize = AtomicUsize::new(0);
static mut WSAIOCTL_ORIG: [u8; 14] = [0u8; 14];

type WsaIoctlFn = unsafe extern "system" fn(
    s: usize,
    code: u32,
    in_buf: *const c_void,
    in_len: u32,
    out_buf: *mut c_void,
    out_len: u32,
    bytes_ret: *mut u32,
    overlapped: *mut c_void,
    completion: *const c_void,
) -> i32;

unsafe fn write_hook(target: *mut u8, dest: u64) {
    use windows_sys::Win32::System::Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE};
    let mut old: u32 = 0;
    VirtualProtect(target as _, 14, PAGE_EXECUTE_READWRITE, &mut old);
    target.write(0xFF);
    target.add(1).write(0x25);
    (target.add(2) as *mut u32).write(0u32);
    (target.add(6) as *mut u64).write(dest);
    VirtualProtect(target as _, 14, old, &mut old);
}

unsafe fn restore_wsaioctl(target: *mut u8) {
    use windows_sys::Win32::System::Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE};
    let mut old: u32 = 0;
    VirtualProtect(target as _, 14, PAGE_EXECUTE_READWRITE, &mut old);
    core::ptr::copy_nonoverlapping(core::ptr::addr_of!(WSAIOCTL_ORIG) as *const u8, target, 14);
    VirtualProtect(target as _, 14, old, &mut old);
}

/// Our ConnectEx wrapper: redirects EA ports to 127.0.0.1
unsafe extern "system" fn hooked_connectex(
    s: usize,
    name: *const u8,
    namelen: i32,
    send_buf: *const c_void,
    send_data_len: u32,
    bytes_sent: *mut u32,
    overlapped: *mut c_void,
) -> i32 {
    let real_fn: ConnectExFn = core::mem::transmute(REAL_CONNECTEX.load(Ordering::Relaxed));

    // Milestone-0 transport watch (self-gates on OPENFUT_TRANSPORT_WATCH).
    crate::transport_watch::note_connect("ConnectEx", name, namelen, s);

    // Share the one redirect implementation (v4 + v6) with connect_hook, so ConnectEx
    // dials get the same IPv6 handling as plain connect().
    if let Some((buf, len)) = crate::connect_hook::redirect_if_ea(name, namelen) {
        return real_fn(
            s,
            buf.as_ptr(),
            len,
            send_buf,
            send_data_len,
            bytes_sent,
            overlapped,
        );
    }
    real_fn(
        s,
        name,
        namelen,
        send_buf,
        send_data_len,
        bytes_sent,
        overlapped,
    )
}

/// Our WSAIoctl hook: when ConnectEx is requested, save the real pointer and return ours
pub unsafe extern "system" fn hooked_wsaioctl(
    s: usize,
    code: u32,
    in_buf: *const c_void,
    in_len: u32,
    out_buf: *mut c_void,
    out_len: u32,
    bytes_ret: *mut u32,
    overlapped: *mut c_void,
    completion: *const c_void,
) -> i32 {
    let addr = WSAIOCTL_ADDR.load(Ordering::Relaxed) as *mut u8;

    // Call the real WSAIoctl via unhook/rehook
    restore_wsaioctl(addr);
    let result = {
        let f: WsaIoctlFn = core::mem::transmute(addr);
        f(
            s, code, in_buf, in_len, out_buf, out_len, bytes_ret, overlapped, completion,
        )
    };
    write_hook(addr, hooked_wsaioctl as *const () as u64);

    // If this was a ConnectEx request that succeeded, swap the pointer
    if result == 0 && code == SIO_GET_EXT_FN && in_len == 16 && !in_buf.is_null() {
        let guid = core::slice::from_raw_parts(in_buf as *const u8, 16);
        if guid == CONNECTEX_GUID && out_len >= 8 && !out_buf.is_null() {
            let out_ptr = out_buf as *mut usize;
            let real_addr = *out_ptr;
            if REAL_CONNECTEX
                .compare_exchange(0, real_addr, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                crate::write_log(&format!(
                    "connectex_hook: intercepted ConnectEx @ {real_addr:#x}\n"
                ));
            }
            // Return our hook instead
            *out_ptr = hooked_connectex as *const () as usize;
        }
    }
    result
}

pub unsafe fn install_wsaioctl_hook() -> bool {
    use windows_sys::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress};
    let ws2 = GetModuleHandleA(b"ws2_32.dll\0".as_ptr());
    if ws2.is_null() {
        return false;
    }
    let fn_ptr = match GetProcAddress(ws2, b"WSAIoctl\0".as_ptr()) {
        Some(f) => f as *mut u8,
        None => return false,
    };
    core::ptr::copy_nonoverlapping(
        fn_ptr,
        core::ptr::addr_of_mut!(WSAIOCTL_ORIG) as *mut u8,
        14,
    );
    WSAIOCTL_ADDR.store(fn_ptr as usize, Ordering::Relaxed);
    write_hook(fn_ptr, hooked_wsaioctl as *const () as u64);
    true
}