libinject/

wrappers.rs

use crate::Boolean;
use crate::connections::Connection;
use crate::drcore::log;
use crate::drwrap::Retval;
use crate::socket::{SocketData, WSABufRaw, WSAOverlappedRaw};
use crate::utils::FromBuf;
use crate::utils::socketaddr_from_socket_id;
use crate::{cli, connections, drcore, drwrap, ffi, instrument, pipe, trajectory};
use enum_iterator::Sequence;
use navigator::callstack::CallStack;
use navigator::netgraph::core::{
    Attempt, ConnectAttempt, ConnectReact, React, RecvAttempt, RecvReact, SendAttempt, SendReact,
};
use navigator::netgraph::trajectory::GraphStep;
use navigator::pipe_event::PipeEvent;
use navigator::socket::{HalfDuplexExchange, HalfDuplexResponse, RecvBuffer};
use navigator::trajectory::PlanningStep;
use std::sync::Arc;
use std::{net::SocketAddr, os::raw::c_void};

const SOCKET_ERROR: i32 = -1;
const WSAEHOSTUNREACH: i32 = 10065;
const WSA_IO_PENDING: i32 = 997;
const WSAECONNRESET: i32 = 10054;
const SIO_GET_EXTENSION_FUNCTION_POINTER: u32 = 0xC8000006;
// WSAID_CONNECTEX GUID: {25a207b9-ddf3-4660-8ee9-76e58c74063e}
const WSAID_CONNECTEX: [u8; 16] = [
    0xb9, 0x07, 0xa2, 0x25, 0xf3, 0xdd, 0x60, 0x46, 0x8e, 0xe9, 0x76, 0xe5, 0x8c, 0x74, 0x06, 0x3e,
];

#[derive(Debug, PartialEq, Sequence)]
pub enum Symbols {
    WSAioctl,
    Connect,
    Send,
    WSASend,
    Recv,
    WSARecv,
}

impl Symbols {
    pub fn pre_fn(
        &self,
    ) -> unsafe extern "C" fn(wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
        match self {
            Symbols::WSAioctl => pre_wsaioctl,
            Symbols::Connect => pre_connect,
            Symbols::Send => pre_send,
            Symbols::WSASend => pre_wsasend,
            Symbols::Recv => pre_recv,
            Symbols::WSARecv => pre_wsarecv,
        }
    }
    pub fn post_fn(
        &self,
    ) -> Option<unsafe extern "C" fn(wrapctx: *mut c_void, _user_data: *mut c_void)> {
        match self {
            Symbols::WSAioctl => Some(post_wsaioctl),
            _ => None,
        }
    }
    pub fn to_str(&self) -> &str {
        match self {
            Symbols::WSAioctl => "wsaioctl",
            Symbols::Connect => "connect",
            Symbols::Send => "send",
            Symbols::WSASend => "wsasend",
            Symbols::Recv => "recv",
            Symbols::WSARecv => "wsarecv",
        }
    }
}

/// Pre-callback for WSAIoctl - save output buffer pointer to wrap ConnectEx in post
pub unsafe extern "C" fn pre_wsaioctl(wrapcxt: *mut c_void, user_data: *mut *mut c_void) {
    log("[WSAIoctl] Pre-callback called");

    // Get all the WSAIoctl arguments
    let io_control_code = unsafe { ffi::drwrap_get_arg(wrapcxt, 1) as u32 };
    let lpv_in = unsafe { ffi::drwrap_get_arg(wrapcxt, 2) as *const u8 };
    let cb_in = unsafe { ffi::drwrap_get_arg(wrapcxt, 3) as u32 };
    let lpv_out = unsafe { ffi::drwrap_get_arg(wrapcxt, 4) as *mut usize };

    log(&format!("[WSAIoctl] Control code: 0x{:X}", io_control_code));

    // Check if this is a SIO_GET_EXTENSION_FUNCTION_POINTER request for ConnectEx
    if io_control_code == SIO_GET_EXTENSION_FUNCTION_POINTER {
        if !lpv_in.is_null() && cb_in >= 16 {
            let guid_slice = unsafe { std::slice::from_raw_parts(lpv_in, 16) };

            if guid_slice == &WSAID_CONNECTEX {
                log("[WSAIoctl] Detected ConnectEx request - will wrap in post-callback");

                // Store the output buffer pointer in user_data so we can access it in post-callback
                unsafe {
                    *user_data = lpv_out as *mut c_void;
                }
                // Let the real WSAIoctl proceed to get the real ConnectEx pointer
                return;
            }
        }
    }

    // For all other cases, let WSAIoctl proceed normally
    log("[WSAIoctl] Not a ConnectEx request - letting WSAIoctl proceed");
}

/// Post-callback for WSAIoctl - wrap the real ConnectEx that was returned
pub unsafe extern "C" fn post_wsaioctl(_wrapcxt: *mut c_void, user_data: *mut c_void) {
    if user_data.is_null() {
        // Not a ConnectEx request, nothing to do
        return;
    }

    log("[WSAIoctl] Post-callback called");

    // user_data contains the lpv_out pointer
    let lpv_out = user_data as *mut usize;

    if !lpv_out.is_null() {
        // Read the real ConnectEx pointer that Windows returned
        let real_connectex_ptr = unsafe { *lpv_out };
        log(&format!(
            "[WSAIoctl] Real ConnectEx pointer: 0x{:X}",
            real_connectex_ptr
        ));

        if real_connectex_ptr != 0 {
            // Now wrap the real ConnectEx function
            unsafe {
                match ffi::drwrap_wrap(real_connectex_ptr as *mut u8, Some(pre_connectex), None)
                    .as_bool()
                {
                    true => log(&format!(
                        "[WSAIoctl] Successfully wrapped real ConnectEx @ 0x{:X}",
                        real_connectex_ptr
                    )),
                    false => log(&format!(
                        "[WSAIoctl] FAILED to wrap real ConnectEx @ 0x{:X}",
                        real_connectex_ptr
                    )),
                }
            }
        }
    }
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn pre_connectex(wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
    unsafe {
        let _send_buffer = drwrap::get_arg(wrapctx, 3);
        let send_data_length = drwrap::get_arg(wrapctx, 4) as u32;

        if send_data_length != 0 {
            log(
                "[connectex] send_data_length is not 0! Implementation required to parse the send payload accompanying this connect",
            );
            unimplemented!(
                "[connectex] send_data_length is not 0! Implementation required to parse the send payload accompanying this connect"
            );
        }

        // Required for `skip_call` on 32-bit.
        let stdcall_args_size = if cfg!(target_pointer_width = "32") {
            // 7 args, all ptr sized = 7 * 4 = 28
            28
        } else {
            0
        };

        let (socket_id, socket_addr) = parse_socket(wrapctx, "connectex");
        let graph_step = parse_connect_graph_step(socket_addr, "connectex");
        pipe::send(&PipeEvent::NetGraph(graph_step.clone()));

        // Assume that calls to connectex are using AsyncOverlapped sockets.
        if let GraphStep::Connect(_, ConnectReact { success: true }) = graph_step {
            let socket_data = SocketData::ServerManaged;
            connections::insert(socket_id, socket_addr, socket_data);
            // Allow the connectex call through to the OS, fakenet will handle the socket.
        } else {
            let retval = Retval::Int(false as i32);
            let error_code = WSAEHOSTUNREACH;
            unimplemented!("TODO: handle connection rejection");
            log(&format!(
                "[connectex] Stored error code {} in TLS for WSAGetLastError hook",
                error_code
            ));
            skip_call_helper(wrapctx, retval, "connectex", stdcall_args_size);
        }
    }
}

/// Windows docs for "connect": https://learn.microsoft.com/en-us/windows/win32/api/Winsock2/nf-winsock2-connect
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pre_connect(wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
    unsafe {
        // Required for `skip_call` on 32-bit.
        let stdcall_args_size = if cfg!(target_pointer_width = "32") {
            // socket_id pointer (4) + sockaddr_ptr (4) + int (4)
            12
        } else {
            0
        };

        let (socket_id, socket_addr) = parse_socket(wrapctx, "connect");
        let graph_step = parse_connect_graph_step(socket_addr, "connect");
        pipe::send(&PipeEvent::NetGraph(graph_step.clone()));

        // Assume that calls to connect are using Sync sockets.
        if let GraphStep::Connect(_, ConnectReact { success: true }) = graph_step {
            if cli::get_args()
                .expect("Args parsed in Libinject init")
                .server_managed
            {
                let socket_data = SocketData::ServerManaged;
                connections::insert(socket_id, socket_addr, socket_data);
            } else {
                let socket_data = SocketData::Sync(RecvBuffer::new());
                connections::insert(socket_id, socket_addr, socket_data);
                let retval = Retval::Int(0);
                skip_call_helper(wrapctx, retval, "connect", stdcall_args_size);
            };
        } else {
            let retval = Retval::Int(SOCKET_ERROR);
            skip_call_helper(wrapctx, retval, "connect", stdcall_args_size);
        }
    }
}

unsafe fn connection_of_socket(socket_id: usize, wrapper_name: &str) -> Arc<Connection> {
    unsafe {
        connections::get(socket_id).unwrap_or_else(|| {
            log(&format!("[{wrapper_name}] Binary tried to call {wrapper_name} with a socket that we have not recorded a connection for! Trying to parse the socket address from the socket_id..."));
            match socketaddr_from_socket_id(socket_id) {
                Ok(socket_addr) => {
                    log(&format!("[{wrapper_name}] Successfully parsed socket address from socket_id. Recording this connection."));
                    connections::insert(socket_id, socket_addr, SocketData::ServerManaged);
                    connections::get(socket_id).unwrap()
                },
                Err(e) => {
                    log(&format!("[{wrapper_name} Failed parsing SocketAddr from socket_id (opaque OS handle): {e}"));
                    log(&format!("[{wrapper_name}] If there is only one connection, we will assume the socket_id has changed but the IP address remains the same."));
                    connections::exactly_one().expect("[{wrapper_name}] There is not exactly one connection")
                }
            }
        })
    }
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn pre_wsasend(wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
    unsafe {
        log("[WSASend] Pre-callback called *********");
        let stdcall_args_size = 28;

        let socket_id = drwrap::get_arg(wrapctx, 0) as usize;
        let connection = connection_of_socket(socket_id, "wsasend");

        let buffers_ptr = drwrap::get_arg(wrapctx, 1) as *mut WSABufRaw;
        let buffer_count = drwrap::get_arg(wrapctx, 2) as usize;
        log(&format!("[WSASend] buffer count: {buffer_count}"));

        if buffers_ptr.is_null() {
            panic!("[WSASend] buffer_ptr is null!");
        }

        let overlapped_ptr = drwrap::get_arg(wrapctx, 5) as *mut WSAOverlappedRaw;
        let completion_routine_ptr = drwrap::get_arg(wrapctx, 6);
        log(&format!(
            "[WSASend] overlapper ptr is null?: {}",
            overlapped_ptr.is_null()
        ));
        log(&format!(
            "[WSASend] completion routine ptr is null?: {}",
            completion_routine_ptr.is_null()
        ));

        let mut payloads = Vec::new();

        for i in 0..buffer_count {
            let wsabuf = buffers_ptr.add(i).read();
            let length = wsabuf.len as usize;
            let buf_ptr = wsabuf.buf;

            if !buf_ptr.is_null() && length > 0 {
                let data = std::slice::from_raw_parts(buf_ptr, length);
                payloads.push(data.to_vec());
                let snippet =
                    std::str::from_utf8(&data[..data.len().min(100)]).unwrap_or("<non-utf8>");
                log(&format!("[WSASend] Buffer {i} (len {length}): {snippet}"));
            } else {
                log(&format!("[WSASend] Buffer {i}: null or zero-length"));
            }
        }

        if payloads.len() > 1 {
            unimplemented!("handle multiple WSABUFs arriving in WSASend")
        }

        let request = payloads
            .into_iter()
            .next()
            .expect("At least one WSABUF is expected when WSASend is called");

        log(&format!("[WSASend parsed request: {:?}", request.clone()));

        connections::record_request(connection.socket.id, request.clone());
        let graph_step = parse_send_graph_step(request.clone(), "WSASend");
        pipe::send(&PipeEvent::NetGraph(graph_step.clone()));

        // Match on the socket type.
        match &connection.socket.data {
            SocketData::ServerManaged => {
                // Emit a callstack on WSASend for async sockets.
                capture_and_send_callstack("WSASend");

                // If blocking the send, return SOCKET_ERROR and set WSAECONNRESET; else generate a response to the request and
                // push the response to the socket.
                if let GraphStep::Send(_, SendReact { success: false }) = graph_step {
                    let retval = Retval::Int(SOCKET_ERROR);
                    let error_code = WSAECONNRESET;
                    unimplemented!("TODO: handle rejecting the send");
                    // skip_call_helper(wrapctx, retval, "WSASend", stdcall_args_size);
                } else {
                    let exchange = parse_exchange(request, "WSASend");
                    connection.socket.push_exchange(exchange.clone());
                }
            }
            _ => unimplemented!("Only server managed sockets are implemented for WSASend"),
        }
    }
}

/// Windows docs for "send": https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-send
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pre_send(wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
    unsafe {
        // Opaque handle for the socket (provided by the OS).
        let socket_id = drwrap::get_arg(wrapctx, 0) as usize;
        let payload_ptr = drwrap::get_arg(wrapctx, 1);
        let payload_size = drwrap::get_arg(wrapctx, 2) as usize;

        // Required for `skip_call` on 32-bit.
        let stdcall_args_size = if cfg!(target_pointer_width = "32") {
            // socket_id pointer (4) + payload_ptr (4) + payload_size int (4) + flags int (4)
            16
        } else {
            0
        };

        let connection = connection_of_socket(socket_id, "send");

        let request = match drcore::safe_read(payload_ptr, payload_size) {
            Ok(payload) => {
                log(&format!(
                    "[send] target: {}, payload: {:?}",
                    connection.to_string(),
                    payload
                ));
                payload
            }
            Err(read_err) => {
                panic!(
                    "[send] Failed to read the payload sent by the binary: {:?}",
                    read_err
                );
            }
        };

        connections::record_request(connection.socket.id, request.clone());
        let graph_step = parse_send_graph_step(request.clone(), "send");
        pipe::send(&PipeEvent::NetGraph(graph_step.clone()));

        // If blocking the send, return SOCKET_ERROR; else generate a response to the request and
        // push the response to the socket.
        if let GraphStep::Send(_, SendReact { success: false }) = graph_step {
            let retval = Retval::Int(SOCKET_ERROR);
            skip_call_helper(wrapctx, retval, "send", stdcall_args_size);
        } else {
            let exchange = parse_exchange(request.clone(), "send");
            connection.socket.push_exchange(exchange);
            match connection.socket.data {
                SocketData::Sync(_) => {
                    let retval = Retval::Int(request.len() as i32);
                    skip_call_helper(wrapctx, retval, "send", stdcall_args_size);
                }
                SocketData::ServerManaged => (),
            }
        }
    }
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn pre_wsarecv(_wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
    unsafe {
        log("[WSARecv] Pre-callback called *********");
        //capture_and_send_callstack("WSARecv");
    }
}

/// Windows docs for "recv": https://learn.microsoft.com/en-us/windows/win32/api/winsock/nf-winsock-recv
#[unsafe(no_mangle)]
pub extern "C" fn pre_recv(wrapctx: *mut c_void, _user_data: *mut *mut c_void) {
    unsafe {
        // Opaque handle for the socket (provided by the OS).
        let socket_id = drwrap::get_arg(wrapctx, 0) as usize;

        // Required for `skip_call` on 32-bit.
        let stdcall_args_size = if cfg!(target_pointer_width = "32") {
            // socket_id pointer (4) + payload_ptr (4) + payload_size int (4) + flags int (4)
            16
        } else {
            0
        };

        let connection = connection_of_socket(socket_id, "recv");

        log(&format!(
            "[recv] called for: {}",
            connection.addr.to_string()
        ));
        // Pointer to the recv buffer that the exe will read from.
        let buf_ptr = drwrap::get_arg(wrapctx, 1);
        // Must write <= buf_size to buf_ptr.
        let buf_size = drwrap::get_arg(wrapctx, 2) as usize;

        capture_and_send_callstack("recv");

        match &connection.socket.data {
            SocketData::Sync(recv_buf) => {
                // 100 * 100ms = 10 second countdown timer
                let mut counter = 100;
                while recv_buf.is_empty() {
                    counter -= 1;
                    ffi::dr_sleep(100);
                    if counter <= 0 {
                        // Timeout reached, send a socket error.
                        let retval = Retval::Int(SOCKET_ERROR);
                        skip_call_helper(wrapctx, retval, "recv", stdcall_args_size);
                        return ();
                    }
                }
                let mut exch_slot_guard = recv_buf
                    .slot
                    .lock()
                    .expect("Recv queue mutex should not be poisoned");

                let exch = exch_slot_guard
                    .take()
                    .expect("while loop exits when recv_buf.is_some");

                match &exch.response {
                    HalfDuplexResponse::SocketError => {
                        let graph_step = GraphStep::Recv(
                            RecvAttempt {
                                payload: exch.request.clone(),
                            },
                            RecvReact::SocketError,
                        );
                        let retval = Retval::Int(SOCKET_ERROR);
                        pipe::send(&PipeEvent::NetGraph(graph_step));
                        skip_call_helper(wrapctx, retval, "recv", stdcall_args_size);
                    }
                    HalfDuplexResponse::Payload(response) => {
                        assert!(response.len() <= buf_size);
                        // Write the response payload to the recv buffer.
                        match drcore::safe_write(buf_ptr, response.clone()) {
                            Ok(()) => log("[recv] successfully wrote to recv buffer"),
                            Err(write_err) => {
                                panic!("[recv] Failed writing payload: {:?}", write_err);
                            }
                        };
                        let graph_step = GraphStep::Recv(
                            RecvAttempt {
                                payload: exch.request.clone(),
                            },
                            RecvReact::Payload(response.clone()),
                        );
                        pipe::send(&PipeEvent::NetGraph(graph_step));
                        connections::record_response(socket_id, response.clone());
                        let retval = Retval::Int(response.len() as i32);
                        skip_call_helper(wrapctx, retval, "recv", stdcall_args_size);
                    }
                };
            }
            // Let the `recv` call hit the OS.
            SocketData::ServerManaged => (),
        };
    }
}

fn parse_socket(wrapctx: *mut c_void, wrapper_name: &str) -> (usize, SocketAddr) {
    // Opaque handle for the socket (provided by the OS).
    let socket_id = drwrap::get_arg(wrapctx, 0) as usize;
    // Pointer to the sockaddr struct.
    let sockaddr_ptr = drwrap::get_arg(wrapctx, 1);
    // Size of the sockaddr struct in bytes.
    let sockaddr_size = drwrap::get_arg(wrapctx, 2) as usize;

    if sockaddr_ptr.is_null() || sockaddr_size == 0 {
        todo!("[{wrapper_name}] Couldn't parse an IP address during Connect, ticket #8")
    }

    // Try to parse a std::net::SocketAddr from the pointer.
    let socket_addr = match drcore::safe_read(sockaddr_ptr, sockaddr_size) {
        Ok(buf) => SocketAddr::from_buf(buf),
        Err(read_err) => {
            log(&format!(
                "[{wrapper_name}] Failed reading sockaddr:\n {:?}",
                read_err
            ));
            None
        }
    };

    let socket_addr = socket_addr.unwrap_or_else(|| {
        todo!("[{wrapper_name}] Couldn't parse an IP address during Connect, ticket #8")
    });
    log(&format!(
        "[{wrapper_name}] attempt to connect to {}",
        socket_addr.to_string()
    ));
    (socket_id, socket_addr)
}

fn parse_connect_graph_step(socket_addr: SocketAddr, wrapper_name: &str) -> GraphStep {
    let observed_attempt = Attempt::Connect(ConnectAttempt { addr: socket_addr });
    match parse_planned_step(observed_attempt, wrapper_name) {
        PlanningStep::Graph(graph_step) => graph_step,
        _ => unreachable!("PlanningStep will be a GraphStep::Connect"),
    }
}

fn parse_exchange(request: Vec<u8>, wrapper_name: &str) -> HalfDuplexExchange {
    let observed_attempt = Attempt::Recv(RecvAttempt {
        payload: request.clone(),
    });
    let planned_step = parse_planned_step(observed_attempt, wrapper_name);

    match planned_step {
        PlanningStep::Graph(GraphStep::Recv(recv_attempt, recv_react)) => match recv_react {
            RecvReact::SocketError => HalfDuplexExchange {
                request: recv_attempt.payload,
                response: HalfDuplexResponse::SocketError,
            },
            RecvReact::Payload(p) => HalfDuplexExchange {
                request: recv_attempt.payload,
                response: HalfDuplexResponse::Payload(p),
            },
        },
        PlanningStep::Fuzz(recv_attempt) => HalfDuplexExchange {
            request: recv_attempt.payload.clone(),
            response: HalfDuplexResponse::Payload(
                navigator::protocol::get().fuzz_response(&recv_attempt.payload, Some(log)),
            ),
        },
        _ => unreachable!(
            "GraphStep will be the Recv variant when an Attempt::Recv is the argument to parse_planning_step"
        ),
    }
}

fn parse_send_graph_step(request: Vec<u8>, wrapper_name: &str) -> GraphStep {
    let observed_attempt = Attempt::Send(SendAttempt { payload: request });
    let planned_step = parse_planned_step(observed_attempt, wrapper_name);
    match planned_step {
        PlanningStep::Graph(graph_step @ GraphStep::Send(_, _)) => graph_step,
        _ => unreachable!(
            "Planned step will be Send when parse_planned_step is called with an Attempt::Send"
        ),
    }
}

unsafe fn capture_and_send_callstack(wrapper_name: &str) {
    unsafe {
        // Capture in-module call stack for this Recv attempt, if available.
        let drctx = ffi::dr_get_current_drcontext();
        let callstack = instrument::stack_frames(drctx).map(|frames| {
            let (base, size) = instrument::module_bounds();
            CallStack::from_raw_frames(base, size, frames)
        });
        log(&format!(
            "[{wrapper_name}] callstack recorded: {}",
            if callstack.is_some() { "Some" } else { "None" }
        ));
        if callstack.is_some() {
            pipe::send(&PipeEvent::CallStack(callstack.clone().unwrap()));
        }
    }
}

unsafe fn skip_call_helper(
    wrapctx: *mut c_void,
    retval: Retval,
    wrapper_name: &str,
    args_size: usize,
) {
    unsafe {
        match drwrap::skip_call(wrapctx, retval.clone(), args_size) {
            true => {
                log(&format!(
                    "[{wrapper_name}] Successfully returned {:?}",
                    retval
                ));
            }
            false => panic!("[{wrapper_name}] skip_call failed"),
        };
    }
}

/// Check the current trajectory for a planned reaction to the `observed_attempt`. If not present,
/// create the default `PlanningStep` for the `observed_attempt`.
fn parse_planned_step(observed_attempt: Attempt, wrapper_name: &str) -> PlanningStep {
    // Get the reaction from the configured trajectory (or default) and create the current step.
    match trajectory::next(&observed_attempt) {
        Ok(Some(current_step)) => {
            log(&format!(
                "[{wrapper_name}] following configured trajectory, current step: {:?}",
                current_step
            ));
            current_step
        }
        // If the end of the trajectory has been reached, or if we are off trajectory, do the default behaviour.
        Ok(None) | Err(_) => {
            let default_react = React::default(&observed_attempt, Some(log));
            let default_step =
                PlanningStep::Graph(GraphStep::new(observed_attempt, default_react).expect(
                    "React::default ensures that the default_react matches the attempt type",
                ));
            default_step
        }
    }
}