pixelpass/src/host/mod.rs

pub mod aec;
pub mod audio;
pub mod audit;
mod capture;
mod observer;
mod pipeline;
mod quality;
mod serve;
pub mod taint;
mod wayland;
mod x11;

use anyhow::{Result, bail};
use iroh::endpoint::Connection;
use iroh::{Endpoint, EndpointAddr};
use iroh_tickets::endpoint::EndpointTicket;
use std::collections::HashMap;
use std::time::Duration;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;

use crate::cli::HostOpts;
use crate::common::{
    bandwidth, config, config::BandwidthStatus, deps, display::DisplayServer, endpoint, output,
    signal, tunnel,
};

use self::pipeline::CaptureHandle;
use self::quality::EffectiveQuality;

/// Messages from per-viewer tasks (and the GUI command channel) to the
/// capture supervisor.
// The shared `Viewer` suffix is the point — these are all viewer lifecycle
// messages — so keep the descriptive names.
#[allow(clippy::enum_variant_names)]
enum SupervisorMsg {
    /// A new viewer wants in. Supervisor replies with the local capture HTTP
    /// port to connect to, or an error string if the host is full or capture
    /// spawn failed. `cancel` is the viewer's own token — the supervisor keeps
    /// it so a later `KickViewer` can tear this viewer's stream down.
    AddViewer {
        id: String,
        cancel: CancellationToken,
        reply: oneshot::Sender<Result<u16, String>>,
    },
    /// A viewer's session ended. Supervisor decrements the count and tears
    /// down capture if it just hit zero.
    RemoveViewer { id: String },
    /// Host asked (via the GUI command channel) to disconnect a viewer by
    /// endpoint id. Cancels that viewer's token; the normal teardown path then
    /// emits the `ViewerLeft`.
    KickViewer { id: String },
}

pub async fn run(opts: HostOpts) -> Result<()> {
    let display = DisplayServer::resolve(opts.display_server);
    deps::check_host_binaries(display, &opts)?;

    if display == DisplayServer::Unknown {
        bail!(
            "could not detect display server (WAYLAND_DISPLAY / DISPLAY / XDG_SESSION_TYPE all unset).\n\
             Use --display-server wayland|x11 to override."
        );
    }

    // Resolve quality first: Auto sizes its bandwidth budget against the viewer
    // cap the host will honor. To avoid a circular dependency (the auto-derived
    // cap itself depends on bitrate), Auto sizes against the user's explicit
    // --max-viewers when given, else a single viewer. The resulting effective
    // bitrate then feeds the cap resolution below.
    let sizing_viewers = opts.max_viewers.filter(|&n| n > 0).unwrap_or(1);
    let quality = quality::resolve(&opts, sizing_viewers);

    let resolution = resolve_max_viewers(&opts, quality.bitrate);
    if resolution.value == 0 {
        bail!("--max-viewers must be at least 1");
    }

    let cancel = signal::install_ctrl_c();

    // Phase 5 dry-run audit, off unless `PIXELPASS_AUDIO_AUDIT=1`. Read-only:
    // it observes the graph and logs what phases 2–4 conclude, creating no
    // links. Bound to a name so the handle lives as long as the run — dropping
    // it stops the observer thread and flushes the final O5 summary.
    let _audio_audit = audit::run::spawn_if_enabled()?;

    let endpoint = endpoint::bind(opts.relay.as_deref()).await?;

    // Relay-only ticket: wait for the home relay to connect, then keep only
    // the endpoint id + relay URL and drop the direct IP candidates. The relay
    // coordinates hole-punching to a direct path right after connect, so this
    // doesn't change whether peers can reach each other — it just keeps the
    // ticket short (~140 vs ~320 chars) and stops it from leaking LAN /
    // Docker-bridge addresses to whoever receives the ticket. Awaiting online()
    // first guarantees the relay URL is actually present (addr() right after
    // bind can return before the relay handshake completes); the 15s cap means
    // a relay outage degrades to a possibly-incomplete ticket rather than a hang
    // (n0 DNS discovery still resolves the id in that case).
    if tokio::time::timeout(Duration::from_secs(15), endpoint.online())
        .await
        .is_err()
    {
        tracing::warn!("home relay not connected within 15s; ticket may be incomplete");
    }
    let addr = endpoint.addr();
    let relay_only =
        EndpointAddr::new(addr.id).with_addrs(addr.addrs.iter().filter(|a| a.is_relay()).cloned());
    let ticket = EndpointTicket::new(relay_only);
    let ticket_str = ticket.to_string();
    let clipboard_ok = opts.interactive && copy_to_clipboard(&ticket_str);
    print_host_banner(&ticket, display, &opts, &quality, &resolution, clipboard_ok);

    output::emit(output::Event::Ticket { value: &ticket_str });
    let display_str = format!("{display:?}");
    let capture = capture_summary(&opts);
    let dims = quality.dimensions_summary();
    let cap_source = resolution.source.label();
    output::emit(output::Event::HostInfo {
        display_server: &display_str,
        capture: &capture,
        quality: &quality.label,
        dimensions: &dims,
        hw_encode: !opts.no_hwencode,
        max_viewers: resolution.value,
        max_viewers_source: &cap_source,
    });

    let (sup_tx, sup_rx) = mpsc::channel::<SupervisorMsg>(16);
    let supervisor = tokio::spawn(supervise(
        opts.clone(),
        quality,
        display,
        resolution.value,
        sup_rx,
    ));

    // Command channel for the GUI front-end: read `kick <endpoint-id>` lines
    // off stdin. Only when machine-driven (`--output json`) — a human host has
    // nothing to type here, and we don't want to swallow terminal input. Runs
    // on a plain OS thread (not a tokio task) so a read parked on stdin can't
    // hold up runtime shutdown on Ctrl+C; the thread dies with the process.
    if output::json_enabled() {
        spawn_kick_listener(sup_tx.clone());
    }

    accept_loop(&endpoint, sup_tx.clone(), cancel.clone()).await;

    drop(sup_tx);
    let _ = supervisor.await;

    endpoint.close().await;
    Ok(())
}

async fn accept_loop(
    endpoint: &Endpoint,
    sup_tx: mpsc::Sender<SupervisorMsg>,
    cancel: CancellationToken,
) {
    loop {
        tokio::select! {
            _ = cancel.cancelled() => {
                tracing::info!("cancellation requested — closing accept loop");
                return;
            }
            accepted = endpoint.accept() => {
                let Some(incoming) = accepted else {
                    tracing::info!("endpoint stopped accepting connections");
                    return;
                };
                let conn = match incoming.await {
                    Ok(c) => c,
                    Err(e) => {
                        tracing::warn!("incoming connection failed: {e:#}");
                        continue;
                    }
                };
                let sup_tx = sup_tx.clone();
                let cancel = cancel.clone();
                tokio::spawn(handle_peer(conn, sup_tx, cancel));
            }
        }
    }
}

async fn handle_peer(
    conn: Connection,
    sup_tx: mpsc::Sender<SupervisorMsg>,
    cancel: CancellationToken,
) {
    let remote = conn.remote_id();
    let id = remote.to_string();
    // This viewer's own kill switch: the supervisor holds a clone so a `kick`
    // can cancel it, and the stream select! below watches it.
    let peer_cancel = CancellationToken::new();

    let (reply_tx, reply_rx) = oneshot::channel();
    let add = SupervisorMsg::AddViewer {
        id: id.clone(),
        cancel: peer_cancel.clone(),
        reply: reply_tx,
    };
    if sup_tx.send(add).await.is_err() {
        tracing::warn!(%remote, "supervisor channel closed; dropping peer");
        return;
    }
    let port = match reply_rx.await {
        Ok(Ok(p)) => p,
        Ok(Err(reason)) => {
            tracing::warn!(%remote, %reason, "refusing viewer");
            eprintln!("[pixelpass] refusing viewer {remote}: {reason}");
            return;
        }
        Err(_) => {
            tracing::warn!(%remote, "supervisor reply dropped; dropping peer");
            return;
        }
    };

    let (quic_send, quic_recv) = match conn.accept_bi().await {
        Ok(s) => s,
        Err(e) => {
            tracing::warn!(%remote, "accept_bi failed: {e:#}");
            let _ = sup_tx.send(SupervisorMsg::RemoveViewer { id }).await;
            return;
        }
    };

    eprintln!("[pixelpass] viewer connected: {remote}");

    let tcp = match serve::connect_to_capture(port, Duration::from_secs(5)).await {
        Ok(t) => t,
        Err(e) => {
            tracing::warn!(%remote, "connect_to_capture failed: {e:#}");
            let _ = sup_tx.send(SupervisorMsg::RemoveViewer { id }).await;
            return;
        }
    };

    let bridge = tunnel::bridge(quic_send, quic_recv, tcp);
    tokio::select! {
        res = bridge => match res {
            Ok(()) => tracing::info!(%remote, "bridge closed cleanly"),
            Err(e) => tracing::info!(%remote, "bridge ended: {e:#}"),
        },
        _ = cancel.cancelled() => {
            tracing::info!(%remote, "cancellation during stream");
        }
        _ = peer_cancel.cancelled() => {
            tracing::info!(%remote, "kicked by host");
        }
    }

    eprintln!("[pixelpass] viewer disconnected: {remote}");
    let _ = sup_tx.send(SupervisorMsg::RemoveViewer { id }).await;
}

/// Read `kick <endpoint-id>` lines off stdin and forward them to the
/// supervisor. Runs on a detached OS thread (see the call site for why). Ends
/// when stdin hits EOF (the GUI closed the pipe) or the supervisor is gone.
fn spawn_kick_listener(sup_tx: mpsc::Sender<SupervisorMsg>) {
    use std::io::BufRead;
    std::thread::spawn(move || {
        let stdin = std::io::stdin();
        for line in stdin.lock().lines().map_while(Result::ok) {
            let Some(id) = line.trim().strip_prefix("kick ") else {
                continue;
            };
            let msg = SupervisorMsg::KickViewer {
                id: id.trim().to_string(),
            };
            // blocking_send is valid here: this is a plain thread, not inside
            // the tokio runtime. An Err means the supervisor closed — stop.
            if sup_tx.blocking_send(msg).is_err() {
                break;
            }
        }
    });
}

/// Owns the single shared CaptureHandle and the active viewer count. Spawns
/// capture lazily on the first AddViewer; tears it down when the count drops
/// back to zero. Enforces the max-viewers cap by refusing AddViewer when
/// the count is already at the cap.
async fn supervise(
    opts: HostOpts,
    quality: EffectiveQuality,
    display: DisplayServer,
    max_viewers: u32,
    mut rx: mpsc::Receiver<SupervisorMsg>,
) {
    let mut handle: Option<CaptureHandle> = None;
    // Active viewers, keyed by endpoint id, holding each one's kill switch.
    // The count is just `viewers.len()`. (A given endpoint connecting twice is
    // a non-case here: each viewer process uses a fresh ephemeral identity.)
    let mut viewers: HashMap<String, CancellationToken> = HashMap::new();

    while let Some(msg) = rx.recv().await {
        match msg {
            SupervisorMsg::AddViewer { id, cancel, reply } => {
                let count = viewers.len() as u32;
                if count >= max_viewers {
                    let reason =
                        format!("host is full ({count} of {max_viewers} viewers connected)");
                    output::emit(output::Event::ViewerRefused { reason: &reason });
                    let _ = reply.send(Err(reason));
                    continue;
                }

                if handle.is_none() {
                    tracing::info!("first viewer arriving — spawning capture");
                    match capture::spawn(display, &opts, &quality).await {
                        Ok(h) => {
                            handle = Some(h);
                            output::emit(output::Event::Capture {
                                state: output::CaptureState::Started,
                            });
                        }
                        Err(e) => {
                            let _ = reply.send(Err(format!("capture spawn failed: {e:#}")));
                            continue;
                        }
                    }
                }

                let port = handle.as_ref().expect("handle was just set").local_port();
                viewers.insert(id.clone(), cancel);
                let active = viewers.len() as u32;
                let _ = reply.send(Ok(port));
                output::emit(output::Event::ViewerJoined {
                    id: &id,
                    active,
                    max: max_viewers,
                });
                tracing::info!(active, cap = max_viewers, "viewer joined");
            }
            SupervisorMsg::RemoveViewer { id } => {
                // A given viewer task only ever sends RemoveViewer once, but the
                // map remove is the source of truth either way.
                if viewers.remove(&id).is_none() {
                    continue;
                }
                let active = viewers.len() as u32;
                output::emit(output::Event::ViewerLeft {
                    id: &id,
                    active,
                    max: max_viewers,
                });
                tracing::info!(active, cap = max_viewers, "viewer left");
                if active == 0
                    && let Some(h) = handle.take()
                {
                    tracing::info!("last viewer left — tearing down capture");
                    h.shutdown().await;
                    output::emit(output::Event::Capture {
                        state: output::CaptureState::Stopped,
                    });
                }
            }
            SupervisorMsg::KickViewer { id } => {
                match viewers.get(&id) {
                    // Cancel the viewer's token; its handle_peer select! wakes,
                    // sends RemoveViewer, and the leave is emitted there.
                    Some(cancel) => {
                        tracing::info!(%id, "kicking viewer");
                        cancel.cancel();
                    }
                    None => tracing::debug!(%id, "kick for unknown/already-gone viewer"),
                }
            }
        }
    }

    if let Some(h) = handle.take() {
        tracing::info!("host shutdown — tearing down capture");
        h.shutdown().await;
        output::emit(output::Event::Capture {
            state: output::CaptureState::Stopped,
        });
    }
}

fn print_host_banner(
    ticket: &EndpointTicket,
    display: DisplayServer,
    opts: &HostOpts,
    quality: &EffectiveQuality,
    resolution: &MaxViewersResolution,
    clipboard_ok: bool,
) {
    eprintln!();
    eprintln!("┌─ PixelPass · host ─────────────────────────────────────────");
    eprintln!("│  display server : {display:?}");
    eprintln!("│  capture        : {}", capture_summary(opts));
    eprintln!(
        "│  quality        : {} — {}",
        quality.label,
        quality.dimensions_summary()
    );
    eprintln!("│                   ({})", quality.note);
    eprintln!(
        "│  hw encode      : {}",
        if opts.no_hwencode {
            "off (software x264)"
        } else {
            "on (VAAPI H.264)"
        }
    );
    eprintln!(
        "│  max viewers    : {} ({})",
        resolution.value,
        resolution.source.label()
    );
    eprintln!("│");
    if clipboard_ok {
        eprintln!("│  Your share code has been copied to your clipboard.");
        eprintln!("│  Send it to your viewer(s). (If clipboard didn't work, the");
        eprintln!("│  code is also shown below for manual copy.)");
    } else {
        eprintln!("│  Share this ticket with your viewer(s):");
    }
    eprintln!("│");
    eprintln!("│      pixelpass {ticket}");
    eprintln!("│");
    eprintln!("│  Capture starts when the first viewer connects, runs while");
    eprintln!("│  any viewer is connected, and tears down when the last one");
    eprintln!("│  leaves. Press Ctrl+C to stop the host entirely.");
    eprintln!("└────────────────────────────────────────────────────────────");
    eprintln!();
}

/// How we arrived at the final viewer cap. Surfaced in the banner so the
/// user can tell at a glance whether the number is what they specified,
/// what their measured upstream supports, or just the fallback default.
struct MaxViewersResolution {
    value: u32,
    source: MaxViewersSource,
}

enum MaxViewersSource {
    /// User passed --max-viewers explicitly.
    UserFlag,
    /// Derived from the saved bandwidth measurement.
    BandwidthMeasurement { safe_mbps: f64 },
    /// No flag, no measurement — falling back.
    DefaultFallback,
}

impl MaxViewersSource {
    fn label(&self) -> String {
        match self {
            MaxViewersSource::UserFlag => "user-specified".to_string(),
            MaxViewersSource::BandwidthMeasurement { safe_mbps } => {
                format!("auto: {safe_mbps:.1} Mbps measured upstream")
            }
            MaxViewersSource::DefaultFallback => {
                "default — run `pixelpass --reconfigure` for a connection-aware value".to_string()
            }
        }
    }
}

fn resolve_max_viewers(opts: &HostOpts, effective_bitrate: u32) -> MaxViewersResolution {
    if let Some(n) = opts.max_viewers {
        return MaxViewersResolution {
            value: n,
            source: MaxViewersSource::UserFlag,
        };
    }
    if let Ok(cfg) = config::load()
        && cfg.bandwidth.status == BandwidthStatus::Measured
        && let Some(upstream) = cfg.bandwidth.upstream_mbps
    {
        let n = bandwidth::recommended_max_viewers(upstream, effective_bitrate);
        return MaxViewersResolution {
            value: n,
            source: MaxViewersSource::BandwidthMeasurement {
                safe_mbps: upstream,
            },
        };
    }
    MaxViewersResolution {
        value: 2,
        source: MaxViewersSource::DefaultFallback,
    }
}

fn copy_to_clipboard(text: &str) -> bool {
    match arboard::Clipboard::new().and_then(|mut cb| cb.set_text(text.to_owned())) {
        Ok(()) => true,
        Err(e) => {
            tracing::warn!("clipboard copy failed: {e}");
            false
        }
    }
}

fn capture_summary(opts: &HostOpts) -> String {
    let mut bits = vec![if opts.window { "window" } else { "fullscreen" }.to_string()];
    if let Some(app) = &opts.app {
        if opts.strict_audio {
            bits.push(format!("app-audio={app} (strict)"));
        } else {
            bits.push(format!("app-audio={app}"));
        }
    } else {
        bits.push("system-audio".to_string());
    }
    bits.join(" + ")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::Quality;

    fn opts(app: Option<&str>, strict_audio: bool) -> HostOpts {
        HostOpts {
            window: false,
            app: app.map(str::to_string),
            strict_audio,
            display_server: None,
            quality: Quality::Auto,
            bitrate: None,
            framerate: None,
            max_height: None,
            no_hwencode: false,
            max_viewers: None,
            interactive: false,
            relay: None,
        }
    }

    #[test]
    fn capture_summary_reflects_audio_mode() {
        assert_eq!(
            capture_summary(&opts(None, false)),
            "fullscreen + system-audio"
        );
        assert_eq!(
            capture_summary(&opts(Some("Firefox"), false)),
            "fullscreen + app-audio=Firefox"
        );
        // strict only shows when an app is selected.
        assert_eq!(
            capture_summary(&opts(Some("Firefox"), true)),
            "fullscreen + app-audio=Firefox (strict)"
        );
        assert_eq!(
            capture_summary(&opts(None, true)),
            "fullscreen + system-audio"
        );
    }

    #[test]
    fn initial_app_audio_is_lost_only_in_strict_app_mode() {
        use crate::common::output::AppAudioState;
        use crate::host::audio::initial_app_audio_state;
        // Strict + app: announce silence up front (loopback suppressed).
        assert_eq!(
            initial_app_audio_state(&opts(Some("Firefox"), true)),
            Some(AppAudioState::Lost)
        );
        // Best-effort app (no strict): loopback covers the gap → no initial event.
        assert_eq!(initial_app_audio_state(&opts(Some("Firefox"), false)), None);
        // Whole-desktop (strict is ignored without --app): no per-app events.
        assert_eq!(initial_app_audio_state(&opts(None, true)), None);
        assert_eq!(initial_app_audio_state(&opts(None, false)), None);
    }
}