lanspread/crates/lanspread-peer/src/services/server.rs

//! QUIC server accept loop.

use std::{future::Future, net::SocketAddr, panic::AssertUnwindSafe, sync::Arc, time::Duration};

use futures::FutureExt as _;
use s2n_quic::{
    Connection,
    Server,
    application,
    provider::endpoint_limits,
    stream::BidirectionalStream,
};
use tokio::{
    sync::{OwnedSemaphorePermit, Semaphore, oneshot},
    task::JoinSet,
};
use tokio_util::sync::CancellationToken;

use crate::{
    context::PeerCtx,
    library::prime_library_manifests,
    quic_runtime::{quic_congestion_controller, quic_server_limits, tracked_quic_io},
    scoped_blocking::scoped_blocking,
    services::{
        advertise::{close_mdns_advertiser, monitor_mdns_events, start_mdns_advertiser},
        stream::handle_peer_stream,
    },
    tls,
};

/// Limits unauthenticated handshake memory before QUIC admission completes.
const MAX_INFLIGHT_HANDSHAKES: usize = 64;
/// Limits established connection scopes owned by the application accept loop.
const MAX_ESTABLISHED_CONNECTIONS: usize = 64;
/// Mirrors the transport stream limit and bounds application stream futures.
const MAX_CONTROL_STREAM_TASKS: usize = 32;
/// Server-wide cap acquired before any length-delimited decoder is allocated.
const MAX_GLOBAL_CONTROL_STREAM_TASKS: usize = 64;
/// Long-lived transfers move from the decoder pool to this smaller pool so
/// saturated bulk egress cannot consume every control-plane permit.
const MAX_GLOBAL_BULK_TRANSFER_TASKS: usize = 48;
/// Application idle bound for an established connection with no active
/// request streams. This is independent of transport keepalive traffic.
const CONNECTION_NO_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(10);

struct BoundedEndpointLimits {
    inner: endpoint_limits::Default,
}

impl endpoint_limits::Limiter for BoundedEndpointLimits {
    fn on_connection_attempt(
        &mut self,
        info: &endpoint_limits::ConnectionAttempt<'_>,
    ) -> endpoint_limits::Outcome {
        if !endpoint_connection_capacity_available(info.connection_count) {
            return endpoint_limits::Outcome::close();
        }
        endpoint_limits::Limiter::on_connection_attempt(&mut self.inner, info)
    }
}

const fn endpoint_connection_capacity_available(connection_count: usize) -> bool {
    connection_count < MAX_ESTABLISHED_CONNECTIONS
}

fn bounded_endpoint_limits() -> eyre::Result<BoundedEndpointLimits> {
    Ok(BoundedEndpointLimits {
        inner: endpoint_limits::Default::builder()
            .with_inflight_handshake_limit(MAX_INFLIGHT_HANDSHAKES)?
            .build()?,
    })
}

/// Runs the QUIC server and mDNS advertiser.
pub async fn run_server_component(
    addr: SocketAddr,
    ctx: PeerCtx,
    ready: oneshot::Sender<SocketAddr>,
) -> eyre::Result<()> {
    // Manifest bodies may be disk-backed on their first access. Resolve only
    // the currently publishable local set before opening the public endpoint;
    // this bounds retained manifest memory by actual local availability.
    let publication = {
        let library = ctx.local_library.read().await;
        library.publication(ctx.catalog.catalog())
    };
    let catalog = Arc::clone(&ctx.catalog);
    scoped_blocking(move || prime_library_manifests(&publication.game_ids, &catalog))?;

    let (io, endpoint_control) = tracked_quic_io(addr)?;

    let server = Server::builder()
        .with_tls(tls::server_provider(&ctx.peer_identity)?)?
        .with_io(io)?
        .with_endpoint_limits(bounded_endpoint_limits()?)?
        .with_limits(quic_server_limits()?)?
        .with_congestion_controller(quic_congestion_controller())?
        .start()?;
    let endpoint_task = endpoint_control.take_started()?;

    run_body_with_cleanup(
        run_server_body(server, ctx, ready),
        endpoint_task.shutdown_and_join(),
    )
    .await
}

async fn run_server_body(
    mut server: Server,
    ctx: PeerCtx,
    ready: oneshot::Sender<SocketAddr>,
) -> eyre::Result<()> {
    let server_addr = server.local_addr()?;
    log::info!("Peer server listening on {server_addr}");

    let mdns_advertiser = start_mdns_advertiser(&ctx, server_addr).await?;
    let mdns_monitor = mdns_advertiser.monitor.clone();
    let server_children_shutdown = ctx.shutdown.child_token();
    let mut mdns_tasks = JoinSet::new();
    let mdns_shutdown = server_children_shutdown.clone();
    mdns_tasks.spawn(async move {
        monitor_mdns_events(mdns_monitor, mdns_shutdown).await;
        Ok(())
    });
    let mut connection_tasks = JoinSet::new();
    let control_stream_permits = Arc::new(Semaphore::new(MAX_GLOBAL_CONTROL_STREAM_TASKS));
    let bulk_transfer_permits = Arc::new(Semaphore::new(MAX_GLOBAL_BULK_TRANSFER_TASKS));
    let ready_addr =
        (*ctx.local_peer_addr.read().await).unwrap_or_else(|| direct_connect_addr(server_addr));

    ready
        .send(ready_addr)
        .map_err(|_| eyre::eyre!("network manager stopped before server readiness"))?;

    let server_result = match AssertUnwindSafe(async {
        loop {
            tokio::select! {
                biased;
                () = ctx.shutdown.cancelled() => break Ok(()),
                result = mdns_tasks.join_next(), if !mdns_tasks.is_empty() => {
                    log_joined_child_result(
                        "mDNS monitor",
                        result.expect("non-empty child set"),
                    );
                    log::warn!("mDNS monitor ended while the QUIC server is still running");
                }
                result = connection_tasks.join_next(), if !connection_tasks.is_empty() => {
                    log_joined_child_result(
                        "peer connection",
                        result.expect("non-empty child set"),
                    );
                }
                connection = server.accept() => {
                    let Some(connection) = connection else {
                        break Err(eyre::eyre!("QUIC server accept loop ended unexpectedly"));
                    };

                    if !has_child_capacity(connection_tasks.len(), MAX_ESTABLISHED_CONNECTIONS) {
                        log::warn!(
                            "Closing excess peer connection from {} at application limit {}",
                            connection.remote_addr().map_or_else(
                                |_| "unknown".to_owned(),
                                |addr| addr.to_string(),
                            ),
                            MAX_ESTABLISHED_CONNECTIONS,
                        );
                        connection.close(application::Error::UNKNOWN);
                        continue;
                    }

                    connection_tasks.spawn(handle_peer_connection(
                        connection,
                        ctx.clone(),
                        server_children_shutdown.clone(),
                        Arc::clone(&control_stream_permits),
                        Arc::clone(&bulk_transfer_permits),
                    ));
                }
            }
        }
    })
    .catch_unwind()
    .await
    {
        Ok(result) => result,
        Err(payload) => Err(eyre::eyre!(
            "QUIC server loop panicked: {}",
            panic_payload_to_string(payload.as_ref())
        )),
    };

    // Stop future work before dropping the accept owner. Connection children
    // cooperatively observe this token and close their stream scopes before
    // they return.
    server_children_shutdown.cancel();
    drop(server);
    drain_joined_child_tasks(&mut connection_tasks, "peer connection").await;
    drain_joined_child_tasks(&mut mdns_tasks, "mDNS monitor").await;
    let mdns_close_result = close_mdns_advertiser(mdns_advertiser);
    combine_server_results(server_result, mdns_close_result, Ok(()))
}

async fn run_body_with_cleanup<Body, Cleanup>(body: Body, cleanup: Cleanup) -> eyre::Result<()>
where
    Body: Future<Output = eyre::Result<()>>,
    Cleanup: Future<Output = eyre::Result<()>>,
{
    let body_result = match AssertUnwindSafe(body).catch_unwind().await {
        Ok(result) => result,
        Err(payload) => Err(eyre::eyre!(
            "QUIC server body panicked: {}",
            panic_payload_to_string(payload.as_ref())
        )),
    };

    // This is the unconditional endpoint-owner epilogue. Even an unexpected
    // panic during setup, serving, or descendant cleanup cannot skip the join.
    let endpoint_result = cleanup.await;
    combine_server_results(body_result, Ok(()), endpoint_result)
}

fn combine_server_results(
    server: eyre::Result<()>,
    mdns: eyre::Result<()>,
    endpoint: eyre::Result<()>,
) -> eyre::Result<()> {
    let mut errors = Vec::new();
    if let Err(error) = server {
        errors.push(format!("QUIC server failed: {error:#}"));
    }
    if let Err(error) = mdns {
        errors.push(format!("mDNS advertiser shutdown failed: {error:#}"));
    }
    if let Err(error) = endpoint {
        errors.push(format!("QUIC endpoint shutdown failed: {error:#}"));
    }

    if errors.is_empty() {
        Ok(())
    } else {
        Err(eyre::eyre!(errors.join("; ")))
    }
}

fn panic_payload_to_string(payload: &(dyn std::any::Any + Send)) -> String {
    if let Some(message) = payload.downcast_ref::<&'static str>() {
        return (*message).to_string();
    }
    if let Some(message) = payload.downcast_ref::<String>() {
        return message.clone();
    }
    "unknown panic payload".to_string()
}

fn direct_connect_addr(server_addr: SocketAddr) -> SocketAddr {
    if server_addr.ip().is_unspecified() {
        return SocketAddr::from(([127, 0, 0, 1], server_addr.port()));
    }
    server_addr
}

async fn handle_peer_connection(
    mut connection: Connection,
    ctx: PeerCtx,
    server_shutdown: CancellationToken,
    control_stream_permits: Arc<Semaphore>,
    bulk_transfer_permits: Arc<Semaphore>,
) -> eyre::Result<()> {
    let remote_addr = connection.remote_addr()?;
    log::info!("{remote_addr} peer connected");

    let connection_shutdown = server_shutdown.child_token();
    let mut stream_tasks = JoinSet::new();
    let connection_result = match AssertUnwindSafe(async {
        loop {
            tokio::select! {
                biased;
                () = connection_shutdown.cancelled() => break Ok(()),
                result = stream_tasks.join_next(), if !stream_tasks.is_empty() => {
                    log_joined_child_result(
                        &format!("{remote_addr} peer stream"),
                        result.expect("non-empty child set"),
                    );
                }
                () = tokio::time::sleep(CONNECTION_NO_STREAM_IDLE_TIMEOUT),
                    if stream_tasks.is_empty() => {
                    log::debug!(
                        "Closing idle peer connection from {remote_addr} after {CONNECTION_NO_STREAM_IDLE_TIMEOUT:?}"
                    );
                    break Ok(());
                }
                stream = connection.accept_bidirectional_stream() => {
                    match stream {
                        Ok(Some(mut stream)) => {
                            if !has_child_capacity(stream_tasks.len(), MAX_CONTROL_STREAM_TASKS) {
                                let _ = stream.stop_sending(application::Error::UNKNOWN);
                                let _ = stream.reset(application::Error::UNKNOWN);
                                continue;
                            }
                            let Ok(control_permit) = Arc::clone(&control_stream_permits)
                                .try_acquire_owned()
                            else {
                                let _ = stream.stop_sending(application::Error::UNKNOWN);
                                let _ = stream.reset(application::Error::UNKNOWN);
                                continue;
                            };
                            let stream_ctx = ctx.clone();
                            let stream_shutdown = connection_shutdown.child_token();
                            stream_tasks.spawn(handle_admitted_peer_stream(
                                stream,
                                stream_ctx,
                                Some(remote_addr),
                                stream_shutdown,
                                control_permit,
                                Arc::clone(&bulk_transfer_permits),
                            ));
                        }
                        Ok(None) => break Ok(()),
                        Err(error) => break Err(error.into()),
                    }
                }
            }
        }
    })
    .catch_unwind()
    .await
    {
        Ok(result) => result,
        Err(payload) => Err(eyre::eyre!(
            "{remote_addr} peer connection scope panicked: {}",
            panic_payload_to_string(payload.as_ref())
        )),
    };

    // Cancel the connection-local scope before closing its QUIC owner. Every
    // accepted stream and outbound transfer derives from this token, so child
    // futures can settle without waiting for process-wide shutdown. Closing
    // the connection also wakes any transport operation already in progress.
    connection_shutdown.cancel();
    connection.close(0u32.into());
    drop(connection);
    let stream_label = format!("{remote_addr} peer stream");
    drain_joined_child_tasks(&mut stream_tasks, &stream_label).await;

    log::info!("{remote_addr} peer disconnected");
    connection_result
}

async fn handle_admitted_peer_stream(
    stream: BidirectionalStream,
    ctx: PeerCtx,
    remote_addr: Option<SocketAddr>,
    stream_shutdown: CancellationToken,
    control_permit: OwnedSemaphorePermit,
    bulk_transfer_permits: Arc<Semaphore>,
) -> eyre::Result<()> {
    handle_peer_stream(
        stream,
        ctx,
        remote_addr,
        stream_shutdown,
        control_permit,
        bulk_transfer_permits,
    )
    .await
}

const fn has_child_capacity(active: usize, limit: usize) -> bool {
    active < limit
}

fn log_child_result(label: &str, result: eyre::Result<()>) {
    if let Err(error) = result {
        log::error!("{label} error: {error}");
    }
}

fn log_joined_child_result(label: &str, result: Result<eyre::Result<()>, tokio::task::JoinError>) {
    match result {
        Ok(result) => log_child_result(label, result),
        Err(error) => log::error!("{label} task failed: {error}"),
    }
}

async fn drain_joined_child_tasks(children: &mut JoinSet<eyre::Result<()>>, label: &str) {
    while let Some(result) = children.join_next().await {
        log_joined_child_result(label, result);
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{
            Arc,
            atomic::{AtomicUsize, Ordering},
        },
        time::Duration,
    };

    use tokio::{
        sync::{Semaphore, mpsc},
        task::JoinSet,
    };
    use tokio_util::sync::CancellationToken;

    use super::{
        CONNECTION_NO_STREAM_IDLE_TIMEOUT,
        MAX_CONTROL_STREAM_TASKS,
        MAX_ESTABLISHED_CONNECTIONS,
        MAX_GLOBAL_BULK_TRANSFER_TASKS,
        MAX_GLOBAL_CONTROL_STREAM_TASKS,
        MAX_INFLIGHT_HANDSHAKES,
        bounded_endpoint_limits,
        drain_joined_child_tasks,
        endpoint_connection_capacity_available,
        has_child_capacity,
        run_body_with_cleanup,
    };

    #[test]
    fn unauthenticated_runtime_bounds_are_explicit_and_closed_at_capacity() {
        assert_eq!(MAX_INFLIGHT_HANDSHAKES, 64);
        assert_eq!(MAX_ESTABLISHED_CONNECTIONS, 64);
        assert_eq!(MAX_CONTROL_STREAM_TASKS, 32);
        assert_eq!(MAX_GLOBAL_CONTROL_STREAM_TASKS, 64);
        assert_eq!(MAX_GLOBAL_BULK_TRANSFER_TASKS, 48);
        assert_eq!(
            u64::try_from(MAX_CONTROL_STREAM_TASKS).expect("stream task bound fits u64"),
            crate::quic_runtime::MAX_OPEN_BIDIRECTIONAL_STREAMS,
        );
        assert!(has_child_capacity(63, MAX_ESTABLISHED_CONNECTIONS));
        assert!(!has_child_capacity(64, MAX_ESTABLISHED_CONNECTIONS));
        assert!(!has_child_capacity(32, MAX_CONTROL_STREAM_TASKS));
    }

    #[test]
    fn endpoint_limiter_rejects_before_the_internal_accept_queue_can_exceed_the_cap() {
        bounded_endpoint_limits().expect("endpoint limiter should build");
        assert!(endpoint_connection_capacity_available(
            MAX_ESTABLISHED_CONNECTIONS - 1
        ));
        assert!(!endpoint_connection_capacity_available(
            MAX_ESTABLISHED_CONNECTIONS
        ));
    }

    #[test]
    fn global_control_permit_saturates_and_releases_across_scopes() {
        let permits = Arc::new(Semaphore::new(MAX_GLOBAL_CONTROL_STREAM_TASKS));
        let held = (0..MAX_GLOBAL_CONTROL_STREAM_TASKS)
            .map(|_| {
                Arc::clone(&permits)
                    .try_acquire_owned()
                    .expect("configured global permit should be available")
            })
            .collect::<Vec<_>>();
        assert!(Arc::clone(&permits).try_acquire_owned().is_err());

        drop(held);
        assert!(Arc::clone(&permits).try_acquire_owned().is_ok());
    }

    #[test]
    fn saturated_bulk_pool_preserves_control_plane_permits() {
        let control = Arc::new(Semaphore::new(MAX_GLOBAL_CONTROL_STREAM_TASKS));
        let bulk = Arc::new(Semaphore::new(MAX_GLOBAL_BULK_TRANSFER_TASKS));
        let held_bulk = (0..MAX_GLOBAL_BULK_TRANSFER_TASKS)
            .map(|_| {
                Arc::clone(&bulk)
                    .try_acquire_owned()
                    .expect("configured bulk permit should be available")
            })
            .collect::<Vec<_>>();
        assert!(Arc::clone(&bulk).try_acquire_owned().is_err());
        assert!(
            Arc::clone(&control).try_acquire_owned().is_ok(),
            "bulk saturation must not consume control-plane admission"
        );
        drop(held_bulk);
    }

    #[tokio::test]
    async fn draining_waits_for_every_childs_natural_cleanup() {
        let owner_closed = CancellationToken::new();
        let cleanup_release = CancellationToken::new();
        let completed = Arc::new(AtomicUsize::new(0));
        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
        let mut children = JoinSet::new();

        for child_id in 0..2 {
            let owner_closed = owner_closed.clone();
            let cleanup_release = cleanup_release.clone();
            let completed = completed.clone();
            let started_tx = started_tx.clone();
            children.spawn(async move {
                started_tx
                    .send(child_id)
                    .expect("test observer should remain available");
                owner_closed.cancelled().await;
                cleanup_release.cancelled().await;
                completed.fetch_add(1, Ordering::SeqCst);
                Ok(())
            });
        }
        drop(started_tx);

        let drain_task = tokio::spawn(async move {
            let mut children = children;
            drain_joined_child_tasks(&mut children, "synthetic child").await;
        });
        for _ in 0..2 {
            tokio::time::timeout(Duration::from_secs(1), started_rx.recv())
                .await
                .expect("child should start")
                .expect("start channel should remain open");
        }

        owner_closed.cancel();
        tokio::task::yield_now().await;
        assert_eq!(completed.load(Ordering::SeqCst), 0);
        assert!(!drain_task.is_finished(), "drain must await child cleanup");

        cleanup_release.cancel();
        tokio::time::timeout(Duration::from_secs(1), drain_task)
            .await
            .expect("drain should finish after cleanup is released")
            .expect("drain task should not panic");
        assert_eq!(completed.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn one_panicking_child_does_not_skip_a_siblings_cleanup() {
        let cleanup_release = CancellationToken::new();
        let completed = Arc::new(AtomicUsize::new(0));
        let mut children = JoinSet::new();
        children.spawn(async {
            panic!("injected child panic");
            #[allow(unreachable_code)]
            Ok(())
        });
        let child_release = cleanup_release.clone();
        let child_completed = completed.clone();
        children.spawn(async move {
            child_release.cancelled().await;
            child_completed.fetch_add(1, Ordering::SeqCst);
            Ok(())
        });

        let drain_task = tokio::spawn(async move {
            drain_joined_child_tasks(&mut children, "synthetic joined child").await;
        });
        tokio::task::yield_now().await;
        assert!(!drain_task.is_finished());
        cleanup_release.cancel();
        tokio::time::timeout(Duration::from_secs(1), drain_task)
            .await
            .expect("drain should await the surviving child")
            .expect("drain task should not panic");
        assert_eq!(completed.load(Ordering::SeqCst), 1);
    }

    #[tokio::test(start_paused = true)]
    async fn an_established_connection_without_streams_has_a_finite_application_idle_bound() {
        let started = tokio::time::Instant::now();
        tokio::time::sleep(CONNECTION_NO_STREAM_IDLE_TIMEOUT).await;
        assert_eq!(started.elapsed(), CONNECTION_NO_STREAM_IDLE_TIMEOUT);
    }

    #[tokio::test]
    async fn server_body_panic_still_awaits_owner_cleanup() {
        let cleanup_finished = Arc::new(AtomicUsize::new(0));
        let cleanup_probe = cleanup_finished.clone();

        let result = run_body_with_cleanup(
            async {
                panic!("injected server body panic");
                #[allow(unreachable_code)]
                Ok(())
            },
            async move {
                tokio::task::yield_now().await;
                cleanup_probe.store(1, Ordering::SeqCst);
                Ok(())
            },
        )
        .await;

        assert!(result.is_err());
        assert_eq!(cleanup_finished.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn nested_scopes_publish_completion_from_children_outward() {
        let global_shutdown = CancellationToken::new();
        let server_shutdown = global_shutdown.child_token();
        let stream_cleanup_release = CancellationToken::new();
        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
        let (order_tx, mut order_rx) = mpsc::unbounded_channel();
        let mut connections = JoinSet::new();

        let connection_shutdown = server_shutdown.child_token();
        let connection_stream_cleanup_release = stream_cleanup_release.clone();
        let connection_order_tx = order_tx.clone();
        connections.spawn(async move {
            let mut streams = JoinSet::new();
            for label in ["stream-a", "stream-b"] {
                let shutdown = connection_shutdown.child_token();
                let cleanup_release = connection_stream_cleanup_release.clone();
                let started_tx = started_tx.clone();
                let order_tx = connection_order_tx.clone();
                streams.spawn(async move {
                    started_tx
                        .send(label)
                        .expect("test observer should remain available");
                    shutdown.cancelled().await;
                    cleanup_release.cancelled().await;
                    order_tx
                        .send(label)
                        .expect("order observer should remain available");
                    Ok(())
                });
            }

            connection_shutdown.cancelled().await;
            drain_joined_child_tasks(&mut streams, "synthetic stream").await;
            connection_order_tx
                .send("peer-disconnected")
                .expect("order observer should remain available");
            Ok(())
        });

        let hierarchy_task = tokio::spawn(async move {
            server_shutdown.cancel();
            drain_joined_child_tasks(&mut connections, "synthetic connection").await;
            order_tx
                .send("server-return")
                .expect("order observer should remain available");
        });
        for _ in 0..2 {
            tokio::time::timeout(Duration::from_secs(1), started_rx.recv())
                .await
                .expect("stream should start")
                .expect("start channel should remain open");
        }
        assert!(!hierarchy_task.is_finished());

        stream_cleanup_release.cancel();
        tokio::time::timeout(Duration::from_secs(1), hierarchy_task)
            .await
            .expect("hierarchy should drain")
            .expect("hierarchy task should not panic");

        let mut order = Vec::new();
        while let Ok(event) = order_rx.try_recv() {
            order.push(event);
        }
        let disconnected = order
            .iter()
            .position(|event| *event == "peer-disconnected")
            .expect("disconnect should be published");
        let server_return = order
            .iter()
            .position(|event| *event == "server-return")
            .expect("server return should be published");
        assert!(
            order[..disconnected]
                .iter()
                .all(|event| event.starts_with("stream-"))
        );
        assert_eq!(disconnected, 2);
        assert_eq!(server_return, 3);
        assert!(
            !global_shutdown.is_cancelled(),
            "server-local shutdown must not cancel its runtime parent"
        );
    }
}