scion-quic 0.6.1

// Copyright 2026 Anapaya Systems
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
//! A sans-I/O inspired implementation of a QUIC SCION-server.
use std::{
    collections::{
        BTreeSet, HashMap,
        hash_map::{Entry, OccupiedEntry},
    },
    net::{IpAddr, SocketAddr},
    ops::Range,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::Instant,
};

use futures::{StreamExt, stream::FuturesUnordered};
use prometheus::IntGauge;
use ring::hmac::Key as HmacKey;
use sciparse::address::ip_socket_addr::ScionSocketIpAddr;
use squiche::{ConnectionId, RecvInfo};
use thiserror::Error;
use tokio::{
    sync::Notify,
    task::{JoinError, JoinHandle},
};
use tokio_util::sync::CancellationToken;

use crate::{
    app::{NoApp, QuicScionApplication},
    quic::connection::{
        ConnectionHandle, IsdAsnPair, QuicScionConn, QuicScionConnDriver, ScionSendInfo,
    },
    socket::{BoxedSocketError, GenericScionUdpSocket},
};

const DIGEST_LEN: usize = ring::digest::SHA256_OUTPUT_LEN;
const HMAC_ALGO: ring::hmac::Algorithm = ring::hmac::HMAC_SHA256;
const MIN_ODCID_LEN: usize = 0;
const MIN_TOKEN_LEN: usize = MIN_ODCID_LEN + DIGEST_LEN;
const MAX_TOKEN_LEN: usize = squiche::MAX_CONN_ID_LEN + DIGEST_LEN;

/// A state machine that processes incoming QUIC packets.
///
/// ## Design goals
///
/// The goal is to separate the concerns of connection establishment and packet
/// routing from connection upgrades (HTTP/3) and the concurrency model.
///
/// ### Packet routing
///
/// After a connection has been established, it is the client's responsibility
/// to register the established connection with the endpoint. The endpoint then
/// routes incoming packets to the respective connection.
///
/// ### Concurrency model
///
/// The endpoint remains _un_opiniated wrt. queuing behavior. The
/// [`QuicScionServerEndpoint`] can be shared with concurrent connection drivers
/// or integrated into an actor-based model.
///
/// ## Limitations
///
/// * **0-RTT-support**: One of the most common use cases that we anticipate is (C)RPC interfaces
///   which typically employ `POST`-HTTP-requests as their HTTP-transport method. `POST` methods are
///   not indempotent and therefore cannot appear as part of the early data.
/// * Outgoing packets generated by the endpoint driver are not paced.
///
/// ## Security & connection establishment
///
/// Upon connecting, the client is forced to reply to a retry token and thus
/// prove possession of its source address. This prevents trivial reflection
/// attacks, at the cost of an additional round trip.
///
/// ## CID construction and connection information
///
/// The CID chosen by the server binds the connection to the SCION socket
/// address of the remote (including the AS number). This naturally avoids
/// collisions in case two remote endhosts have the same IP address, but
/// different ASes. However, note that addresses reported in logs, etc. are
/// still just the IP-addresses of the remote endhosts, as the underlying
/// `squiche` library uses standard socket addresses.
pub struct QuicScionServerEndpoint<T> {
    established: HashMap<ConnectionId<'static>, T>,
    establishing_set: EstablishingSet,
    local_addr: ScionSocketIpAddr,
    cid_generator: CidGenerator,
    token_generator: TokenGenerator,
    // Scrub space for token generation. Avoids repeated allocations.
    token_scrub_space: [u8; MAX_TOKEN_LEN],
    config: squiche::Config,
    metrics: Metrics,
}

impl<T> QuicScionServerEndpoint<T> {
    /// Creates a new [`QuicScionServerEndpoint`].
    ///
    /// ## Parameters
    ///
    /// * `rnd_seed` is used to seed the token and CID generator.
    /// * `config` configures new connections.
    /// * `local_addr` is the address of the local listening socket.
    pub fn new(
        rnd_seed: [u8; 32],
        config: squiche::Config,
        local_addr: ScionSocketIpAddr,
        metrics: Metrics,
    ) -> Self {
        let key = HmacKey::new(HMAC_ALGO, &rnd_seed);
        let mut seed1 = [0u8; 32];
        let mut seed2 = [0u8; 32];
        seed1.copy_from_slice(ring::hmac::sign(&key, &[0x01]).as_ref());
        seed2.copy_from_slice(ring::hmac::sign(&key, &[0x02]).as_ref());

        Self {
            established: Default::default(),
            establishing_set: Default::default(),
            local_addr,
            cid_generator: CidGenerator::new(seed1),
            token_generator: TokenGenerator::new(seed2),
            token_scrub_space: [0u8; MAX_TOKEN_LEN],
            config,
            metrics,
        }
    }

    /// Processes a single incoming QUIC packet.
    ///
    /// If a new connection is returned from `recv`, it is the callers
    /// responsibility to immediately dispatch packets sent by the
    /// connection. Omitting this might lead to spurious timeouts.
    ///
    /// ## Parameters
    ///
    /// * `recv_buf`: A reference to the buffer that contains the newly received QUIC packet.
    /// * `send_buf`: If `Ok(RecvResult::Send(_))` is returned, contains the QUIC packet to be sent
    ///   back to the source.
    /// * `from`: The address of the remote the packet was received from.
    pub fn recv(
        &mut self,
        recv_buf: &mut [u8],
        send_buf: &mut [u8],
        from: ScionSocketIpAddr,
    ) -> RecvResult<'_, T> {
        // Parse QUIC Header
        let hdr = squiche::Header::from_slice(recv_buf, squiche::MAX_CONN_ID_LEN)
            .map_err(PacketProcessError::InvalidHeader)?;

        tracing::trace!(?hdr.scid, ?hdr.dcid, ?from, "Received QUIC packet");

        // check if the connection is established
        if let Entry::Occupied(e) = self.established.entry(hdr.dcid.clone()) {
            return Ok(RecvOutcome::ConnEvent(e.into_mut()));
        }

        let local_addr = self.local_addr.socket_addr();
        let remote_addr = from.socket_addr();

        let res = self.establishing_set.update(&hdr.dcid, |c| {
            c.inner
                .recv(
                    recv_buf,
                    RecvInfo {
                        from: remote_addr,
                        to: local_addr,
                    },
                )
                .map(|_| ())
        })?;
        self.metrics
            .establishing_connections_gauge
            .set(self.establishing_set.conn_map.len() as i64);
        match res {
            EstablishingOutcome::Established(scion_quic_conn) => {
                return Ok(RecvOutcome::EstablishedConn(scion_quic_conn));
            }
            EstablishingOutcome::Establishing(cid) => {
                // SAFETY: Just updated the entry
                let c = self.establishing_set.conn_map.get_mut(&cid).unwrap();
                return Ok(RecvOutcome::Establishing(EstablishingScionQuicConn(c)));
            }
            EstablishingOutcome::Done => {}
        }

        // If the packets belongs to a connection that is currently being
        // established, it is received on that connection. If the connection
        // transitions to is_established, it is returned. Otherwise, a
        // reference is returned asking the caller to send all queued
        // outgoing packets.

        // At this point, we can safety ignore non-initial packets because we
        // forced the client to send all possible (out-of-order) fragments under
        // the same, server chosen DCID.
        if hdr.ty != squiche::Type::Initial {
            return Err(PacketProcessError::ExpectedInitialPacket(hdr.dcid));
        }

        // The send information used for immediate responses.
        let send_info = ScionSendInfo::new(self.local_addr, from, Instant::now());
        // Check version support
        if !squiche::version_is_supported(hdr.version) {
            let len = squiche::negotiate_version(&hdr.scid, &hdr.dcid, send_buf)
                .map_err(PacketProcessError::VersionNegotiationError)?;

            return Ok(RecvOutcome::Send(len, send_info));
        }

        // Condition: Initial && Version matches

        // SAFETY: we are dealing with an initial packet which has the token set.
        let token = hdr.token.expect("token is always set");

        let dcid = &hdr.dcid;
        let mut cid_builder = self.cid_generator.build_cid();
        cid_builder
            .set_isd_asn(from.isd_asn().to_u64())
            .set_ip_addr(remote_addr.ip())
            .set_port(remote_addr.port());
        // Do stateless retry if the client didn't send a token.
        if token.is_empty() {
            tracing::trace!("Doing stateless retry");
            #[allow(clippy::absurd_extreme_comparisons)]
            if hdr.dcid.len() < MIN_ODCID_LEN {
                return Err(PacketProcessError::DcidTooShort(hdr.dcid.len()));
            }
            cid_builder.set_odcid(dcid);
            let cid = cid_builder.build();
            let new_token_len = self
                .token_generator
                .generate(dcid, &mut self.token_scrub_space);

            // As all the input (e.g. version) is checked, retry() should never
            // fail.
            let len = squiche::retry(
                &hdr.scid,
                dcid,
                &cid, /* the new dcid assigned by the server */
                &self.token_scrub_space[..new_token_len],
                hdr.version,
                send_buf,
            )?;

            return Ok(RecvOutcome::Send(len, send_info));
        }

        // condition: We have a token
        // We can already check whether the DCID of the packet has the
        // expected length, before we recompute the expected SCID itself.
        if dcid.len() != squiche::MAX_CONN_ID_LEN {
            return Err(PacketProcessError::InvalidDestinationConnectionId);
        }

        // extract and verify odcid
        let odcid_len = self
            .token_generator
            .verifier()
            .verify_and_extract_odcid(&token, &mut self.token_scrub_space)?;
        let odcid = &self.token_scrub_space[..odcid_len];

        // recompute the scid
        cid_builder.set_odcid(odcid);
        let expected_scid = cid_builder.build();

        // The DCID chosen by the client MUST be the one that we re-computed
        // from the source IP, port and the original ODCID that the client
        // committed to.
        if hdr.dcid != expected_scid {
            return Err(PacketProcessError::InvalidDestinationConnectionId);
        }

        let odcid = ConnectionId::from_ref(odcid);
        let mut conn = squiche::accept(
            &expected_scid,
            Some(&odcid),
            local_addr,
            remote_addr,
            &mut self.config,
        )
        .map_err(PacketProcessError::AcceptError)?;

        conn.recv(
            recv_buf,
            squiche::RecvInfo {
                from: remote_addr,
                to: local_addr,
            },
        )?;
        let res = self.establishing_set.insert(
            dcid,
            Box::new(QuicScionConn {
                asn_pair: IsdAsnPair {
                    from: self.local_addr.isd_asn(),
                    to: from.isd_asn(),
                },
                inner: conn,
                app: NoApp,
            }),
        );
        self.metrics
            .establishing_connections_gauge
            .set(self.establishing_set.conn_map.len() as i64);
        match res {
            EstablishingOutcome::Established(scion_quic_conn) => {
                Ok(RecvOutcome::EstablishedConn(scion_quic_conn))
            }
            EstablishingOutcome::Establishing(cid) => {
                // SAFETY: just inserted
                Ok(RecvOutcome::Establishing(EstablishingScionQuicConn(
                    self.establishing_set.conn_map.get_mut(&cid).unwrap(),
                )))
            }
            EstablishingOutcome::Done => Ok(RecvOutcome::Done),
        }
    }

    /// Returns the instant at which the next establishing-connection timeout
    /// fires, if any.
    pub fn timeout(&self) -> Option<Instant> {
        self.establishing_set.timeout()
    }

    /// Should be called whenever a timeout happened.
    ///
    /// This method should be called repeatedly until
    /// `Ok(TimeoutOutcome::Done)` is returned.
    pub fn on_timeout(&mut self, now: Instant) -> TimeoutOutcome<'_> {
        let outcome = self.establishing_set.on_timeout(now);
        // `on_timeout` may have closed and removed an establishing connection,
        // so re-sync the gauge here as well. Otherwise the gauge keeps the
        // stale value last written by `recv`, since the timeout path never
        // touches it.
        self.metrics
            .establishing_connections_gauge
            .set(self.establishing_set.conn_map.len() as i64);
        match outcome {
            EstablishingOutcome::Established(scion_quic_conn) => {
                let scid = scion_quic_conn.inner.source_id();
                tracing::error!(
                    ?scid,
                    "establishing connection transitioned to established on timeout"
                );
                TimeoutOutcome::Done
            }
            EstablishingOutcome::Establishing(cid) => {
                // ask the caller to send out packets from this connection
                // SAFETY: just established that the connection exists
                TimeoutOutcome::Establishing(EstablishingScionQuicConn(
                    self.establishing_set.conn_map.get_mut(&cid).unwrap(),
                ))
            }
            EstablishingOutcome::Done => TimeoutOutcome::Done,
        }
    }

    /// Register a handle for an established connection.
    ///
    /// After the [Self::recv()] returned an established connection, the
    /// caller can use this method to register a caller-chosen connection
    /// handle for the connection id `cid`.
    pub fn register_handle(&mut self, cid: ConnectionId<'static>, conn: T) -> Option<T> {
        let res = self.established.insert(cid, conn);
        if res.is_none() {
            self.metrics.routed_source_cids_gauge.inc();
        }
        res
    }

    /// Remove a handle for an established connection.
    pub fn remove_conn(&mut self, cid: ConnectionId<'static>) -> Option<T> {
        let res = self.established.remove(&cid);
        if res.is_some() {
            self.metrics.routed_source_cids_gauge.dec();
        }
        res
    }
}

/// Drives a [`QuicScionServerEndpoint`] on top of a tokio runtime and a
/// [`GenericScionUdpSocket`].
///
/// It is responsible for ...
/// * ... dispatching incoming packets to the endpoint state machine.
/// * ... launching connection drivers for established connections.
/// * ... (de)registering connections with the endpoint.
pub struct QuicScionEndpointDriver<F, A = NoApp>
where
    A: QuicScionApplication,
{
    established_conn: F,
    config: A::Config,
    socket: Arc<dyn GenericScionUdpSocket>,
    local_addr: SocketAddr,
    quic_scion_endpoint: QuicScionServerEndpoint<ConnectionHandle<A>>,
    send_buf: Box<[u8; DEFAULT_SEND_BUF_SIZE]>,
    recv_buf: Box<[u8; DEFAULT_RECV_BUF_SIZE]>,
    spawned_connections: FuturesUnordered<JoinWithId<Result<(), BoxedSocketError>>>,
}

const DEFAULT_SEND_BUF_SIZE: usize = 65535;
const DEFAULT_RECV_BUF_SIZE: usize = 65535;

impl<F, A> QuicScionEndpointDriver<F, A>
where
    A: QuicScionApplication + 'static,
    A::Config: Default,
    F: Fn(ConnectionHandle<A>) + Send + Sync,
{
    /// Creates a new [`QuicScionEndpointDriver`] using a default application
    /// configuration.
    ///
    /// The application protocol `A` run by each connection is determined by the
    /// handle type accepted by `established_conn` (and stored by
    /// `quic_scion_endpoint`); it defaults to [`NoApp`] for plain QUIC. Use
    /// [`Self::with_config`] to provide a non-default configuration.
    ///
    /// ## Parameters
    ///
    /// * `quic_scion_endpoint` the endpoint state machine to drive.
    /// * `scion_udp_socket` the socket used to send and receive packets.
    /// * `established_conn` is invoked once for every newly established connection.
    ///
    /// ## Panics
    ///
    /// This panics if the local address is not an IPv4 or IPv6 address.
    pub fn new(
        quic_scion_endpoint: QuicScionServerEndpoint<ConnectionHandle<A>>,
        scion_udp_socket: Arc<dyn GenericScionUdpSocket>,
        established_conn: F,
    ) -> Self {
        Self::with_config(
            quic_scion_endpoint,
            scion_udp_socket,
            established_conn,
            A::Config::default(),
        )
    }
}

impl<F, A> QuicScionEndpointDriver<F, A>
where
    A: QuicScionApplication + 'static,
    F: Fn(ConnectionHandle<A>) + Send + Sync,
{
    /// Creates a new [`QuicScionEndpointDriver`].
    ///
    /// Each connection's application instance is constructed via
    /// [`QuicScionApplication::on_established`], passing `config` by reference,
    /// once the connection is established.
    ///
    /// ## Parameters
    ///
    /// * `quic_scion_endpoint` the endpoint state machine to drive.
    /// * `scion_udp_socket` the socket used to send and receive packets.
    /// * `established_conn` is invoked once for every newly established connection.
    /// * `config` is the shared application configuration.
    ///
    /// ## Panics
    ///
    /// This panics if the local address is not an IPv4 or IPv6 address.
    pub fn with_config(
        quic_scion_endpoint: QuicScionServerEndpoint<ConnectionHandle<A>>,
        scion_udp_socket: Arc<dyn GenericScionUdpSocket>,
        established_conn: F,
        config: A::Config,
    ) -> Self {
        let local_addr = scion_udp_socket.local_addr().socket_addr();

        Self {
            quic_scion_endpoint,
            established_conn,
            config,
            socket: scion_udp_socket,
            local_addr,
            send_buf: Box::new([0u8; DEFAULT_SEND_BUF_SIZE]),
            recv_buf: Box::new([0u8; DEFAULT_RECV_BUF_SIZE]),
            spawned_connections: Default::default(),
        }
    }

    /// Runs the driver until `cancel_token` is cancelled or a fatal socket
    /// error occurs.
    // XXX(dsd): The current implementation binds the driver to the tokio
    // runtime. However, a generalization only requires delegating task spawning
    // and timeout generation to an abstract runtime.
    pub async fn run(mut self, cancel_token: CancellationToken) -> Result<(), BoxedSocketError> {
        let start_time = Instant::now();
        let timeout = tokio::time::sleep_until(start_time.into());
        tokio::pin!(timeout);
        let mut timeout_inst: Option<Instant> = None;

        while !cancel_token.is_cancelled() {
            tokio::select! {
                /* BEGIN I/O */
                res = self.socket.recv_from(self.recv_buf.as_mut()) => {
                    // We treat I/O errors as fatal here.
                    let (recv_size, recv_from) = res?;
                    self.handle_recv(recv_size, recv_from).await?;
                },
                _ = (&mut timeout), if timeout_inst.is_some() => {
                    self.handle_timeout(timeout_inst.expect("timeout_inst.is_some() is branch condition")).await?;
                },
                r = self.spawned_connections.select_next_some(), if !self.spawned_connections.is_empty() => {
                    self.handle_closed_connection(r);
                },
                /* END I/O */
            }
            let new_timeout = self.quic_scion_endpoint.timeout();
            if timeout_inst != new_timeout {
                timeout_inst = new_timeout;
                if let Some(t) = new_timeout {
                    timeout.as_mut().reset(t.into());
                }
            }
        }
        Ok(())
    }

    #[inline]
    async fn handle_recv(
        &mut self,
        recv_size: usize,
        recv_from: ScionSocketIpAddr,
    ) -> Result<(), BoxedSocketError> {
        match self.quic_scion_endpoint.recv(
            &mut self.recv_buf.as_mut()[..recv_size],
            self.send_buf.as_mut(),
            recv_from,
        ) {
            Ok(RecvOutcome::ConnEvent(c)) => {
                {
                    let mut conn = c.lock();
                    if let Err(err) = conn.inner.recv(
                        &mut self.recv_buf[..recv_size],
                        RecvInfo {
                            from: recv_from.socket_addr(),
                            to: self.local_addr,
                        },
                    ) {
                        let scid = conn.inner.source_id();
                        tracing::error!(?scid, ?err, "error receiving for connection")
                    }
                }
                c.notify();
            }
            Ok(RecvOutcome::Send(n, send_info)) => {
                // XXX(dsd): Here, we bail out on any I/O-error, but ideally
                // want to be more lenient in case of recoverable errors.
                //
                // Currently, however, the GenericScionUdpSocket trait does not
                // expose the underlying error type, so we cannot distinguish
                // between recoverable and fatal errors.
                self.socket
                    .send_to(&self.send_buf[..n], send_info.to)
                    .await?;
            }
            Ok(RecvOutcome::Establishing(mut c)) => {
                loop {
                    match c.send(self.send_buf.as_mut()) {
                        Ok((n, send_info)) => {
                            // XXX(dsd): Here, we bail out on any I/O-error, but
                            // ideally want to be more lenient in case of
                            // recoverable errors.
                            //
                            // Currently, however, the GenericScionUdpSocket
                            // trait does not expose the underlying error type,
                            // so we cannot distinguish between recoverable and
                            // fatal errors.
                            self.socket
                                .send_to(&self.send_buf[..n], send_info.to)
                                .await?;
                        }
                        Err(squiche::Error::Done) => break,
                        Err(err) => {
                            tracing::error!(
                                ?err,
                                "error performing send operation on establishing connection"
                            )
                        }
                    }
                }
            }
            Ok(RecvOutcome::EstablishedConn(c)) => {
                let conn_id = c.inner.source_id().into_owned();
                // The endpoint establishes connections without an application
                // layer; construct the application now that the connection is
                // established and attach it.
                let QuicScionConn {
                    asn_pair,
                    mut inner,
                    ..
                } = *c;
                let app = A::on_established(&mut inner, &self.config);
                let conn = QuicScionConn {
                    asn_pair,
                    inner,
                    app,
                };
                let conn_handle = ConnectionHandle::new(Notify::new(), conn);
                let jh = tokio::spawn({
                    let driver = QuicScionConnDriver::new(conn_handle.clone(), self.socket.clone());
                    async move { driver.run().await }
                });
                self.spawned_connections.push({
                    let conn_id = conn_id.clone();
                    JoinWithId::new(conn_id, jh)
                });
                self.quic_scion_endpoint
                    .register_handle(conn_id, conn_handle.clone());
                (self.established_conn)(conn_handle);
            }
            Ok(RecvOutcome::Done) => {}
            Err(err) => {
                tracing::info!(?err, "error driving endpoint");
            }
        }
        Ok(())
    }

    async fn handle_timeout(&mut self, now: Instant) -> Result<(), BoxedSocketError> {
        while let TimeoutOutcome::Establishing(mut establishing_scion_quic_conn) =
            self.quic_scion_endpoint.on_timeout(now)
        {
            match establishing_scion_quic_conn.send(self.send_buf.as_mut()) {
                Ok((n, send_info)) => {
                    self.socket
                        .send_to(&self.send_buf[..n], send_info.to)
                        .await?;
                }
                Err(squiche::Error::Done) => {}
                Err(err) => {
                    tracing::error!(?err, "sending packet for establishing connection");
                }
            }
        }
        Ok(())
    }

    fn handle_closed_connection(
        &mut self,
        r: (
            ConnectionId<'static>,
            Result<Result<(), BoxedSocketError>, JoinError>,
        ),
    ) {
        let (conn_id, res): (_, Result<_, JoinError>) = r;
        match res {
            Ok(Err(err)) => {
                tracing::error!(?err, "Connection driver returned socket error");
            }
            Err(err) => {
                tracing::error!(?err, "Connection driver failed with exception");
            }
            Ok(_) => {}
        }
        self.quic_scion_endpoint.remove_conn(conn_id);
    }
}

struct JoinWithId<R> {
    conn_id: ConnectionId<'static>,
    handle: JoinHandle<R>,
}

impl<R> JoinWithId<R> {
    fn new(conn_id: ConnectionId<'static>, handle: JoinHandle<R>) -> Self {
        Self { conn_id, handle }
    }
}

impl<R> Future for JoinWithId<R> {
    // Always yields conn_id, even on JoinError (cancel/panic).
    type Output = (ConnectionId<'static>, Result<R, JoinError>);

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // We never move `handle` out, only poll it in place.
        // `conn_id` is Clone (or we clone it on completion).
        let this = self.get_mut();
        match Pin::new(&mut this.handle).poll(cx) {
            Poll::Ready(res) => Poll::Ready((this.conn_id.clone(), res)),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// The result of a receive operation.
pub type RecvResult<'a, T> = Result<RecvOutcome<'a, T>, PacketProcessError>;

/// The outcome of a receive operation, if successful.
#[non_exhaustive]
pub enum RecvOutcome<'a, T> {
    /// A connection that has been fully established.
    EstablishedConn(Box<QuicScionConn>),
    /// An incoming packet triggered a response from a connection that is
    /// currently being established.
    /// It is the callers responsibility to send all queued packets on the
    /// contained connection, by repeatedly calling `send()` until
    /// `Err(Done)` is returned.
    Establishing(EstablishingScionQuicConn<'a>),
    /// This is only returned when the connection is not yet established.
    Send(usize, ScionSendInfo),
    /// A packet has arrived for the given connection.
    ConnEvent(&'a mut T),
    /// Nothing to be done.
    Done,
}

/// The result of calling [`QuicScionServerEndpoint::on_timeout`].
pub enum TimeoutOutcome<'a> {
    /// A timeout of a connection that is currently being established was
    /// triggered.
    /// It is the callers responsibility to send all queued packets on the
    /// contained connection, by repeatedly calling `send()` until
    /// `Err(Done)` is returned.
    Establishing(EstablishingScionQuicConn<'a>),
    /// No further action is required.
    Done,
}

/// A SCION QUIC connection that is currently being established.
///
/// Handed out by the endpoint so the caller can flush the connection's queued
/// outgoing packets while it is still completing its handshake.
pub struct EstablishingScionQuicConn<'a>(&'a mut QuicScionConn);

impl<'a> EstablishingScionQuicConn<'a> {
    /// Writes the next queued outgoing packet for this establishing connection
    /// into `send_buf`, returning its length and the SCION send info.
    pub fn send(&mut self, send_buf: &mut [u8]) -> squiche::Result<(usize, ScionSendInfo)> {
        self.0.send(send_buf)
    }
}

/// Metrics for the connection endpoint.
#[derive(Debug, Clone)]
pub struct Metrics {
    /// The number of connections that are currently being established.
    pub establishing_connections_gauge: IntGauge,

    /// The number of currently registered connections.
    pub routed_source_cids_gauge: IntGauge,
}

impl Metrics {
    /// Creates a new [`Metrics`] instance disconnected from any Prometheus registry.
    pub fn new_without_registry() -> Self {
        Self {
            establishing_connections_gauge: IntGauge::new(
                "quic_scion_establishing_connections",
                "The number of connections that are currently being established.",
            )
            .unwrap(),
            routed_source_cids_gauge: IntGauge::new(
                "quic_scion_routed_source_cids",
                "The number of currently registered connections.",
            )
            .unwrap(),
        }
    }
}

/// An error that occurred while processing an incoming packet.
#[derive(Debug, Error)]
pub enum PacketProcessError {
    /// Could not parse the local/remote address.
    #[error("failed to parse local/remote address")]
    InvalidAddress,
    /// An initial packet was expected (for the contained DCID), but a packet of
    /// another type was received.
    #[error("expected initial packet: {0:?}")]
    ExpectedInitialPacket(ConnectionId<'static>),
    /// The address validation token was invalid.
    #[error("invalid address validation token: {0}")]
    InvalidToken(#[from] TokenVerificationError),
    /// The DCID chosen by the client was invalid.
    #[error("invalid destination connection ID")]
    InvalidDestinationConnectionId,
    /// The QUIC header could not be parsed.
    #[error("invalid header: {0}")]
    InvalidHeader(squiche::Error),
    /// An error occurred during connection establishment.
    #[error("connection error during establishment: {0}")]
    ConnectionError(#[from] squiche::Error),
    /// Version negotiation failed.
    #[error("failed to negotiate version: {0}")]
    VersionNegotiationError(squiche::Error),
    /// Accepting the connection failed.
    #[error("failed to accept connection: {0}")]
    AcceptError(squiche::Error),
    /// The client-chosen DCID is shorter than the configured minimum.
    #[error("client-chosen dcid too short: {0}")]
    DcidTooShort(usize),
}

/// An error that occurred while verifying an address validation token.
#[derive(Debug, Error)]
pub enum TokenVerificationError {
    /// Token is larger than the maximum allowed size.
    /// (Maximum CID length + Tag length)
    #[error("Token too long")]
    TokenTooLong,
    /// Token is too short.
    /// (Minimum CID length + Tag length)
    #[error("Token too short")]
    TokenTooShort,
    /// Signature verification failed.
    #[error("Signature verification failed")]
    InvalidSignature,
}

// Represents the set of connections that are currently being established.
#[derive(Default)]
struct EstablishingSet {
    // Invariant: conn.is_establishing(), ∀ conn ∊ being_established
    conn_map: HashMap<ConnectionId<'static>, Box<QuicScionConn>>,
    // A priority queue of connections that are currently being established.
    //
    // Let `(t, c)` be the item with lowest priority (i.e. next timeout) in
    // `being_established`. Calling `on_timeout(now)` with `now >= t` will
    // call `on_timeout` on the connection with connection id `c`. The
    // postcondition of `on_timeout(now)` is that the connection with
    // connection id `c` is either `!is_establishing` or it is re-inserted
    // with a timeout instant `> now`.
    timeouts: BTreeSet<(Instant, OrdId<'static>)>,
}

impl EstablishingSet {
    fn on_timeout(&mut self, now: Instant) -> EstablishingOutcome {
        let (_, cid) = if self.timeouts.first().is_some_and(|(t, _)| *t <= now) {
            self.timeouts.pop_first().unwrap()
        } else {
            return EstablishingOutcome::Done;
        };

        let Entry::Occupied(mut conn) = self.conn_map.entry(cid.0.clone()) else {
            return EstablishingOutcome::Done;
        };
        conn.get_mut().inner.on_timeout();
        // Assumption: a connection _cannot_ transition to established based
        // on a timeout.
        if conn.get_mut().inner.is_closed() {
            tracing::trace!(cid=?cid.0, "establishing connection closed");
            let _ = conn.remove();
            return EstablishingOutcome::Done;
        }
        // Re-insert the connection with the current timeout.
        if let Some(t) = conn.get().inner.timeout_instant() {
            self.timeouts.insert((t, cid.clone()));
        }

        EstablishingOutcome::Establishing(cid.0)
    }

    fn timeout(&self) -> Option<Instant> {
        self.timeouts.first().map(|(t, _)| *t)
    }

    fn update<F>(
        &mut self,
        dcid: &ConnectionId<'static>,
        mut update: F,
    ) -> Result<EstablishingOutcome, squiche::Error>
    where
        F: FnMut(&mut QuicScionConn) -> Result<(), squiche::Error>,
    {
        let Entry::Occupied(mut conn) = self.conn_map.entry(dcid.clone()) else {
            return Ok(EstablishingOutcome::Done);
        };
        update(conn.get_mut())?;
        Ok(Self::process_conn(&mut self.timeouts, dcid, conn))
    }

    fn insert(
        &mut self,
        dcid: &ConnectionId<'static>,
        conn: Box<QuicScionConn>,
    ) -> EstablishingOutcome {
        let conn = self.conn_map.entry(dcid.clone()).insert_entry(conn);
        Self::process_conn(&mut self.timeouts, dcid, conn)
    }

    fn process_conn(
        timeouts: &mut BTreeSet<(Instant, OrdId<'static>)>,
        dcid: &ConnectionId<'static>,
        conn: OccupiedEntry<'_, ConnectionId<'static>, Box<QuicScionConn>>,
    ) -> EstablishingOutcome {
        if conn.get().inner.is_established() {
            return EstablishingOutcome::Established(conn.remove());
        }
        if let Some(t) = conn.get().inner.timeout_instant() {
            timeouts.insert((t, OrdId(dcid.clone())));
        }
        EstablishingOutcome::Establishing(dcid.clone())
    }
}

/// The outcome of feeding a packet/timeout to an establishing connection.
enum EstablishingOutcome {
    Established(Box<QuicScionConn>),
    Establishing(ConnectionId<'static>),
    Done,
}

/// A [`ConnectionId`] wrapper that is ordered by its raw bytes, so it can be
/// used in the timeout priority queue.
#[derive(Debug, Clone, PartialEq, Eq)]
struct OrdId<'a>(ConnectionId<'a>);

impl<'a> std::cmp::PartialOrd for OrdId<'a> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<'a> std::cmp::Ord for OrdId<'a> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.as_ref().cmp(other.0.as_ref())
    }
}

/// Constructs the server-chosen connection IDs.
struct CidGenerator {
    signing_key: HmacKey,
}

impl CidGenerator {
    fn new(key_seed: [u8; 32]) -> Self {
        let signing_key = HmacKey::new(HMAC_ALGO, &key_seed);
        Self { signing_key }
    }

    fn build_cid(&self) -> CidBuilder<'_> {
        CidBuilder {
            signing_key: &self.signing_key,
            message: [0u8; CidBuilder::TOTAL_LEN],
        }
    }
}

/// Builds a server-chosen connection ID that commits to the remote's SCION
/// address (ISD-ASN, IP, port) and the original destination connection ID.
struct CidBuilder<'a> {
    signing_key: &'a HmacKey,
    message: [u8; CidBuilder::TOTAL_LEN],
}

impl<'a> CidBuilder<'a> {
    const TOTAL_LEN: usize = Self::DCID_OFFSET_RANGE.end;
    const ASN_OFFSET_RANGE: Range<usize> = //
        0..size_of::<u64>();
    const IP_ADDR_OFFSET_RANGE: Range<usize> =
        Self::ASN_OFFSET_RANGE.end..(Self::ASN_OFFSET_RANGE.end + 16usize);
    const PORT_OFFSET_RANGE: Range<usize> =
        Self::IP_ADDR_OFFSET_RANGE.end..(Self::IP_ADDR_OFFSET_RANGE.end + size_of::<u16>());
    const DCID_OFFSET_RANGE: Range<usize> =
        Self::PORT_OFFSET_RANGE.end..(Self::PORT_OFFSET_RANGE.end + squiche::MAX_CONN_ID_LEN);

    fn set_isd_asn(&mut self, asn: u64) -> &mut Self {
        self.message[Self::ASN_OFFSET_RANGE].copy_from_slice(&asn.to_be_bytes());
        self
    }

    fn set_ip_addr(&mut self, ip_addr: IpAddr) -> &mut Self {
        Self::write_addr_bytes(ip_addr, &mut self.message[Self::IP_ADDR_OFFSET_RANGE]);
        self
    }

    fn set_port(&mut self, port: u16) -> &mut Self {
        self.message[Self::PORT_OFFSET_RANGE].copy_from_slice(&port.to_be_bytes());
        self
    }

    fn set_odcid(&mut self, dcid: &[u8]) -> &mut Self {
        // Connection IDs are variable length (0..=MAX_CONN_ID_LEN), so the
        // ODCID is not guaranteed to fill the whole region (e.g. stock quiche
        // clients use 16-byte CIDs). Write it into the start of the fixed-size
        // region and leave the remaining bytes zero (the message buffer is
        // zero-initialised in `build_cid`). The retry and verification paths
        // write the same ODCID, so the derived CID stays stable.
        debug_assert!(
            dcid.len() <= squiche::MAX_CONN_ID_LEN,
            "odcid longer than MAX_CONN_ID_LEN"
        );
        let dcid_region = &mut self.message[Self::DCID_OFFSET_RANGE];
        dcid_region[..dcid.len()].copy_from_slice(dcid);
        self
    }

    fn build(self) -> squiche::ConnectionId<'static> {
        let sig = ring::hmac::sign(self.signing_key, &self.message);
        sig.as_ref()[..squiche::MAX_CONN_ID_LEN].to_vec().into()
    }

    fn write_addr_bytes(ip_addr: IpAddr, buf: &mut [u8]) {
        assert!(buf.len() >= 16, "buffer must be at least 16 bytes");

        let ipv6 = match ip_addr {
            IpAddr::V4(v4) => v4.to_ipv6_mapped(),
            IpAddr::V6(v6) => v6,
        };

        buf[..16].copy_from_slice(&ipv6.octets());
    }
}

/// Generates address validation (retry) tokens.
struct TokenGenerator {
    signing_key: HmacKey,
}

impl TokenGenerator {
    fn new(key: [u8; 32]) -> Self {
        let signing_key = HmacKey::new(HMAC_ALGO, &key);
        Self { signing_key }
    }

    /// Generate a token for the odcid and write the token to `out` and
    /// return the length of the token.
    ///
    /// Format: TAG || odcid
    fn generate(&self, odcid: &[u8], out: &mut [u8]) -> usize {
        let end = DIGEST_LEN + odcid.len();
        out[..DIGEST_LEN].copy_from_slice(ring::hmac::sign(&self.signing_key, odcid).as_ref());
        out[DIGEST_LEN..end].copy_from_slice(odcid);
        DIGEST_LEN + odcid.len()
    }

    fn verifier(&self) -> TokenVerifier<'_> {
        let signing_key = &self.signing_key;
        TokenVerifier { signing_key }
    }
}

/// Verifies address validation (retry) tokens and extracts the committed ODCID.
struct TokenVerifier<'a> {
    signing_key: &'a HmacKey,
}

impl<'a> TokenVerifier<'a> {
    fn verify_and_extract_odcid(
        &self,
        token: &[u8],
        out: &mut [u8],
    ) -> Result<usize, TokenVerificationError> {
        if token.len() < MIN_TOKEN_LEN {
            return Err(TokenVerificationError::TokenTooShort);
        }

        if token.len() > MAX_TOKEN_LEN {
            return Err(TokenVerificationError::TokenTooLong);
        }

        let (tag, odcid) = token.split_at(DIGEST_LEN);

        ring::hmac::verify(self.signing_key, odcid, tag)
            .map_err(|_| TokenVerificationError::InvalidSignature)?;

        out[..odcid.len()].copy_from_slice(odcid);
        Ok(odcid.len())
    }
}

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

    fn make_generator() -> TokenGenerator {
        TokenGenerator::new([0x42u8; 32])
    }

    #[test]
    fn roundtrip() {
        let g = make_generator();
        let odcid = b"test-cid-1234567";
        let mut token_buf = [0u8; MAX_TOKEN_LEN];
        let mut odcid_out = [0u8; MAX_TOKEN_LEN];

        let token_len = g.generate(odcid, &mut token_buf);
        let odcid_len = g
            .verifier()
            .verify_and_extract_odcid(&token_buf[..token_len], &mut odcid_out)
            .unwrap();

        assert_eq!(&odcid_out[..odcid_len], odcid);
    }

    #[test]
    fn tampered_tag_rejected() {
        let g = make_generator();
        let mut token_buf = [0u8; MAX_TOKEN_LEN];
        let token_len = g.generate(b"some-cid", &mut token_buf);

        // Flip a bit in the tag
        token_buf[0] ^= 0x01;

        let mut odcid_out = [0u8; MAX_TOKEN_LEN];
        assert!(matches!(
            g.verifier()
                .verify_and_extract_odcid(&token_buf[..token_len], &mut odcid_out),
            Err(TokenVerificationError::InvalidSignature)
        ));
    }

    #[test]
    fn tampered_odcid_rejected() {
        let g = make_generator();
        let mut token_buf = [0u8; MAX_TOKEN_LEN];
        let token_len = g.generate(b"some-cid", &mut token_buf);

        // Flip a bit in the odcid payload
        token_buf[DIGEST_LEN] ^= 0x01;

        let mut odcid_out = [0u8; MAX_TOKEN_LEN];
        assert!(matches!(
            g.verifier()
                .verify_and_extract_odcid(&token_buf[..token_len], &mut odcid_out),
            Err(TokenVerificationError::InvalidSignature)
        ));
    }

    #[test]
    fn token_too_short_rejected() {
        let g = make_generator();
        let mut odcid_out = [0u8; MAX_TOKEN_LEN];
        // Only a partial tag, no odcid
        let short_token = [0u8; DIGEST_LEN - 1];
        assert!(matches!(
            g.verifier()
                .verify_and_extract_odcid(&short_token, &mut odcid_out),
            Err(TokenVerificationError::TokenTooShort)
        ));
    }

    #[test]
    fn token_too_long_rejected() {
        let g = make_generator();
        let mut odcid_out = [0u8; MAX_TOKEN_LEN];
        let long_token = [0u8; MAX_TOKEN_LEN + 1];
        assert!(matches!(
            g.verifier()
                .verify_and_extract_odcid(&long_token, &mut odcid_out),
            Err(TokenVerificationError::TokenTooLong)
        ));
    }

    #[test]
    fn wrong_key_rejected() {
        let g = make_generator();
        let mut token_buf = [0u8; MAX_TOKEN_LEN];
        let token_len = g.generate(b"some-cid", &mut token_buf);

        // Verify with a different key
        let other_gen = TokenGenerator::new([0x99u8; 32]);
        let mut odcid_out = [0u8; MAX_TOKEN_LEN];
        assert!(matches!(
            other_gen
                .verifier()
                .verify_and_extract_odcid(&token_buf[..token_len], &mut odcid_out),
            Err(TokenVerificationError::InvalidSignature)
        ));
    }
}