zlayer-overlay 0.14.2

//! Local TCP listener + bidirectional pump for `--forward` mappings.
//!
//! A `--forward LOCAL_IP:PORT=OVERLAY_IP:PORT` mapping parses into a
//! [`ForwardSpec`]. [`serve_forward`] then binds a local
//! [`tokio::net::TcpListener`] and, for every accepted connection, opens a
//! tunneled [`TcpConn`] to the overlay peer via an injected [`TunnelDialer`]
//! and pumps bytes in both directions until either side closes.
//!
//! The [`TunnelDialer`] seam is why *one* pump serves both edge transports:
//!
//! - **netstack mode** wraps the smoltcp netstack connect path (Wave-2),
//! - **device mode** "degrades to a plain proxy" — a passthrough dialer that
//!   `connect`s a real [`tokio::net::TcpStream`] to the overlay address — yet
//!   reuses the exact same [`serve_forward`]. DRY: the transport differs, the
//!   accept/pump plumbing does not.

use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::Arc;

use async_trait::async_trait;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

use crate::edge::{EdgeError, TcpConn};

/// Chunk size for host→overlay reads. Each read produces one `Vec<u8>` sent
/// across the tunnel; 64 KiB balances syscall count against per-chunk alloc.
const PUMP_CHUNK_BYTES: usize = 64 * 1024;

/// One `--forward` mapping: bind `local`, tunnel each connection to `remote`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ForwardSpec {
    /// Host-side address the edge client binds and listens on.
    pub local: SocketAddr,
    /// Overlay-side address every accepted connection is tunneled to.
    pub remote: SocketAddr,
}

impl FromStr for ForwardSpec {
    type Err = EdgeError;

    /// Parse `LOCAL_IP:PORT=OVERLAY_IP:PORT`.
    ///
    /// # Errors
    ///
    /// Returns [`EdgeError::Forward`] when the `=` separator is missing, when
    /// either side is not a valid `IP:PORT` [`SocketAddr`] (IPv6 must use the
    /// bracketed `[addr]:port` form), or when either port is `0`.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (local_raw, remote_raw) = s.split_once('=').ok_or_else(|| {
            EdgeError::Forward(format!(
                "forward spec {s:?} is missing the '=' separator \
                 (expected LOCAL_IP:PORT=OVERLAY_IP:PORT)"
            ))
        })?;
        let local = parse_side(local_raw, "local")?;
        let remote = parse_side(remote_raw, "remote")?;
        Ok(Self { local, remote })
    }
}

/// Parse one `IP:PORT` side of a forward spec, rejecting port `0`.
fn parse_side(raw: &str, which: &str) -> Result<SocketAddr, EdgeError> {
    let trimmed = raw.trim();
    let addr = SocketAddr::from_str(trimmed).map_err(|e| {
        EdgeError::Forward(format!(
            "forward spec {which} side {trimmed:?} is not a valid IP:PORT address ({e}); \
             IPv6 must be bracketed, e.g. [::1]:8080"
        ))
    })?;
    if addr.port() == 0 {
        return Err(EdgeError::Forward(format!(
            "forward spec {which} side {trimmed:?} has port 0, which is not a valid forward port"
        )));
    }
    Ok(addr)
}

/// Opens a tunneled connection to an overlay address.
///
/// Injected into [`serve_forward`] so that both edge transports share one
/// accept/pump loop. In netstack mode the implementation wraps the smoltcp
/// netstack connect path; in device mode it is a passthrough that dials a real
/// [`TcpStream`]. Both yield a [`TcpConn`] byte pipe. The concrete impls are
/// provided by the Wave-2 client; this module only defines the seam.
#[async_trait]
pub trait TunnelDialer: Send + Sync {
    /// Open a connection to `dst` across the tunnel.
    ///
    /// # Errors
    ///
    /// Returns an [`EdgeError`] when the tunneled connection cannot be
    /// established (handshake failure, connection refused, timeout, etc.).
    async fn dial(&self, dst: SocketAddr) -> Result<TcpConn, EdgeError>;
}

/// Bind `spec.local` and forward every accepted connection to `spec.remote`
/// through `dial`.
///
/// Each accepted connection is handled on its own spawned task by [`pump`], so
/// a single connection failing (dial error or mid-stream I/O error) never
/// tears down the listener. Runs until the listener's `accept` returns an
/// error, which is surfaced to the caller (the Wave-2 supervisor owns
/// restart/backoff policy).
///
/// # Errors
///
/// Returns [`EdgeError::Forward`] if the local listener cannot be bound, or if
/// `accept` fails.
pub async fn serve_forward<D>(spec: ForwardSpec, dial: D) -> Result<(), EdgeError>
where
    D: TunnelDialer + 'static,
{
    let listener = TcpListener::bind(spec.local).await.map_err(|e| {
        EdgeError::Forward(format!(
            "failed to bind forward listener on {}: {e}",
            spec.local
        ))
    })?;
    tracing::debug!(local = %spec.local, remote = %spec.remote, "edge forward listening");

    let dial = Arc::new(dial);
    loop {
        let (stream, peer) = listener.accept().await.map_err(|e| {
            EdgeError::Forward(format!("forward accept on {} failed: {e}", spec.local))
        })?;
        let dial = Arc::clone(&dial);
        let remote = spec.remote;
        tokio::spawn(async move {
            tracing::debug!(%peer, %remote, "edge forward connection accepted");
            let conn = match dial.dial(remote).await {
                Ok(conn) => conn,
                Err(e) => {
                    tracing::debug!(%peer, %remote, error = %e, "edge forward dial failed");
                    return;
                }
            };
            pump(stream, conn).await;
            tracing::debug!(%peer, %remote, "edge forward connection closed");
        });
    }
}

/// Bidirectionally pump one accepted host `stream` against a tunneled `conn`
/// until either direction closes, then shut down cleanly.
///
/// - Host→overlay: read from `stream`, send each chunk as a `Vec<u8>` on
///   `conn.tx`; on EOF or read error, drop `tx` to half-close the tunneled
///   send side.
/// - Overlay→host: forward each `conn.rx` chunk to `stream`'s write half; a
///   `None` from `rx` (peer closed) shuts down the write half.
///
/// Both directions run concurrently on this one task via `join!`, so no extra
/// tasks are spawned per connection.
async fn pump(stream: TcpStream, conn: TcpConn) {
    let TcpConn { tx, mut rx } = conn;
    let (mut read_half, mut write_half) = stream.into_split();

    // Host → overlay: read from the local socket, forward chunks into the
    // tunnel. Dropping `tx` at the end half-closes the tunneled send side.
    let host_to_overlay = async move {
        let mut buf = vec![0u8; PUMP_CHUNK_BYTES];
        loop {
            let n = match read_half.read(&mut buf).await {
                Ok(0) | Err(_) => break, // clean EOF or read error: stop reading
                Ok(n) => n,
            };
            if tx.send(buf[..n].to_vec()).await.is_err() {
                break; // overlay side dropped its receiver
            }
        }
        drop(tx);
    };

    // Overlay → host: drain tunnel chunks into the local socket. `rx` yielding
    // `None` means the remote peer closed, so we shut down the write half.
    let overlay_to_host = async move {
        while let Some(chunk) = rx.recv().await {
            if write_half.write_all(&chunk).await.is_err() {
                break; // local socket closed for writing
            }
        }
        let _ = write_half.shutdown().await;
    };

    tokio::join!(host_to_overlay, overlay_to_host);
}

#[cfg(test)]
mod tests {
    use super::{serve_forward, ForwardSpec, TunnelDialer};
    use crate::edge::{EdgeError, TcpConn};
    use async_trait::async_trait;
    use std::net::SocketAddr;
    use std::str::FromStr;
    use std::time::Duration;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    #[test]
    fn forward_spec_parses_ipv4() {
        let spec = ForwardSpec::from_str("127.0.0.1:8080=10.42.0.5:80").unwrap();
        assert_eq!(spec.local, "127.0.0.1:8080".parse::<SocketAddr>().unwrap());
        assert_eq!(spec.remote, "10.42.0.5:80".parse::<SocketAddr>().unwrap());
    }

    #[test]
    fn forward_spec_parses_bracketed_ipv6() {
        let spec = ForwardSpec::from_str("[::1]:8080=[fd00::2]:443").unwrap();
        assert_eq!(spec.local, "[::1]:8080".parse::<SocketAddr>().unwrap());
        assert_eq!(spec.remote, "[fd00::2]:443".parse::<SocketAddr>().unwrap());
    }

    #[test]
    fn forward_spec_rejects_missing_separator() {
        let err = ForwardSpec::from_str("127.0.0.1:8080").unwrap_err();
        assert!(err.to_string().contains("'='"), "{err}");
    }

    #[test]
    fn forward_spec_rejects_missing_port() {
        let err = ForwardSpec::from_str("127.0.0.1=10.0.0.1:80").unwrap_err();
        assert!(matches!(err, EdgeError::Forward(_)), "{err}");
    }

    #[test]
    fn forward_spec_rejects_bad_ip() {
        let err = ForwardSpec::from_str("not-an-ip:80=10.0.0.1:80").unwrap_err();
        assert!(matches!(err, EdgeError::Forward(_)), "{err}");
    }

    #[test]
    fn forward_spec_rejects_zero_port() {
        let local_zero = ForwardSpec::from_str("127.0.0.1:0=10.0.0.1:80").unwrap_err();
        assert!(local_zero.to_string().contains("port 0"), "{local_zero}");
        let remote_zero = ForwardSpec::from_str("127.0.0.1:80=10.0.0.1:0").unwrap_err();
        assert!(remote_zero.to_string().contains("port 0"), "{remote_zero}");
    }

    /// A [`TunnelDialer`] whose [`TcpConn`] loops everything the pump sends
    /// straight back — an in-memory echo, no real network involved.
    struct EchoDialer;

    #[async_trait]
    impl TunnelDialer for EchoDialer {
        async fn dial(&self, _dst: SocketAddr) -> Result<TcpConn, EdgeError> {
            // `tx` is what the pump writes host→overlay; `rx` is what the pump
            // reads overlay→host. Wire them through an echo task so bytes the
            // pump sends come back on the return channel. Bounded channels keep
            // the test deterministic.
            let (host_tx, mut host_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(16);
            let (back_tx, back_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(16);
            tokio::spawn(async move {
                while let Some(chunk) = host_rx.recv().await {
                    if back_tx.send(chunk).await.is_err() {
                        break;
                    }
                }
                // host_rx closed (pump dropped its tx on EOF) → drop back_tx →
                // pump's rx yields None → pump shuts down the write half.
            });
            Ok(TcpConn {
                tx: host_tx,
                rx: back_rx,
            })
        }
    }

    #[tokio::test]
    async fn serve_forward_echoes_through_dialer() {
        // Reserve a free local port, then hand it to serve_forward to bind.
        let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let port = probe.local_addr().unwrap().port();
        drop(probe);
        let local: SocketAddr = ([127, 0, 0, 1], port).into();
        let spec = ForwardSpec {
            local,
            remote: ([10, 0, 0, 1], 80).into(),
        };

        tokio::spawn(serve_forward(spec, EchoDialer));

        // Connect once the listener is up (retry-bound to stay deterministic).
        let mut stream = None;
        for _ in 0..100 {
            if let Ok(s) = tokio::net::TcpStream::connect(local).await {
                stream = Some(s);
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
        let mut stream = stream.expect("forward listener never came up");

        stream.write_all(b"hello edge").await.unwrap();
        // Half-close the write side so the pump sees EOF and tears the echo
        // path down cleanly after the bytes round-trip.
        stream.shutdown().await.unwrap();

        let mut got = Vec::new();
        stream.read_to_end(&mut got).await.unwrap();
        assert_eq!(got, b"hello edge");
    }
}