fix(hook): preserve WinSock connect errors
This commit is contained in:
@@ -12,6 +12,6 @@ fn main() {
|
||||
{
|
||||
let definition =
|
||||
PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap()).join("version.def");
|
||||
println!("cargo:rustc-link-arg={}", definition.display());
|
||||
println!("cargo:rustc-cdylib-link-arg={}", definition.display());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
#![cfg(windows)]
|
||||
#![allow(dead_code)]
|
||||
|
||||
// Compile the production connect hook directly into an executable test target.
|
||||
// The hook crate itself is a cdylib, whose unit-test artifact remains a DLL and
|
||||
// therefore cannot be executed by the native Windows test runner.
|
||||
fn write_log(_: &str) {}
|
||||
|
||||
#[path = "../src/connect_hook.rs"]
|
||||
mod connect_hook;
|
||||
Reference in New Issue
Block a user