zlayer-agent 0.14.0

//! Runtime-agnostic userspace L4 forwarder for first-class host-shared overlay
//! IPs: bind `<overlay_ip>:<port>` and splice each connection to the
//! container's local delivery address. Used by the macOS Seatbelt, native-VZ,
//! and libkrun runtimes after overlayd adds the per-container `utun` alias.
//!
//! Unlike `macos_vz_linux`'s vsock forwarders, this module splices over plain
//! kernel sockets: a `TcpStream::connect(target)` per inbound TCP connection,
//! and a per-client `UdpSocket::connect(target)` for UDP. That makes the
//! transport OS-independent, so the TCP/UDP logic compiles and unit-tests on
//! Linux even though its only production callers are macOS runtimes.
//
// The production callers (Seatbelt / native-VZ / libkrun overlay attach) are
// macOS-only, so on a non-macOS build the forwarders have no caller and would
// trip `dead_code`. They are NOT incomplete: macOS exercises every path, and
// the unit tests below drive the TCP path on Linux — so the lint is only
// suppressed on non-macOS, non-test builds, where the code is genuinely unused.
#![cfg_attr(not(any(test, target_os = "macos")), allow(dead_code))]
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::task::AbortHandle;
use zlayer_spec::PortProtocol;

/// Idle timeout after which a UDP client session (its dedicated upstream socket
/// and relay task) is reclaimed. Matches the vsock UDP forwarder's 60s window.
const UDP_SESSION_TIMEOUT: Duration = Duration::from_secs(60);

/// Largest single UDP datagram relayed in either direction. Matches the maximum
/// IPv4/IPv6 UDP payload, so no legitimate datagram is truncated; also bounds a
/// single allocation on read.
const MAX_UDP_DATAGRAM: usize = 65_535;

/// Spawn a forwarder binding `bind` and relaying to `target` for `proto`.
///
/// Returns an [`AbortHandle`]; abort it on container teardown to stop the
/// listener and all in-flight splices. The bind retries briefly to absorb the
/// race where the `utun` alias is not yet kernel-visible (`EADDRNOTAVAIL`).
#[must_use]
pub fn spawn_overlay_forward(
    bind: SocketAddr,
    target: SocketAddr,
    proto: PortProtocol,
) -> AbortHandle {
    tokio::spawn(async move {
        match proto {
            PortProtocol::Udp => run_udp(bind, target).await,
            // Tcp (and any future default) splice over a plain TCP stream.
            PortProtocol::Tcp => run_tcp(bind, target).await,
        }
    })
    .abort_handle()
}

/// Bind a TCP listener on `bind`, retrying up to ~10 times at 100ms intervals to
/// absorb the window where overlayd's freshly-added `utun` alias is not yet
/// visible to `bind(2)` (`EADDRNOTAVAIL`).
async fn bind_tcp_retry(bind: SocketAddr) -> std::io::Result<tokio::net::TcpListener> {
    let mut last_err = None;
    for _ in 0..10 {
        match tokio::net::TcpListener::bind(bind).await {
            Ok(l) => return Ok(l),
            Err(e) => {
                last_err = Some(e);
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
        }
    }
    // At least one attempt ran, so `last_err` is always populated.
    Err(last_err.unwrap_or_else(|| {
        std::io::Error::new(
            std::io::ErrorKind::AddrNotAvailable,
            "tcp bind retries exhausted",
        )
    }))
}

/// Bind a UDP socket on `bind`, retrying up to ~10 times at 100ms intervals to
/// absorb the same `utun`-alias-visibility race as [`bind_tcp_retry`]
/// (`EADDRNOTAVAIL`).
async fn bind_udp_retry(bind: SocketAddr) -> std::io::Result<tokio::net::UdpSocket> {
    let mut last_err = None;
    for _ in 0..10 {
        match tokio::net::UdpSocket::bind(bind).await {
            Ok(s) => return Ok(s),
            Err(e) => {
                last_err = Some(e);
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
        }
    }
    // At least one attempt ran, so `last_err` is always populated.
    Err(last_err.unwrap_or_else(|| {
        std::io::Error::new(
            std::io::ErrorKind::AddrNotAvailable,
            "udp bind retries exhausted",
        )
    }))
}

/// Accept loop: bind `bind`, then for each inbound connection open a fresh
/// `TcpStream::connect(target)` and bidirectionally copy bytes between the two.
/// Best-effort: bind/accept/connect errors are logged, not fatal.
async fn run_tcp(bind: SocketAddr, target: SocketAddr) {
    let listener = match bind_tcp_retry(bind).await {
        Ok(l) => l,
        Err(e) => {
            tracing::warn!(%bind, %target, error = %e, "overlay forwarder: TCP bind failed");
            return;
        }
    };
    tracing::info!(%bind, %target, "overlay forwarder: TCP listener up");
    loop {
        match listener.accept().await {
            Ok((mut inbound, _peer)) => {
                tokio::spawn(async move {
                    match tokio::net::TcpStream::connect(target).await {
                        Ok(mut outbound) => {
                            let _ =
                                tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await;
                        }
                        Err(e) => tracing::debug!(
                            %target,
                            error = %e,
                            "overlay forwarder: connect to delivery target failed"
                        ),
                    }
                });
            }
            Err(e) => {
                tracing::debug!(error = %e, "overlay forwarder: accept failed");
            }
        }
    }
}

/// A live per-client UDP relay session: the dedicated upstream socket connected
/// to `target`, plus the last-activity timestamp used by the idle sweep.
struct UdpSession {
    /// Upstream socket bound to an ephemeral local port and `connect`ed to the
    /// delivery target, so its reader task only ever sees target traffic.
    upstream: Arc<tokio::net::UdpSocket>,
    /// Last time this client sent OR received a datagram (for idle expiry).
    last_activity: tokio::time::Instant,
}

/// Bidirectional UDP relay: bind one socket on `bind`; the first datagram from a
/// given client address opens a dedicated upstream `UdpSocket::connect(target)`
/// and a reader task that pumps target replies back to that client. Subsequent
/// datagrams from a known client reuse its upstream socket. Idle sessions are
/// reclaimed after [`UDP_SESSION_TIMEOUT`]. Best-effort: bind/relay errors are
/// logged, not fatal.
async fn run_udp(bind: SocketAddr, target: SocketAddr) {
    // The bind can hit the same alias-visibility race as TCP; retry briefly.
    let inbound = match bind_udp_retry(bind).await {
        Ok(s) => Arc::new(s),
        Err(e) => {
            tracing::warn!(%bind, %target, error = %e, "overlay forwarder: UDP bind failed");
            return;
        }
    };
    tracing::info!(%bind, %target, "overlay forwarder: UDP relay up");

    // Sessions keyed by client address, shared with each upstream reader task
    // (which bumps `last_activity` on the way back) and the idle sweep.
    let sessions: Arc<Mutex<HashMap<SocketAddr, UdpSession>>> =
        Arc::new(Mutex::new(HashMap::new()));

    // The idle sweep runs inline via `select!` (below) rather than as a separate
    // task, so the whole relay — inbound socket, session map, and sweep — lives
    // and dies with this one future. Aborting the [`AbortHandle`] returned by
    // `spawn_overlay_forward` therefore tears everything down at once; no child
    // sweeper task can outlive it.
    let mut sweep = tokio::time::interval(Duration::from_secs(10));
    let mut buf = vec![0u8; MAX_UDP_DATAGRAM];
    loop {
        let (len, client_addr) = tokio::select! {
            _ = sweep.tick() => {
                let now = tokio::time::Instant::now();
                let mut map = sessions.lock().await;
                map.retain(|_, s| now.duration_since(s.last_activity) < UDP_SESSION_TIMEOUT);
                continue;
            }
            r = inbound.recv_from(&mut buf) => match r {
                Ok(v) => v,
                Err(e) => {
                    tracing::debug!(error = %e, "overlay forwarder: UDP recv failed");
                    continue;
                }
            },
        };
        let datagram = &buf[..len];

        // Fast path: an existing session already owns an upstream socket.
        {
            let mut map = sessions.lock().await;
            if let Some(s) = map.get_mut(&client_addr) {
                s.last_activity = tokio::time::Instant::now();
                if let Err(e) = s.upstream.send(datagram).await {
                    tracing::debug!(
                        client = %client_addr,
                        %target,
                        error = %e,
                        "overlay forwarder: upstream send failed; dropping session"
                    );
                    map.remove(&client_addr);
                }
                continue;
            }
        }

        // New client: bind an ephemeral upstream socket and connect it to the
        // target, then spawn a reader that pumps target replies back out the
        // shared inbound socket to this client.
        let upstream = match new_udp_upstream(target).await {
            Ok(s) => Arc::new(s),
            Err(e) => {
                tracing::debug!(
                    client = %client_addr,
                    %target,
                    error = %e,
                    "overlay forwarder: failed to open upstream UDP socket"
                );
                continue;
            }
        };

        // Spawn the guest->client reader for this session.
        spawn_udp_reader(
            Arc::clone(&upstream),
            Arc::clone(&inbound),
            Arc::clone(&sessions),
            client_addr,
        );

        // Deliver the first datagram, then record the session.
        if let Err(e) = upstream.send(datagram).await {
            tracing::debug!(
                client = %client_addr,
                %target,
                error = %e,
                "overlay forwarder: first upstream send failed"
            );
            continue;
        }
        let mut map = sessions.lock().await;
        map.insert(
            client_addr,
            UdpSession {
                upstream,
                last_activity: tokio::time::Instant::now(),
            },
        );
    }
}

/// Spawn the guest->client reader task for a single UDP session: pump target
/// replies back out the shared `inbound` socket to `client_addr`, bumping the
/// session's `last_activity` on each datagram, until the upstream socket
/// errors/closes or idles past [`UDP_SESSION_TIMEOUT`].
fn spawn_udp_reader(
    upstream: Arc<tokio::net::UdpSocket>,
    inbound: Arc<tokio::net::UdpSocket>,
    sessions: Arc<Mutex<HashMap<SocketAddr, UdpSession>>>,
    client_addr: SocketAddr,
) {
    tokio::spawn(async move {
        let mut rbuf = vec![0u8; MAX_UDP_DATAGRAM];
        loop {
            // Bound the reader's lifetime by the idle window so it always
            // self-terminates — even if the parent forward future is aborted
            // at teardown (which can't reach in to cancel this child task).
            // Socket closed/error, or idle past the session window: stop.
            let Ok(Ok(n)) =
                tokio::time::timeout(UDP_SESSION_TIMEOUT, upstream.recv(&mut rbuf)).await
            else {
                break;
            };
            if inbound.send_to(&rbuf[..n], client_addr).await.is_err() {
                break;
            }
            // Bump activity so a chatty server keeps the session alive.
            let mut map = sessions.lock().await;
            match map.get_mut(&client_addr) {
                Some(s) => s.last_activity = tokio::time::Instant::now(),
                // Session was reclaimed out from under us; exit.
                None => break,
            }
        }
    });
}

/// Bind an ephemeral upstream UDP socket on the unspecified address matching the
/// target's address family and `connect` it to `target`, so its `recv`/`send`
/// only ever touch target traffic.
async fn new_udp_upstream(target: SocketAddr) -> std::io::Result<tokio::net::UdpSocket> {
    let local: SocketAddr = if target.is_ipv6() {
        (std::net::Ipv6Addr::UNSPECIFIED, 0).into()
    } else {
        (std::net::Ipv4Addr::UNSPECIFIED, 0).into()
    };
    let sock = tokio::net::UdpSocket::bind(local).await?;
    sock.connect(target).await?;
    Ok(sock)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{Ipv4Addr, SocketAddr};
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    /// Spawn a TCP echo server on `127.0.0.1:0` and return its bound address.
    async fn spawn_tcp_echo() -> SocketAddr {
        let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
            .await
            .expect("bind echo");
        let addr = listener.local_addr().expect("echo addr");
        tokio::spawn(async move {
            loop {
                let Ok((mut sock, _)) = listener.accept().await else {
                    break;
                };
                tokio::spawn(async move {
                    let mut buf = [0u8; 1024];
                    loop {
                        match sock.read(&mut buf).await {
                            Ok(0) | Err(_) => break,
                            Ok(n) => {
                                if sock.write_all(&buf[..n]).await.is_err() {
                                    break;
                                }
                            }
                        }
                    }
                });
            }
        });
        addr
    }

    #[tokio::test]
    async fn tcp_forward_round_trips_bytes() {
        let echo_addr = spawn_tcp_echo().await;

        // Bind the forwarder's listener to an ephemeral 127.0.0.1 port. We bind
        // first to learn the port, then hand the address to the forwarder, which
        // re-binds it via its retry loop (the kernel lets us re-bind the freed
        // ephemeral port). Both bind+target are loopback so this runs on Linux CI.
        let probe = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
            .await
            .expect("probe bind");
        let bind_addr = probe.local_addr().expect("probe addr");
        drop(probe);

        let handle = spawn_overlay_forward(bind_addr, echo_addr, PortProtocol::Tcp);

        // Connect to the forwarder, retrying briefly while its listener comes up.
        let mut client = None;
        for _ in 0..50 {
            if let Ok(c) = tokio::net::TcpStream::connect(bind_addr).await {
                client = Some(c);
                break;
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
        let mut client = client.expect("connect to forwarder");

        let payload = b"hello-overlay-forward";
        client.write_all(payload).await.expect("write");
        client.flush().await.expect("flush");

        let mut got = vec![0u8; payload.len()];
        client.read_exact(&mut got).await.expect("read echo");
        assert_eq!(&got, payload);

        handle.abort();
    }
}