sail-rs 0.3.3

//! Local port forwarding for interactive sessions. A [`PortForward`] listens on
//! the user's machine and forwards each accepted connection to the same port on
//! the box's localhost, so a sandbox's localhost servers (dev servers, OAuth
//! login callbacks) are reachable locally.

use std::net::{Ipv4Addr, Ipv6Addr};
use std::sync::Arc;
use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;

use crate::error::SailError;
use crate::pb::workerproxy::v1 as pb;
use crate::worker::WorkerProxy;

/// Buffered tunnel frames from the local socket toward the guest; bounds memory
/// and applies backpressure when the uplink is slow.
const TUNNEL_CHANNEL_CAP: usize = 32;
/// Read chunk for the local-socket read loop.
const TUNNEL_READ_CHUNK: usize = 32 * 1024;

/// A live local port forward. Dropping it stops accepting new connections;
/// connections already open finish on their own.
pub struct PortForward {
    local_port: u16,
    accept: tokio::task::JoinHandle<()>,
}

impl Drop for PortForward {
    fn drop(&mut self) {
        self.accept.abort();
    }
}

impl PortForward {
    /// The local port the forward is listening on (`127.0.0.1`, and `::1` when
    /// available). This equals the requested port, or an OS-picked free port when
    /// 0 was requested.
    pub fn local_port(&self) -> u16 {
        self.local_port
    }
}

/// Start forwarding `127.0.0.1:local_port` to the sailbox's guest-local
/// `remote_port`. The listener binds up front (so a bind failure surfaces here),
/// then accepts in the background. Passing `local_port` 0 lets the OS pick a free
/// port, reported by [`PortForward::local_port`].
pub async fn forward_port(
    worker: Arc<WorkerProxy>,
    endpoint: String,
    sailbox_id: String,
    local_port: u16,
    remote_port: u16,
) -> Result<PortForward, SailError> {
    let v4 = TcpListener::bind((Ipv4Addr::LOCALHOST, local_port))
        .await
        .map_err(|err| SailError::Internal {
            message: format!("bind 127.0.0.1:{local_port}: {err}"),
        })?;
    let bound = v4.local_addr().map_or(local_port, |addr| addr.port());
    // Also listen on the IPv6 loopback on the same port: some systems resolve
    // `localhost` to `::1`, so catching both keeps a forwarded URL reachable
    // however the browser resolves it.
    let v6 = match TcpListener::bind((Ipv6Addr::LOCALHOST, bound)).await {
        Ok(v6) => Some(v6),
        // The IPv6 loopback port is held by another service. A browser resolving
        // `localhost` to `::1` would reach that service instead of this forward,
        // and a login redirect to the exact port cannot be rewritten afterward, so
        // treat the forward as unavailable rather than own only half the port.
        Err(err) if err.kind() == std::io::ErrorKind::AddrInUse => {
            return Err(SailError::Internal {
                message: format!("bind [::1]:{bound}: {err}"),
            });
        }
        // IPv6 loopback is otherwise unavailable (e.g. a host without IPv6). The
        // v4 listener suffices, since `localhost` then resolves to 127.0.0.1.
        Err(_) => None,
    };
    let accept = tokio::spawn(async move {
        loop {
            let accepted = match &v6 {
                Some(v6) => tokio::select! {
                    conn = v4.accept() => conn,
                    conn = v6.accept() => conn,
                },
                None => v4.accept().await,
            };
            match accepted {
                Ok((conn, _)) => {
                    let worker = Arc::clone(&worker);
                    let endpoint = endpoint.clone();
                    let sailbox_id = sailbox_id.clone();
                    tokio::spawn(async move {
                        let _ =
                            tunnel_connection(worker, &endpoint, &sailbox_id, remote_port, conn)
                                .await;
                    });
                }
                // A transient accept error (e.g. fd exhaustion) must not kill the
                // forward the caller still holds; back off briefly and keep going.
                Err(_) => tokio::time::sleep(Duration::from_millis(50)).await,
            }
        }
    });
    Ok(PortForward {
        local_port: bound,
        accept,
    })
}

/// Pump one accepted local connection through a `TunnelSailboxPort` stream: the
/// handshake names the sailbox and port, then bytes flow both ways until either
/// side closes.
async fn tunnel_connection(
    worker: Arc<WorkerProxy>,
    endpoint: &str,
    sailbox_id: &str,
    remote_port: u16,
    conn: TcpStream,
) -> Result<(), SailError> {
    let (tx, rx) = mpsc::channel::<pb::TunnelSailboxPortRequest>(TUNNEL_CHANNEL_CAP);
    // The handshake frame goes first; the channel preserves order.
    let handshake = pb::TunnelSailboxPortRequest {
        sailbox_id: sailbox_id.to_string(),
        port: u32::from(remote_port),
        data: Vec::new(),
        eof: false,
    };
    if tx.send(handshake).await.is_err() {
        return Ok(());
    }
    let (mut read_half, mut write_half) = conn.into_split();

    // Start the local -> guest pump before awaiting the response. The guest
    // dials the service and, for a request/response protocol (HTTP dev servers,
    // OAuth callbacks), waits for request bytes before it responds; the response
    // await resolves only once the server sends its first frame, so the request
    // must already be flowing or the first connection deadlocks. Half-close on
    // EOF.
    let uplink = tokio::spawn(async move {
        let mut buf = vec![0u8; TUNNEL_READ_CHUNK];
        loop {
            match read_half.read(&mut buf).await {
                Ok(n) if n > 0 => {
                    if tx.send(data_frame(buf[..n].to_vec())).await.is_err() {
                        break;
                    }
                }
                // No more request bytes, whether a clean EOF or a local reset
                // (a browser aborting a forwarded request): half-close with an
                // eof frame so the guest CloseWrites its dialed connection and
                // can finish replying. Dropping the stream without this frame
                // would not carry the half-close, leaving a request/response
                // service waiting on the guest until the shell exits.
                _ => {
                    let _ = tx.send(eof_frame()).await;
                    break;
                }
            }
        }
    });

    let response = match open_tunnel_stream(&worker, endpoint, rx).await {
        Ok(response) => response,
        Err(err) => {
            uplink.abort();
            return Err(err);
        }
    };
    let mut server = response.into_inner();

    // guest -> local: write frames to the socket, half-close on EOF.
    while let Ok(Some(frame)) = server.message().await {
        if !frame.data.is_empty() && write_half.write_all(&frame.data).await.is_err() {
            break;
        }
        if frame.eof {
            let _ = write_half.shutdown().await;
        }
    }
    uplink.abort();
    Ok(())
}

/// The port a URL targets on a server in the sandbox, if any: its host is
/// `localhost`, the loopback `127.0.0.1`, or a wildcard (`0.0.0.0` or `::`), each
/// of which the guest reaches at `127.0.0.1`. The port is forwarded and the URL
/// rewritten to the bound local address for opening.
pub fn forwardable_local_port(url: &str) -> Option<u16> {
    let parsed = url::Url::parse(url).ok()?;
    if is_forwardable_local_host(&parsed) {
        return parsed.port_or_known_default();
    }
    None
}

// A host the guest reaches at 127.0.0.1: `localhost`, the loopback 127.0.0.1
// itself, the IPv4 wildcard 0.0.0.0, or the IPv6 wildcard `::`. A `::` listener is
// dual-stack and so answers on 127.0.0.1, which is what the port scanner forwards
// and the tunnel dials. Other 127.0.0.0/8 addresses and the IPv6 loopback `::1`
// are excluded, since the tunnel dials 127.0.0.1 specifically.
fn is_forwardable_local_host(url: &url::Url) -> bool {
    match url.host() {
        Some(url::Host::Domain(domain)) => domain == "localhost",
        Some(url::Host::Ipv4(ip)) => ip == Ipv4Addr::LOCALHOST || ip.is_unspecified(),
        Some(url::Host::Ipv6(ip)) => ip.is_unspecified(),
        _ => false,
    }
}

/// Whether `url` targets a loopback address the sandbox meant for itself but the
/// tunnel does not dial (any loopback other than 127.0.0.1: a different
/// 127.0.0.0/8 address, the IPv6 loopback `::1`, or an IPv4-mapped loopback like
/// `::ffff:127.0.0.1`). Such a URL must not be opened against the user's own
/// loopback.
pub fn is_unforwardable_loopback_url(url: &str) -> bool {
    let Ok(parsed) = url::Url::parse(url) else {
        return false;
    };
    match parsed.host() {
        Some(url::Host::Ipv4(ip)) => ip.is_loopback() && ip != Ipv4Addr::LOCALHOST,
        // An IPv4-mapped address (`::ffff:a.b.c.d`) reaches the same host as
        // `a.b.c.d`, and the port scanner never forwards a mapped bind, so any
        // mapped loopback is one the tunnel cannot reach.
        Some(url::Host::Ipv6(ip)) => match ip.to_ipv4_mapped() {
            Some(v4) => v4.is_loopback(),
            None => ip.is_loopback(),
        },
        _ => false,
    }
}

/// If `url` targets a forwarded port on a sandbox server, rewrite it to reach the
/// forward. A wildcard bind host (`0.0.0.0` or `::`) is not itself connectable,
/// so it is pointed at `localhost`, which resolves to the forward's loopback. A
/// `localhost` or `127.0.0.1` host is left exactly as the guest wrote it, so an
/// HTTPS certificate for that name still verifies (and localhost-scoped cookies
/// and origins are preserved); the forward listens on the same port, so nothing
/// else changes. Returns the URL unchanged when it does not target a forwardable
/// local host or the port is not forwarded.
pub fn rewrite_loopback_url(url: &str, local_port_for: impl Fn(u16) -> Option<u16>) -> String {
    let Ok(mut parsed) = url::Url::parse(url) else {
        return url.to_string();
    };
    if !is_forwardable_local_host(&parsed) {
        return url.to_string();
    }
    let Some(remote) = parsed.port_or_known_default() else {
        return url.to_string();
    };
    let Some(local) = local_port_for(remote) else {
        return url.to_string();
    };
    let wildcard = matches!(parsed.host(), Some(url::Host::Ipv4(ip)) if ip.is_unspecified())
        || matches!(parsed.host(), Some(url::Host::Ipv6(ip)) if ip.is_unspecified());
    if !wildcard {
        return url.to_string();
    }
    if parsed.set_host(Some("localhost")).is_ok() && parsed.set_port(Some(local)).is_ok() {
        parsed.to_string()
    } else {
        url.to_string()
    }
}

/// The loopback callback an external login URL will redirect to, read from a
/// `redirect_uri`/`redirect_url` query parameter.
pub enum RedirectCallback {
    /// A forwardable localhost port. The login completes once that port is
    /// forwarded, so the shell holds the browser open until then.
    Forwardable(u16),
    /// A loopback the tunnel cannot dial (a non-`.0.0.1` 127/8 address, `::1`, or
    /// an IPv4-mapped loopback). The provider's redirect would hit the user's own
    /// machine rather than the sandbox, so the login cannot complete.
    Unreachable,
}

/// Classify the loopback callback an external login URL redirects to, if any, so
/// the shell can decline to open a login whose redirect would not reach the
/// sandbox.
pub fn redirect_callback(url: &str) -> Option<RedirectCallback> {
    let parsed = url::Url::parse(url).ok()?;
    for (key, value) in parsed.query_pairs() {
        if key == "redirect_uri" || key == "redirect_url" {
            if let Some(port) = forwardable_local_port(&value) {
                return Some(RedirectCallback::Forwardable(port));
            }
            if is_unforwardable_loopback_url(&value) {
                return Some(RedirectCallback::Unreachable);
            }
        }
    }
    None
}

/// Open the `TunnelSailboxPort` stream, feeding it the request frames from `rx`.
type TunnelResponse = tonic::Response<tonic::Streaming<pb::TunnelSailboxPortResponse>>;

async fn open_tunnel_stream(
    worker: &Arc<WorkerProxy>,
    endpoint: &str,
    rx: mpsc::Receiver<pb::TunnelSailboxPortRequest>,
) -> Result<TunnelResponse, SailError> {
    let request = worker.request_for(ReceiverStream::new(rx), &[], /* timeout */ None)?;
    worker
        .client_for(endpoint)?
        .tunnel_sailbox_port(request)
        .await
        .map_err(|status| SailError::from_file_rpc_status(&status))
}

fn data_frame(data: Vec<u8>) -> pb::TunnelSailboxPortRequest {
    pb::TunnelSailboxPortRequest {
        sailbox_id: String::new(),
        port: 0,
        data,
        eof: false,
    }
}

fn eof_frame() -> pb::TunnelSailboxPortRequest {
    pb::TunnelSailboxPortRequest {
        sailbox_id: String::new(),
        port: 0,
        data: Vec::new(),
        eof: true,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn forwardable_local_port_covers_loopback_wildcard_and_defaults() {
        assert_eq!(forwardable_local_port("http://localhost:5173/"), Some(5173));
        assert_eq!(
            forwardable_local_port("http://127.0.0.1:9000/x"),
            Some(9000)
        );
        // Wildcard binds are reachable at 127.0.0.1 in the guest, including the
        // dual-stack IPv6 wildcard `::`.
        assert_eq!(forwardable_local_port("http://0.0.0.0:3000/"), Some(3000));
        assert_eq!(forwardable_local_port("http://[::]:3000/"), Some(3000));
        // A default-port target uses the scheme's known default.
        assert_eq!(forwardable_local_port("http://localhost/cb"), Some(80));
        assert_eq!(forwardable_local_port("https://localhost/cb"), Some(443));
        // Not forwardable: a non-local URL, the IPv6 loopback `::1`, or a
        // 127.0.0.0/8 address other than 127.0.0.1 (the tunnel dials 127.0.0.1).
        assert_eq!(forwardable_local_port("https://example.com:8080/"), None);
        assert_eq!(forwardable_local_port("http://[::1]:5173/"), None);
        assert_eq!(
            forwardable_local_port("http://[::ffff:127.0.0.1]:5173/"),
            None
        );
        assert_eq!(forwardable_local_port("http://127.0.0.2:3000/"), None);
    }

    #[test]
    fn is_unforwardable_loopback_url_matches_undialed_loopbacks() {
        // A loopback the tunnel does not dial: `::1`, 127/8 other than .0.0.1, or
        // an IPv4-mapped loopback (which `is_loopback` alone does not catch).
        assert!(is_unforwardable_loopback_url("http://[::1]:5173/"));
        assert!(is_unforwardable_loopback_url("http://127.0.0.2:3000/"));
        assert!(is_unforwardable_loopback_url(
            "http://[::ffff:127.0.0.1]:3000/"
        ));
        // Forwardable and external hosts are not flagged.
        assert!(!is_unforwardable_loopback_url("http://localhost:8080/"));
        assert!(!is_unforwardable_loopback_url("http://127.0.0.1:8080/"));
        assert!(!is_unforwardable_loopback_url("http://0.0.0.0:8080/"));
        assert!(!is_unforwardable_loopback_url("http://[::]:8080/"));
        assert!(!is_unforwardable_loopback_url("https://example.com/"));
    }

    #[test]
    fn rewrite_loopback_url_only_rewrites_wildcard_hosts() {
        let mut forwards = HashMap::new();
        forwards.insert(3000u16, 3000u16); // forwards are one-to-one
        forwards.insert(8080u16, 8080u16);
        let lookup = |remote: u16| forwards.get(&remote).copied();

        // A localhost or 127.0.0.1 host is left exactly as written, so an HTTPS
        // cert for that name still verifies and localhost origins are preserved.
        assert_eq!(
            rewrite_loopback_url("https://localhost:8080/app", lookup),
            "https://localhost:8080/app"
        );
        assert_eq!(
            rewrite_loopback_url("http://127.0.0.1:3000/x", lookup),
            "http://127.0.0.1:3000/x"
        );
        // A wildcard host is not connectable, so it is pointed at localhost.
        assert_eq!(
            rewrite_loopback_url("http://0.0.0.0:3000/x", lookup),
            "http://localhost:3000/x"
        );
        assert_eq!(
            rewrite_loopback_url("http://[::]:3000/x", lookup),
            "http://localhost:3000/x"
        );
        // Unforwarded and non-loopback URLs are left alone.
        assert_eq!(
            rewrite_loopback_url("http://localhost:9999/x", lookup),
            "http://localhost:9999/x"
        );
        assert_eq!(
            rewrite_loopback_url("https://example.com:8080/x", lookup),
            "https://example.com:8080/x"
        );
    }

    #[test]
    fn redirect_callback_classifies_the_redirect_uri() {
        use RedirectCallback::{Forwardable, Unreachable};
        // A forwardable localhost callback returns its port (redirect_uri or
        // redirect_url, encoded or plain, localhost or 127.0.0.1).
        assert!(matches!(
            redirect_callback(
                "https://example.com/oauth?client_id=x&redirect_uri=http%3A%2F%2Flocalhost%3A43222%2Fcallback"
            ),
            Some(Forwardable(43222))
        ));
        assert!(matches!(
            redirect_callback("https://id.example/a?redirect_url=http://127.0.0.1:5000/cb"),
            Some(Forwardable(5000))
        ));
        // A loopback callback the tunnel can't reach is unreachable, not opened.
        assert!(matches!(
            redirect_callback("https://example.com/oauth?redirect_uri=http://[::1]:5173/cb"),
            Some(Unreachable)
        ));
        assert!(matches!(
            redirect_callback("https://example.com/oauth?redirect_uri=http://127.0.0.2:5173/cb"),
            Some(Unreachable)
        ));
        // A remote redirect, or none at all, is not a loopback callback.
        assert!(
            redirect_callback("https://example.com/oauth?redirect_uri=https://example.com/cb")
                .is_none()
        );
        assert!(redirect_callback("https://example.com/login").is_none());
    }
}