|
|
|
@@ -33,6 +33,31 @@ static CONNECT_ADDR: AtomicUsize = AtomicUsize::new(0);
|
|
|
|
|
// Original 14 bytes saved before we overwrite them
|
|
|
|
|
static mut ORIGINAL_BYTES: [u8; 14] = [0u8; 14];
|
|
|
|
|
|
|
|
|
|
/// Restores the real WinSock call's thread-local last error after detour repair,
|
|
|
|
|
/// logging, and other instrumentation have run. Callers inspect this value after
|
|
|
|
|
/// `SOCKET_ERROR`; leaking a logger/VirtualProtect error changes connect semantics.
|
|
|
|
|
struct WsaLastErrorGuard(i32);
|
|
|
|
|
|
|
|
|
|
impl WsaLastErrorGuard {
|
|
|
|
|
unsafe fn capture() -> Self {
|
|
|
|
|
use windows_sys::Win32::Networking::WinSock::WSAGetLastError;
|
|
|
|
|
Self(WSAGetLastError())
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn value(&self) -> i32 {
|
|
|
|
|
self.0
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl Drop for WsaLastErrorGuard {
|
|
|
|
|
fn drop(&mut self) {
|
|
|
|
|
unsafe {
|
|
|
|
|
use windows_sys::Win32::Networking::WinSock::WSASetLastError;
|
|
|
|
|
WSASetLastError(self.0);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// For WSAConnect IAT fallback
|
|
|
|
|
type WsaConnectFn = unsafe extern "system" fn(
|
|
|
|
|
s: usize,
|
|
|
|
@@ -191,6 +216,8 @@ pub unsafe extern "system" fn hooked_connect(s: usize, name: *const u8, namelen:
|
|
|
|
|
core::mem::transmute(addr);
|
|
|
|
|
f(s, buf.as_ptr(), len)
|
|
|
|
|
};
|
|
|
|
|
// Named binding held until `return r`: its Drop restores the WSA error after `write_hook`.
|
|
|
|
|
let _last_error = WsaLastErrorGuard::capture();
|
|
|
|
|
write_hook(addr, hooked_connect as *const () as u64);
|
|
|
|
|
return r;
|
|
|
|
|
} else {
|
|
|
|
@@ -202,17 +229,15 @@ pub unsafe extern "system" fn hooked_connect(s: usize, name: *const u8, namelen:
|
|
|
|
|
let f: unsafe extern "system" fn(usize, *const u8, i32) -> i32 = core::mem::transmute(addr);
|
|
|
|
|
f(s, call_name, call_len)
|
|
|
|
|
};
|
|
|
|
|
let last_error = WsaLastErrorGuard::capture();
|
|
|
|
|
write_hook(addr, hooked_connect as *const () as u64);
|
|
|
|
|
if namelen >= 8 {
|
|
|
|
|
let sa = &*(call_name as *const SockaddrIn);
|
|
|
|
|
if sa.sin_family == AF_INET {
|
|
|
|
|
let err = if r != 0 {
|
|
|
|
|
use windows_sys::Win32::Networking::WinSock::WSAGetLastError;
|
|
|
|
|
WSAGetLastError()
|
|
|
|
|
} else {
|
|
|
|
|
0
|
|
|
|
|
};
|
|
|
|
|
crate::write_log(&format!("connect_hook: result={r} wsa_err={err}\n"));
|
|
|
|
|
let logged_error = if r != 0 { last_error.value() } else { 0 };
|
|
|
|
|
crate::write_log(&format!(
|
|
|
|
|
"connect_hook: result={r} wsa_err={logged_error}\n"
|
|
|
|
|
));
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
r
|
|
|
|
@@ -260,3 +285,24 @@ pub unsafe fn install_inline_connect_hook() -> bool {
|
|
|
|
|
write_hook(connect_fn, hooked_connect as *const () as u64);
|
|
|
|
|
true
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
|
mod tests {
|
|
|
|
|
use super::WsaLastErrorGuard;
|
|
|
|
|
use windows_sys::Win32::Networking::WinSock::{
|
|
|
|
|
WSAGetLastError, WSASetLastError, WSAEWOULDBLOCK,
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn restores_winsock_last_error_after_instrumentation() {
|
|
|
|
|
unsafe {
|
|
|
|
|
WSASetLastError(WSAEWOULDBLOCK);
|
|
|
|
|
{
|
|
|
|
|
let guard = WsaLastErrorGuard::capture();
|
|
|
|
|
assert_eq!(guard.value(), WSAEWOULDBLOCK);
|
|
|
|
|
WSASetLastError(0);
|
|
|
|
|
}
|
|
|
|
|
assert_eq!(WSAGetLastError(), WSAEWOULDBLOCK);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|