wireshift_fallback/

net_ops.rs

//! Network operation implementations for the fallback backend.

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",
                    ));
                }
                // Event-driven wait on listener readiness instead of a fixed
                // 10ms sleep. The old sleep added up to a full 10ms of latency
                // to every connection that arrived mid-quantum (Law 7); poll(2)
                // returns the instant the listening socket has a pending
                // connection, so accept is serviced in sub-millisecond time.
                // The timeout is capped at a short quantum so cancellation and
                // the overall deadline are still re-checked promptly, and it
                // never exceeds the time left until `deadline`.
                wait_for_listener_readiness(&listener, remaining);
            }
            Err(error) => {
                return Err(Error::io(
                    "accept failed",
                    error,
                    "ensure the listener is bound and still accepting connections",
                ));
            }
        }
    }
}

/// Block until the listening socket has a pending connection or the (capped)
/// timeout elapses, using `poll(2)` for event-driven readiness.
///
/// The poll timeout is capped at a short quantum so the accept loop still
/// re-checks cancellation and the overall deadline promptly, and never exceeds
/// the caller's remaining budget. A pending connection makes the listener
/// readable immediately, so `poll` returns without waiting the full quantum -
/// eliminating the fixed-sleep connection latency (Law 7). Spurious wakeups or
/// `poll` errors (e.g. `EINTR`) simply return; the caller re-attempts `accept`,
/// which surfaces any genuine listener error loudly.
fn wait_for_listener_readiness(listener: &TcpListener, remaining: Duration) {
    use std::os::unix::io::AsRawFd;

    // Cap at 10ms to preserve the previous cancellation-check cadence, but never
    // wait longer than the time left before the deadline.
    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,
    };
    // SAFETY: `pfd` is a single, live, correctly-initialized pollfd; `poll`
    // reads `nfds` entries and writes `revents`. The listener outlives the call.
    unsafe {
        libc::poll(std::ptr::addr_of_mut!(pfd), 1, wait_ms);
    }
}

/// Apply an optional read timeout to `stream`, failing CLOSED on error.
///
/// The previous `let _ = stream.set_read_timeout(..)` discarded a configuration
/// failure and then performed a read with NO timeout - which can hang the
/// worker thread indefinitely (Law 10). Propagating the error surfaces the
/// misconfiguration loudly instead of silently dropping the timeout.
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() {
        // A zero-duration read timeout is invalid at the OS level; the previous
        // `let _ = set_read_timeout(..)` swallowed that error and left the read
        // with no timeout (a potential hang). The fix must surface it (Law 10).
        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() {
        // Guard against over-failing: a valid timeout and `None` both succeed.
        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() {
        // Start accept BEFORE any peer connects so it parks in the readiness
        // wait (WouldBlock -> poll). A peer then connects after a delay; the
        // event-driven poll must return the instant the listener is readable,
        // so the connect->accept-return latency is far below the old fixed 10ms
        // sleep quantum. Asserting < 5ms discriminates the fixed-sleep bug while
        // tolerating normal scheduler jitter.
        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) // hold the client open past accept
        });

        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:?}"
        );
    }
}