moqtap-proxy 0.2.1

//! Transparent proxy orchestrator — accept loop and session management.
//!
//! The proxy binds a single listener that accepts raw-QUIC MoQT and
//! WebTransport clients simultaneously, negotiated via ALPN. Each
//! accepted connection is handed to a [`ProxySession`] that forwards
//! traffic to the configured upstream relay.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use tokio_util::sync::CancellationToken;

use crate::error::ProxyError;
use crate::event::{ProxyEvent, SessionId};
use crate::hook::{NoOpHook, ProxyHook};
use crate::listener::{AcceptedConn, Listener, ListenerConfig};
use crate::observer::ProxyObserver;
use crate::session::{ProxySession, ProxySessionConfig};

/// Configuration for the transparent proxy.
pub struct ProxyConfig {
    /// Listener configuration (bind address, certs).
    pub listener: ListenerConfig,
    /// Per-session configuration (upstream address, TLS, transport).
    pub session: ProxySessionConfig,
}

/// A transparent MoQT proxy that accepts client connections and forwards
/// traffic to an upstream relay.
///
/// Each accepted connection spawns a [`ProxySession`] that handles
/// bidirectional stream forwarding with inline MoQT frame parsing. The
/// client-facing transport (raw QUIC vs WebTransport) is chosen by the
/// client via ALPN and dispatched automatically — no configuration
/// required.
pub struct TransparentProxy {
    config: ProxyConfig,
    observer: Arc<dyn ProxyObserver>,
    hook: Arc<dyn ProxyHook>,
    cancel: CancellationToken,
    next_session_id: AtomicU64,
}

impl TransparentProxy {
    /// Create a new proxy with the given configuration and observer.
    pub fn new(config: ProxyConfig, observer: Arc<dyn ProxyObserver>) -> Self {
        Self {
            config,
            observer,
            hook: Arc::new(NoOpHook),
            cancel: CancellationToken::new(),
            next_session_id: AtomicU64::new(1),
        }
    }

    /// Create a new proxy with a custom hook for frame mutation.
    pub fn with_hook(
        config: ProxyConfig,
        observer: Arc<dyn ProxyObserver>,
        hook: Arc<dyn ProxyHook>,
    ) -> Self {
        Self {
            config,
            observer,
            hook,
            cancel: CancellationToken::new(),
            next_session_id: AtomicU64::new(1),
        }
    }

    /// Returns a cancellation token that can be used to trigger shutdown.
    pub fn cancel_token(&self) -> CancellationToken {
        self.cancel.clone()
    }

    /// Run the proxy accept loop. Blocks until cancelled or a fatal
    /// listener error occurs.
    pub async fn run(&self) -> Result<(), ProxyError> {
        let listener = Listener::bind(ListenerConfig {
            bind_addr: self.config.listener.bind_addr,
            cert_chain: self.config.listener.cert_chain.clone(),
            key_der: self.config.listener.key_der.clone_key(),
        })?;

        loop {
            tokio::select! {
                result = listener.accept() => {
                    self.dispatch(result?);
                }
                _ = self.cancel.cancelled() => {
                    listener.close();
                    return Ok(());
                }
            }
        }
    }

    /// Spawn a session for an accepted connection, picking the right
    /// entry point based on the negotiated transport.
    fn dispatch(&self, accepted: AcceptedConn) {
        match accepted {
            AcceptedConn::Quic { conn, alpn } => {
                let session_id = self.next_session_id();
                let client_addr = conn.remote_address();
                self.emit_session_started(session_id, client_addr, "QUIC");
                let session = self.new_session(session_id, alpn);
                tokio::spawn(async move {
                    let _ = session.run(conn).await;
                });
            }
            #[cfg(feature = "webtransport")]
            AcceptedConn::WebTransport(conn) => {
                let session_id = self.next_session_id();
                let client_addr = conn.remote_address();
                self.emit_session_started(session_id, client_addr, "WebTransport");
                // WebTransport carries no moqt-* ALPN (always h3) so the
                // session falls back to config.draft and/or SETUP-peek.
                let session = self.new_session(session_id, Vec::new());
                tokio::spawn(async move {
                    let _ = session.run_webtransport(conn).await;
                });
            }
        }
    }

    // ── Helpers ────────────────────────────────────────────────

    fn next_session_id(&self) -> SessionId {
        SessionId(self.next_session_id.fetch_add(1, Ordering::Relaxed))
    }

    fn emit_session_started(
        &self,
        session_id: SessionId,
        client_addr: std::net::SocketAddr,
        client_transport: &str,
    ) {
        if self.observer.wants_events() {
            self.observer.on_event(&ProxyEvent::SessionStarted {
                session_id,
                client_addr,
                client_transport: client_transport.to_string(),
            });
        }
    }

    fn new_session(&self, session_id: SessionId, client_alpn: Vec<u8>) -> ProxySession {
        ProxySession::new(
            session_id,
            ProxySessionConfig {
                draft: self.config.session.draft,
                upstream_transport: self.config.session.upstream_transport.clone(),
                upstream_addr: self.config.session.upstream_addr.clone(),
                skip_upstream_cert_verify: self.config.session.skip_upstream_cert_verify,
                upstream_ca_certs: self.config.session.upstream_ca_certs.clone(),
                upstream_connect_timeout_secs: self.config.session.upstream_connect_timeout_secs,
            },
            client_alpn,
            Arc::clone(&self.observer),
            Arc::clone(&self.hook),
            self.cancel.child_token(),
        )
    }
}