supermachine 0.7.97

// Host-side TLS termination. Spike22 accepts HTTPS on a host port,
// terminates the TLS handshake using rustls, and bridges the
// decrypted bytes to the guest via the muxer's existing host-side
// TCP listener (the port TsiListener::bind allocates on TSI_LISTEN).
//
// The guest sees plain HTTP — no TLS code in the guest at all.
//
// Each accepted TLS connection runs on its own thread with a single
// poll loop driving rustls's ServerConnection plus a plain TCP
// socket to the guest's TSI listener. Fine for moderate concurrency;
// we can swap in async/epoll later if needed.

// Portable: the terminator is pure rustls + std::net over the portable vsock
// muxer (which auto-binds the guest TSI listener's host TCP port on both
// backends via `TsiListener::bind`). On HVF it runs in the worker subprocess
// (CLI-driven); on KVM it runs in-process via `Vm::expose_tls`.
#![cfg(any(
    all(target_os = "macos", target_arch = "aarch64"),
    all(target_os = "linux", target_arch = "x86_64")
))]

use std::fmt;
use std::fs::File;
use std::io::{BufReader, ErrorKind, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::os::unix::net::UnixStream;
use std::sync::Arc;
use std::time::Duration;

use rustls::pki_types::CertificateDer;
use rustls::{ServerConfig, ServerConnection};

use crate::devices::virtio::vsock::device::Vsock;

pub struct TlsConfig {
    pub listen_addr: String,
    /// Optional: target a specific guest vm_port. If None (the usual
    /// case for HTTP guests), the terminator uses the FIRST TSI
    /// listener the muxer has registered — which works because the
    /// guest's TSI driver assigns random vm_ports per bind() that
    /// don't match what userspace asked for.
    pub vm_port: Option<u32>,
    pub cert_path: String,
    pub key_path: String,
}

#[derive(Debug)]
pub enum StartError {
    Bind {
        addr: String,
        source: std::io::Error,
    },
    Config(String),
    LocalAddr {
        addr: String,
        source: std::io::Error,
    },
    ThreadSpawn {
        name: String,
        source: std::io::Error,
    },
}

impl fmt::Display for StartError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StartError::Bind { addr, source } => write!(f, "TLS: bind {addr}: {source}"),
            StartError::Config(e) => write!(f, "TLS: load cert/key: {e}"),
            StartError::LocalAddr { addr, source } => {
                write!(f, "TLS: local_addr {addr}: {source}")
            }
            StartError::ThreadSpawn { name, source } => {
                write!(f, "spawn thread {name}: {source}")
            }
        }
    }
}

impl std::error::Error for StartError {}

/// Start the terminator and return the bound listen address (useful when
/// `listen_addr` requests port 0 and the caller needs the OS-assigned port).
pub fn start(cfg: TlsConfig, vsock: Arc<Vsock>) -> Result<std::net::SocketAddr, StartError> {
    let server_config =
        build_server_config(&cfg.cert_path, &cfg.key_path).map_err(StartError::Config)?;
    let listener = TcpListener::bind(&cfg.listen_addr).map_err(|source| StartError::Bind {
        addr: cfg.listen_addr.clone(),
        source,
    })?;
    let local = listener
        .local_addr()
        .map_err(|source| StartError::LocalAddr {
            addr: cfg.listen_addr.clone(),
            source,
        })?;
    eprintln!(
        "  TLS terminator on {local} -> guest vm_port={:?}",
        cfg.vm_port
    );

    let server_config = Arc::new(server_config);
    let vm_port_opt = cfg.vm_port;

    let name = "tls-acceptor".to_string();
    std::thread::Builder::new()
        .name(name.clone())
        .spawn(move || {
            for stream in listener.incoming() {
                let stream = match stream {
                    Ok(s) => s,
                    Err(e) => {
                        eprintln!("[tls] accept err: {e}");
                        continue;
                    }
                };
                let _ = stream.set_nodelay(true);
                let cfg = server_config.clone();
                let vsock_c = vsock.clone();
                std::thread::Builder::new()
                    .name("tls-conn".to_string())
                    .spawn(move || handle_conn(stream, cfg, vsock_c, vm_port_opt))
                    .ok();
            }
        })
        .map_err(|source| StartError::ThreadSpawn { name, source })?;
    Ok(local)
}

/// Sentry-backend variant of [`start`]: accept HTTPS on `cfg.listen_addr`,
/// terminate with rustls, and bridge the decrypted plaintext through a caller
/// supplied connector. The no-virt sentry's workload sockets live in the
/// persistent supervisor's LoopNet, so the connector usually dials the
/// supervisor's bridge stream rather than a host TCP port.
pub fn start_to_addr(cfg: TlsConfig, guest_port: u16) -> Result<std::net::SocketAddr, StartError> {
    let server_config =
        build_server_config(&cfg.cert_path, &cfg.key_path).map_err(StartError::Config)?;
    let listener = TcpListener::bind(&cfg.listen_addr).map_err(|source| StartError::Bind {
        addr: cfg.listen_addr.clone(),
        source,
    })?;
    let local = listener
        .local_addr()
        .map_err(|source| StartError::LocalAddr {
            addr: cfg.listen_addr.clone(),
            source,
        })?;
    eprintln!("  TLS terminator on {local} -> guest 127.0.0.1:{guest_port}");

    let server_config = Arc::new(server_config);
    let name = "tls-acceptor".to_string();
    std::thread::Builder::new()
        .name(name.clone())
        .spawn(move || {
            for stream in listener.incoming() {
                let stream = match stream {
                    Ok(s) => s,
                    Err(e) => {
                        eprintln!("[tls] accept err: {e}");
                        continue;
                    }
                };
                let _ = stream.set_nodelay(true);
                let cfg = server_config.clone();
                std::thread::Builder::new()
                    .name("tls-conn".to_string())
                    .spawn(move || handle_conn_to_addr(stream, cfg, guest_port))
                    .ok();
            }
        })
        .map_err(|source| StartError::ThreadSpawn { name, source })?;
    Ok(local)
}

/// Like [`start_to_addr`], but the plaintext backend is an already-connected
/// Unix stream. Used by the sentry backend to connect every TLS client to the
/// persistent supervisor's LoopNet bridge.
pub fn start_to_unix_stream<F>(
    cfg: TlsConfig,
    target: String,
    connect_plain: F,
) -> Result<std::net::SocketAddr, StartError>
where
    F: Fn() -> std::io::Result<UnixStream> + Send + Sync + 'static,
{
    let server_config =
        build_server_config(&cfg.cert_path, &cfg.key_path).map_err(StartError::Config)?;
    let listener = TcpListener::bind(&cfg.listen_addr).map_err(|source| StartError::Bind {
        addr: cfg.listen_addr.clone(),
        source,
    })?;
    let local = listener
        .local_addr()
        .map_err(|source| StartError::LocalAddr {
            addr: cfg.listen_addr.clone(),
            source,
        })?;
    eprintln!("  TLS terminator on {local} -> {target}");

    let server_config = Arc::new(server_config);
    let connect_plain = Arc::new(connect_plain);
    let target = Arc::new(target);
    let name = "tls-acceptor".to_string();
    std::thread::Builder::new()
        .name(name.clone())
        .spawn(move || {
            for stream in listener.incoming() {
                let stream = match stream {
                    Ok(s) => s,
                    Err(e) => {
                        eprintln!("[tls] accept err: {e}");
                        continue;
                    }
                };
                let _ = stream.set_nodelay(true);
                let cfg = server_config.clone();
                let connect_plain = connect_plain.clone();
                let target = target.clone();
                std::thread::Builder::new()
                    .name("tls-conn".to_string())
                    .spawn(move || handle_conn_to_unix_stream(stream, cfg, connect_plain, target))
                    .ok();
            }
        })
        .map_err(|source| StartError::ThreadSpawn { name, source })?;
    Ok(local)
}

/// Per-connection bridge for [`start_to_unix_stream`]: connect to the sentry
/// supervisor bridge, then drive the same rustls↔plaintext [`pump`] as
/// [`handle_conn`].
fn handle_conn_to_unix_stream(
    mut tls_sock: TcpStream,
    cfg: Arc<ServerConfig>,
    connect_plain: Arc<dyn Fn() -> std::io::Result<UnixStream> + Send + Sync>,
    target: Arc<String>,
) {
    let mut conn = match ServerConnection::new(cfg) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("[tls] ServerConnection: {e}");
            return;
        }
    };
    // Finish the TLS handshake BEFORE dialing the guest (see
    // [`complete_handshake`]) — otherwise an eager/early-closing backend races
    // the half-built session and the client sees `SSL_ERROR_ZERO_RETURN`.
    if let Err(e) = complete_handshake(&mut conn, &mut tls_sock) {
        if crate::trace::enabled("tls") {
            eprintln!("[tls] handshake: {e}");
        }
        return;
    }

    let plain_sock = match connect_plain() {
        Ok(s) => s,
        Err(e) => {
            eprintln!("[tls] connect {target}: {e}");
            return;
        }
    };

    let _ = tls_sock.set_nonblocking(true);
    let _ = plain_sock.set_nonblocking(true);
    if let Err(e) = pump(&mut conn, tls_sock, plain_sock) {
        if crate::trace::enabled("tls") {
            eprintln!("[tls] pump end: {e}");
        }
    }
}

/// Per-connection bridge for [`start_to_addr`]: dial a plaintext TCP listener at
/// `127.0.0.1:guest_port`, then drive the same rustls↔plaintext [`pump`] as
/// [`handle_conn`].
fn handle_conn_to_addr(mut tls_sock: TcpStream, cfg: Arc<ServerConfig>, guest_port: u16) {
    let mut conn = match ServerConnection::new(cfg) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("[tls] ServerConnection: {e}");
            return;
        }
    };
    if let Err(e) = complete_handshake(&mut conn, &mut tls_sock) {
        if crate::trace::enabled("tls") {
            eprintln!("[tls] handshake: {e}");
        }
        return;
    }

    let plain_sock = match TcpStream::connect(("127.0.0.1", guest_port)) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("[tls] connect 127.0.0.1:{guest_port}: {e}");
            return;
        }
    };
    let _ = plain_sock.set_nodelay(true);

    let _ = tls_sock.set_nonblocking(true);
    let _ = plain_sock.set_nonblocking(true);
    if let Err(e) = pump(&mut conn, tls_sock, plain_sock) {
        if crate::trace::enabled("tls") {
            eprintln!("[tls] pump end: {e}");
        }
    }
}

fn build_server_config(cert_path: &str, key_path: &str) -> Result<ServerConfig, String> {
    let _ = rustls::crypto::ring::default_provider().install_default();

    let cert_file = File::open(cert_path).map_err(|e| format!("open cert {cert_path}: {e}"))?;
    let mut cr = BufReader::new(cert_file);
    let certs: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut cr)
        .collect::<Result<_, _>>()
        .map_err(|e| format!("parse certs: {e}"))?;
    if certs.is_empty() {
        return Err(format!("no certs in {cert_path}"));
    }

    let key_file = File::open(key_path).map_err(|e| format!("open key {key_path}: {e}"))?;
    let mut kr = BufReader::new(key_file);
    let key = rustls_pemfile::private_key(&mut kr)
        .map_err(|e| format!("parse key: {e}"))?
        .ok_or_else(|| format!("no private key in {key_path}"))?;

    ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(certs, key)
        .map_err(|e| format!("ServerConfig: {e}"))
}

/// Drive the rustls server handshake to completion on a BLOCKING socket BEFORE
/// the backend is dialed. A TLS terminator MUST finish the client handshake
/// first: otherwise a backend that speaks first, or one that EOFs early, races
/// the half-built TLS session and the [`pump`] flushes a `close_notify` before
/// the `ServerHello` ever goes out — the client sees `SSL_ERROR_ZERO_RETURN`
/// (observed with an eager `printf | nc` backend). Handshaking up front is also
/// strictly better in production: we never dial the backend for a connection
/// whose TLS handshake fails (port scans, bad SNI, slow-loris). Returns `Err`
/// if the client aborts mid-handshake.
fn complete_handshake(conn: &mut ServerConnection, sock: &mut TcpStream) -> std::io::Result<()> {
    sock.set_nonblocking(false)?;
    while conn.is_handshaking() {
        // Flush our pending flight (ServerHello, …) before blocking on a read,
        // so the peer can make progress and we don't deadlock waiting for input.
        while conn.wants_write() {
            if conn.write_tls(sock)? == 0 {
                return Err(std::io::Error::new(
                    ErrorKind::WriteZero,
                    "tls handshake: wrote 0",
                ));
            }
        }
        if conn.is_handshaking() && conn.wants_read() {
            if conn.read_tls(sock)? == 0 {
                return Err(std::io::Error::new(
                    ErrorKind::UnexpectedEof,
                    "client closed during TLS handshake",
                ));
            }
            conn.process_new_packets()
                .map_err(|e| std::io::Error::new(ErrorKind::Other, format!("{e}")))?;
        }
    }
    // Flush the final handshake bytes (e.g. the server's last flight).
    while conn.wants_write() {
        if conn.write_tls(sock)? == 0 {
            break;
        }
    }
    Ok(())
}

fn handle_conn(
    mut tls_sock: TcpStream,
    cfg: Arc<ServerConfig>,
    vsock: Arc<Vsock>,
    vm_port: Option<u32>,
) {
    let mut conn = match ServerConnection::new(cfg) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("[tls] ServerConnection: {e}");
            return;
        }
    };
    // Complete the TLS handshake BEFORE dialing the backend (see
    // [`complete_handshake`]). Only then do we wait for / connect the guest.
    if let Err(e) = complete_handshake(&mut conn, &mut tls_sock) {
        if crate::trace::enabled("tls") {
            eprintln!("[tls] handshake: {e}");
        }
        return;
    }

    let host_port = match wait_for_host_port(&vsock, vm_port) {
        Some(p) => p,
        None => {
            eprintln!("[tls] no host port (vm_port={vm_port:?}) after wait");
            return;
        }
    };
    let plain_sock = match TcpStream::connect(("127.0.0.1", host_port)) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("[tls] connect 127.0.0.1:{host_port}: {e}");
            return;
        }
    };
    let _ = plain_sock.set_nodelay(true);

    let _ = tls_sock.set_nonblocking(true);
    let _ = plain_sock.set_nonblocking(true);
    if let Err(e) = pump(&mut conn, tls_sock, plain_sock) {
        if crate::trace::enabled("tls") {
            eprintln!("[tls] pump end: {e}");
        }
    }
}

/// Bridge bytes between TLS client and plaintext guest.
trait PlainBackend: Read + Write {
    fn shutdown_both(&self) -> std::io::Result<()>;
}

impl PlainBackend for TcpStream {
    fn shutdown_both(&self) -> std::io::Result<()> {
        self.shutdown(std::net::Shutdown::Both)
    }
}

impl PlainBackend for UnixStream {
    fn shutdown_both(&self) -> std::io::Result<()> {
        self.shutdown(std::net::Shutdown::Both)
    }
}

fn pump<P: PlainBackend>(
    conn: &mut ServerConnection,
    mut tls_sock: TcpStream,
    mut plain_sock: P,
) -> std::io::Result<()> {
    let mut buf = [0u8; 16 * 1024];
    let mut closed_tls = false;
    let mut closed_plain = false;
    loop {
        let mut did_work = false;

        // 1. Pull encrypted bytes from client into rustls.
        if conn.wants_read() && !closed_tls {
            match conn.read_tls(&mut tls_sock) {
                Ok(0) => {
                    closed_tls = true;
                }
                Ok(_) => {
                    did_work = true;
                    if let Err(e) = conn.process_new_packets() {
                        return Err(std::io::Error::new(ErrorKind::Other, format!("{e}")));
                    }
                }
                Err(e) if e.kind() == ErrorKind::WouldBlock => {}
                Err(e) => return Err(e),
            }
        }

        // 2. Pull decrypted plaintext, forward to guest.
        loop {
            match conn.reader().read(&mut buf) {
                Ok(0) => break,
                Ok(n) => {
                    did_work = true;
                    if plain_sock.write_all(&buf[..n]).is_err() {
                        closed_plain = true;
                        break;
                    }
                }
                Err(e) if e.kind() == ErrorKind::WouldBlock => break,
                Err(_) => break,
            }
        }

        // 3. Pull bytes from guest, push into rustls writer.
        if !closed_plain {
            match plain_sock.read(&mut buf) {
                Ok(0) => {
                    closed_plain = true;
                    conn.send_close_notify();
                }
                Ok(n) => {
                    did_work = true;
                    conn.writer().write_all(&buf[..n])?;
                }
                Err(e) if e.kind() == ErrorKind::WouldBlock => {}
                Err(_) => {
                    closed_plain = true;
                    conn.send_close_notify();
                }
            }
        }

        // 4. Push encrypted bytes back to client.
        if conn.wants_write() {
            match conn.write_tls(&mut tls_sock) {
                Ok(_) => {
                    did_work = true;
                }
                Err(e) if e.kind() == ErrorKind::WouldBlock => {}
                Err(e) => return Err(e),
            }
        }

        // Termination: both sides drained + closed.
        if closed_tls && closed_plain && !conn.wants_write() {
            break;
        }
        // If there's still data to flush either way, keep going.
        if closed_tls && !conn.wants_write() && !closed_plain {
            // Client side is closed; drain whatever guest sends, then exit.
            // Already handled in (3); fall through to the sleep.
        }

        if !did_work {
            // Avoid spin: brief sleep when both sides idle.
            std::thread::sleep(Duration::from_micros(200));
        }
    }
    let _ = tls_sock.shutdown(std::net::Shutdown::Both);
    let _ = plain_sock.shutdown_both();
    Ok(())
}

fn wait_for_host_port(vsock: &Vsock, vm_port: Option<u32>) -> Option<u16> {
    let lookup = || match vm_port {
        Some(p) => vsock.muxer().host_port_for_vm_port(p),
        None => vsock.muxer().first_host_port(),
    };
    for _ in 0..50 {
        if let Some(p) = lookup() {
            return Some(p);
        }
        std::thread::sleep(Duration::from_millis(20));
    }
    lookup()
}