use std::io::{Read, Write};
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use wireshift_core::buffer::{Buffer, Submitted};
use wireshift_core::op::CompletionPayload;
use wireshift_core::{Error, Result};
macro_rules! retry_eintr {
($expr:expr) => {
loop {
match $expr {
Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
result => break result,
}
}
};
}
pub(crate) fn execute_connect(
addr: SocketAddr,
timeout: Option<Duration>,
canceled: Arc<AtomicBool>,
) -> Result<CompletionPayload> {
let stream = TcpStream::connect_timeout(&addr, timeout.unwrap_or(Duration::from_secs(5)))
.map_err(|error| {
Error::io(
"connect failed",
error,
"ensure the remote host is reachable",
)
})?;
if canceled.load(Ordering::Relaxed) {
let _ = stream.shutdown(Shutdown::Both);
return Err(Error::canceled(
"connect canceled after socket creation",
"avoid canceling the request after the socket connects",
));
}
Ok(CompletionPayload::Stream(stream))
}
pub(crate) fn execute_accept(
listener: TcpListener,
canceled: Arc<AtomicBool>,
timeout: Option<Duration>,
) -> Result<CompletionPayload> {
listener.set_nonblocking(true).map_err(|error| {
Error::io(
"accept setup failed",
error,
"ensure the listener remains valid",
)
})?;
let deadline = Instant::now() + timeout.unwrap_or(Duration::from_secs(5));
loop {
if canceled.load(Ordering::Relaxed) {
return Err(Error::canceled(
"accept canceled before a peer connected",
"avoid canceling the request while a connection is expected",
));
}
match retry_eintr!(listener.accept()) {
Ok((stream, _)) => return Ok(CompletionPayload::Stream(stream)),
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(Error::timeout(
"accept timed out waiting for an inbound peer",
"connect a client before waiting on the accept request",
));
}
wait_for_listener_readiness(&listener, remaining);
}
Err(error) => {
return Err(Error::io(
"accept failed",
error,
"ensure the listener is bound and still accepting connections",
));
}
}
}
}
fn wait_for_listener_readiness(listener: &TcpListener, remaining: Duration) {
use std::os::unix::io::AsRawFd;
let wait = remaining.min(Duration::from_millis(10));
let wait_ms = i32::try_from(wait.as_millis()).unwrap_or(i32::MAX);
let mut pfd = libc::pollfd {
fd: listener.as_raw_fd(),
events: libc::POLLIN,
revents: 0,
};
unsafe {
libc::poll(std::ptr::addr_of_mut!(pfd), 1, wait_ms);
}
}
fn apply_read_timeout(stream: &TcpStream, timeout: Option<Duration>) -> Result<()> {
if let Some(t) = timeout {
stream.set_read_timeout(Some(t)).map_err(|error| {
Error::io(
"recv set_read_timeout failed",
error,
"ensure the socket is valid and the read timeout is a non-zero duration",
)
})?;
}
Ok(())
}
pub(crate) fn execute_send(
mut stream: TcpStream,
buffer: Buffer<Submitted>,
) -> Result<CompletionPayload> {
let bytes = retry_eintr!(stream.write(buffer.backend_ref())).map_err(|error| {
Error::io(
"send failed",
error,
"ensure the TCP stream remains writable while sending data",
)
})?;
Ok(CompletionPayload::Bytes(bytes))
}
pub(crate) fn execute_recv(
mut stream: TcpStream,
mut buffer: Buffer<Submitted>,
canceled: Arc<AtomicBool>,
timeout: Option<Duration>,
) -> Result<CompletionPayload> {
apply_read_timeout(&stream, timeout)?;
if canceled.load(Ordering::Relaxed) {
return Err(Error::canceled(
"recv canceled before data transfer started",
"avoid canceling the request before the backend starts receiving data",
));
}
let bytes = retry_eintr!(stream.read(buffer.backend_mut())).map_err(|error| {
Error::io(
"recv failed",
error,
"ensure the TCP stream remains readable while receiving data",
)
})?;
let completed = buffer.into_completed(bytes)?;
Ok(CompletionPayload::Recv {
stream,
buffer: completed,
bytes,
})
}
#[cfg(test)]
mod net_ops_tests {
use super::{apply_read_timeout, execute_accept};
use std::net::{TcpListener, TcpStream};
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::{Duration, Instant};
use wireshift_core::op::CompletionPayload;
#[test]
fn apply_read_timeout_surfaces_invalid_zero_duration() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let client = TcpStream::connect(addr).unwrap();
let result = apply_read_timeout(&client, Some(Duration::from_secs(0)));
let err = result.expect_err("zero-duration read timeout must be rejected loudly");
assert!(
err.to_string().to_lowercase().contains("timeout"),
"error must identify the read-timeout configuration failure, got: {err}"
);
}
#[test]
fn apply_read_timeout_accepts_valid_duration_and_none() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let client = TcpStream::connect(addr).unwrap();
apply_read_timeout(&client, Some(Duration::from_millis(50)))
.expect("a valid read timeout must be accepted");
apply_read_timeout(&client, None).expect("no timeout must be a no-op");
}
#[test]
fn accept_wakes_on_connection_without_fixed_sleep_latency() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let canceled = Arc::new(AtomicBool::new(false));
let connector = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(25));
let connect_at = Instant::now();
let stream = TcpStream::connect(addr).unwrap();
(connect_at, stream) });
let payload = execute_accept(listener, canceled, Some(Duration::from_secs(2)))
.expect("accept must succeed once a peer connects");
let accept_done = Instant::now();
let (connect_at, _client) = connector.join().unwrap();
assert!(
matches!(payload, CompletionPayload::Stream(_)),
"accept must yield the accepted stream"
);
let latency = accept_done.saturating_duration_since(connect_at);
assert!(
latency < Duration::from_millis(5),
"accept must wake on listener readiness (sub-quantum), not after a fixed 10ms sleep; \
connect->accept latency was {latency:?}"
);
}
}