pixelpass/src/host/serve.rs

//! Display-server-agnostic serving layer: takes a capture child's stdout
//! producing MPEG-TS bytes and fans them out to N concurrent HTTP viewers
//! on a localhost port. One reader task pumps stdout chunks into a
//! tokio::sync::broadcast channel; the accept loop spawns one drain task
//! per accepted TCP connection. Slow consumers see Lagged and skip ahead;
//! MPEG-TS resyncs at the next keyframe.
//!
//! Backends (host/wayland.rs, future host/x11.rs) build their own gst
//! pipeline and hand the resulting ChildStdout to [`Serve::bind`].

use anyhow::{Context, Result, bail};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::process::ChildStdout;
use tokio::sync::broadcast;
use tokio::task::JoinHandle;
use tokio::time::{Instant, sleep};

/// Broadcast-channel capacity in chunks. Each chunk is up to 64 KiB from
/// the capture child's stdout, so 16 chunks ≈ 1 MiB ≈ ~2 s of buffered
/// jitter at typical bitrates. A viewer that falls behind by more than
/// this gets Lagged and skips ahead — MPEG-TS recovers at the next
/// keyframe.
const FANOUT_CAPACITY: usize = 16;

/// Size of each chunk read from the capture child's stdout.
const READ_CHUNK: usize = 64 * 1024;

/// Owns the localhost HTTP listener and the two long-running tasks that
/// pump bytes from a capture child to all connected viewers.
pub struct Serve {
    port: u16,
    reader: Option<JoinHandle<()>>,
    server: Option<JoinHandle<()>>,
}

impl Serve {
    /// Bind a localhost listener on a random port, set up the broadcast
    /// fanout, and spawn the reader + accept-loop tasks. The provided
    /// `stdout` is assumed to produce MPEG-TS bytes.
    pub async fn bind(stdout: ChildStdout) -> Result<Self> {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .context("could not bind local capture HTTP listener")?;
        let port = listener.local_addr()?.port();

        let (tx, _) = broadcast::channel::<Arc<Vec<u8>>>(FANOUT_CAPACITY);
        let reader = tokio::spawn(pump_to_broadcast(stdout, tx.clone()));
        let server = tokio::spawn(run_accept_loop(listener, tx));

        Ok(Self {
            port,
            reader: Some(reader),
            server: Some(server),
        })
    }

    pub fn local_port(&self) -> u16 {
        self.port
    }

    /// Abort the reader and accept-loop tasks. Backends typically call this
    /// after killing their capture child so the reader sees stdout EOF and
    /// exits on its own; the abort is a backstop.
    pub async fn shutdown(mut self) {
        if let Some(task) = self.reader.take() {
            task.abort();
        }
        if let Some(task) = self.server.take() {
            task.abort();
        }
    }
}

impl Drop for Serve {
    fn drop(&mut self) {
        if let Some(task) = self.reader.as_ref() {
            task.abort();
        }
        if let Some(task) = self.server.as_ref() {
            task.abort();
        }
    }
}

/// Connect to the local capture HTTP listener, retrying until it's up or
/// we time out. Returns the connected socket — the bridge layer pipes
/// QUIC↔this socket once it's open.
pub async fn connect_to_capture(port: u16, max_wait: Duration) -> Result<TcpStream> {
    let deadline = Instant::now() + max_wait;
    loop {
        match TcpStream::connect(("127.0.0.1", port)).await {
            Ok(stream) => return Ok(stream),
            Err(_) if Instant::now() < deadline => {
                sleep(Duration::from_millis(50)).await;
            }
            Err(e) => bail!("capture HTTP listener never came up on 127.0.0.1:{port}: {e}"),
        }
    }
}

/// Read the capture child's stdout in chunks and broadcast each to all
/// current subscribers. `broadcast::send` returns Err when there are no
/// receivers; we ignore it so the capture child isn't backpressured
/// waiting for a viewer.
async fn pump_to_broadcast(mut stdout: ChildStdout, tx: broadcast::Sender<Arc<Vec<u8>>>) {
    let mut buf = vec![0u8; READ_CHUNK];
    loop {
        match stdout.read(&mut buf).await {
            Ok(0) => {
                tracing::info!("capture stdout EOF — fanout reader exiting");
                return;
            }
            Ok(n) => {
                let chunk = Arc::new(buf[..n].to_vec());
                let _ = tx.send(chunk);
            }
            Err(e) => {
                tracing::warn!("capture stdout read error: {e}");
                return;
            }
        }
    }
}

async fn run_accept_loop(listener: TcpListener, tx: broadcast::Sender<Arc<Vec<u8>>>) {
    loop {
        let sock = match listener.accept().await {
            Ok((s, _)) => s,
            Err(e) => {
                // Most accept errors are transient (EMFILE from a brief FD spike,
                // EINTR, etc.). Bailing on the first one would kill the entire
                // viewer fanout for the rest of the session.
                tracing::warn!("capture HTTP accept failed (continuing): {e}");
                continue;
            }
        };
        let rx = tx.subscribe();
        tokio::spawn(serve_one_viewer(sock, rx));
    }
}

async fn serve_one_viewer(mut sock: TcpStream, mut rx: broadcast::Receiver<Arc<Vec<u8>>>) {
    if !drain_http_request(&mut sock).await {
        return;
    }

    const RESPONSE: &[u8] = b"HTTP/1.1 200 OK\r\n\
        Content-Type: video/mp2t\r\n\
        Cache-Control: no-cache, no-store\r\n\
        Connection: close\r\n\
        \r\n";
    if sock.write_all(RESPONSE).await.is_err() {
        return;
    }

    loop {
        match rx.recv().await {
            Ok(chunk) => {
                if sock.write_all(&chunk).await.is_err() {
                    return;
                }
            }
            Err(broadcast::error::RecvError::Lagged(skipped)) => {
                tracing::warn!(
                    skipped,
                    "viewer fanout lagged — MPEG-TS will resync at next keyframe"
                );
                continue;
            }
            Err(broadcast::error::RecvError::Closed) => return,
        }
    }
}

async fn drain_http_request(sock: &mut TcpStream) -> bool {
    let mut buf = [0u8; 1024];
    let mut total = Vec::with_capacity(512);
    loop {
        match sock.read(&mut buf).await {
            Ok(0) => return false,
            Ok(n) => total.extend_from_slice(&buf[..n]),
            Err(_) => return false,
        }
        if total.windows(4).any(|w| w == b"\r\n\r\n") {
            return true;
        }
        if total.len() > 16 * 1024 {
            return false;
        }
    }
}