asupersync 0.3.5

//! Connection routing and management for native QUIC endpoint.
//!
//! This module provides connection-ID routing, timer scheduling integration,
//! and connection lifecycle management for the ATP native QUIC endpoint.
//! It bridges the gap between the UDP endpoint packet I/O and individual
//! QUIC connection state machines.

#![allow(dead_code)]

use crate::cx::Cx;
use crate::net::atp::quic::AtpPacketProtection;
use crate::net::quic_core::{
    ConnectionId, LongPacketType, PacketHeader, QuicCoreError, ShortHeader,
};
use crate::net::quic_native::{
    NativeQuicConnection, NativeQuicConnectionConfig, OutgoingPacket, ReceivedPacket,
};
use crate::net::quic_native::{
    NativeQuicConnectionError, PacketNumberSpace, PacketProtectionRequest, PacketProtectionSpace,
    ProtectedPacket, ProtectionProof, TranscriptHash,
};
use crate::time::Sleep;
use crate::types::outcome::Outcome;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::time::{Duration, Instant};

const DEFAULT_MAX_CONNECTIONS: usize = 4096;

/// Connection routing table that maps connection IDs to active QUIC connections.
#[derive(Debug)]
pub struct ConnectionRouter {
    /// Map from destination connection ID to connection handle.
    connections: HashMap<ConnectionId, ConnectionHandle>,
    /// Maximum active connections accepted by this router.
    max_connections: usize,
    /// Next connection ID counter for generating new connections.
    next_connection_id: u64,
    /// Connection configuration template.
    config_template: NativeQuicConnectionConfig,
    /// Monotonic clock origin for connection timer APIs that use microseconds.
    clock_origin: Instant,
}

/// Handle to a managed QUIC connection with timing and lifecycle state.
#[derive(Debug)]
pub struct ConnectionHandle {
    /// The underlying QUIC connection state machine.
    connection: NativeQuicConnection,
    /// Packet-protection provider for 1-RTT UDP handoff.
    packet_protection: Option<ConnectionPacketProtection>,
    /// Remote peer address.
    peer_addr: SocketAddr,
    /// Last activity timestamp for connection timeout tracking.
    last_activity: Instant,
    /// Connection establishment timestamp.
    established_at: Option<Instant>,
    /// Pending timer deadline for this connection.
    next_timer_deadline: Option<Instant>,
}

/// Native QUIC connection removed from the router for application-level handoff.
#[derive(Debug)]
pub struct AcceptedNativeQuicConnection {
    /// Connection ID that owned the routed connection.
    pub connection_id: ConnectionId,
    /// The native QUIC connection state machine.
    pub connection: NativeQuicConnection,
    /// Remote peer address associated with the accepted connection.
    pub peer_addr: SocketAddr,
}

struct ConnectionPacketProtection {
    protection: AtpPacketProtection,
}

impl std::fmt::Debug for ConnectionPacketProtection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConnectionPacketProtection")
            .field("provider_kind", &self.protection.provider_kind())
            .finish_non_exhaustive()
    }
}

/// Timer event for a specific connection.
#[derive(Debug, Clone)]
pub struct ConnectionTimerEvent {
    /// Connection ID this timer event belongs to.
    pub connection_id: ConnectionId,
    /// Type of timer that fired.
    pub timer_type: TimerType,
    /// Deadline when this timer was scheduled to fire.
    pub deadline: Instant,
}

/// Types of timers used by QUIC connections.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimerType {
    /// Probe timeout (PTO) for loss recovery.
    ProbeTimeout,
    /// ACK delay timer.
    AckDelay,
    /// Connection idle timeout.
    IdleTimeout,
    /// Connection draining timeout.
    DrainTimeout,
    /// Keep-alive probe.
    KeepAlive,
}

/// Result of routing a received packet to a connection.
#[derive(Debug)]
pub enum RoutingResult {
    /// Packet was successfully routed to an existing connection.
    Routed {
        /// Connection ID packet was routed to.
        connection_id: ConnectionId,
        /// Outgoing packets generated by processing this packet.
        outgoing_packets: Vec<OutgoingPacket>,
    },
    /// Packet is a new connection attempt (e.g., Initial packet).
    NewConnection {
        /// Suggested connection ID for the new connection.
        connection_id: ConnectionId,
        /// Remote address that originated the first datagram.
        peer_addr: SocketAddr,
        /// Original Initial packet that triggered connection creation.
        triggering_packet: ReceivedPacket,
        /// Initial outgoing packets for handshake response.
        outgoing_packets: Vec<OutgoingPacket>,
    },
    /// Packet should be dropped (invalid CID, stateless reset, etc.).
    Drop {
        /// Reason for dropping the packet.
        reason: String,
    },
}

/// Errors from connection routing operations.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionRouterError {
    /// Operation was cancelled via Cx.
    Cancelled,
    /// Connection ID not found in routing table.
    ConnectionNotFound(ConnectionId),
    /// Connection is in invalid state for the operation.
    InvalidConnectionState {
        /// Connection ID.
        connection_id: ConnectionId,
        /// Description of invalid state.
        reason: String,
    },
    /// Unable to create new connection.
    ConnectionCreationFailed(String),
    /// Timer scheduling failed.
    TimerSchedulingFailed(String),
    /// Packet reached a connection but failed state-machine processing.
    PacketProcessingFailed {
        /// Connection ID.
        connection_id: ConnectionId,
        /// Processing error.
        reason: String,
    },
    /// Application-data frames were ready, but no 1-RTT packet protection was installed.
    PacketProtectionUnavailable {
        /// Connection ID.
        connection_id: ConnectionId,
    },
}

impl std::fmt::Display for ConnectionRouterError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Cancelled => write!(f, "operation cancelled"),
            Self::ConnectionNotFound(cid) => write!(f, "connection not found: {cid:?}"),
            Self::InvalidConnectionState {
                connection_id,
                reason,
            } => {
                write!(
                    f,
                    "invalid connection state for {connection_id:?}: {reason}"
                )
            }
            Self::ConnectionCreationFailed(msg) => write!(f, "connection creation failed: {msg}"),
            Self::TimerSchedulingFailed(msg) => write!(f, "timer scheduling failed: {msg}"),
            Self::PacketProcessingFailed {
                connection_id,
                reason,
            } => {
                write!(
                    f,
                    "packet processing failed for {connection_id:?}: {reason}"
                )
            }
            Self::PacketProtectionUnavailable { connection_id } => {
                write!(
                    f,
                    "packet protection unavailable for application-data packet on {connection_id:?}"
                )
            }
        }
    }
}

impl std::error::Error for ConnectionRouterError {}

impl ConnectionRouter {
    /// Create a new connection router with the given configuration template.
    pub fn new(config_template: NativeQuicConnectionConfig) -> Self {
        Self::with_max_connections(config_template, DEFAULT_MAX_CONNECTIONS)
    }

    /// Create a connection router with an explicit active-connection cap.
    ///
    /// A zero cap is normalized to one connection so the router never accepts an
    /// unbounded configuration by accident.
    pub fn with_max_connections(
        config_template: NativeQuicConnectionConfig,
        max_connections: usize,
    ) -> Self {
        Self {
            connections: HashMap::new(),
            max_connections: max_connections.max(1),
            next_connection_id: 1,
            config_template,
            clock_origin: Instant::now(),
        }
    }

    /// Route a received packet to the appropriate connection.
    pub async fn route_packet(
        &mut self,
        cx: &Cx,
        packet: ReceivedPacket,
    ) -> Result<RoutingResult, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let routing_info = match self.decode_routing_info(&packet) {
            Ok(info) => info,
            Err(err) => {
                return Ok(RoutingResult::Drop {
                    reason: format!("invalid QUIC header: {err}"),
                });
            }
        };
        let connection_id = routing_info.destination_cid;
        let now_micros = self.instant_micros(packet.receive_time);

        if let Some(handle) = self.connections.get_mut(&connection_id) {
            handle.last_activity = Instant::now();
            handle
                .connection
                .on_datagram_received(cx, packet.data.len() as u64)
                .map_err(|err| ConnectionRouterError::PacketProcessingFailed {
                    connection_id,
                    reason: err.to_string(),
                })?;
            let payload = packet.data.get(routing_info.header_len..).ok_or_else(|| {
                ConnectionRouterError::PacketProcessingFailed {
                    connection_id,
                    reason: "header length exceeded datagram length".to_string(),
                }
            })?;
            let plaintext_payload = if routing_info.space == PacketNumberSpace::ApplicationData {
                unprotect_1rtt_packet(
                    cx,
                    connection_id,
                    handle,
                    &packet.data[..routing_info.header_len],
                    payload,
                    routing_info.packet_number,
                    routing_info.key_phase,
                )
                .await?
            } else {
                payload.to_vec()
            };
            handle
                .connection
                .process_packet_payload(
                    cx,
                    routing_info.space,
                    routing_info.packet_number,
                    &plaintext_payload,
                    now_micros,
                )
                .map_err(|err| ConnectionRouterError::PacketProcessingFailed {
                    connection_id,
                    reason: err.to_string(),
                })?;
            let outgoing_packets = drain_connection_frames(
                cx,
                connection_id,
                handle,
                routing_info.space,
                packet.src_addr,
                packet.receive_time,
                now_micros,
            )
            .await?;
            Self::refresh_connection_timer(
                cx,
                connection_id,
                handle,
                self.clock_origin,
                now_micros,
                packet.receive_time,
            )?;
            cx.trace(&format!(
                "Routed packet from {} to connection {connection_id:?}",
                packet.src_addr
            ));

            Ok(RoutingResult::Routed {
                connection_id,
                outgoing_packets,
            })
        } else if routing_info.kind == PacketRoutingKind::Initial
            && self.connections.len() >= self.max_connections
        {
            Ok(RoutingResult::Drop {
                reason: format!(
                    "connection limit reached: active={}, max={}",
                    self.connections.len(),
                    self.max_connections
                ),
            })
        } else if routing_info.kind == PacketRoutingKind::Initial {
            let new_connection_id = connection_id;

            cx.trace(&format!(
                "New connection attempt from {} assigned ID {new_connection_id:?}",
                packet.src_addr
            ));

            Ok(RoutingResult::NewConnection {
                connection_id: new_connection_id,
                peer_addr: packet.src_addr,
                triggering_packet: packet,
                outgoing_packets: Vec::new(),
            })
        } else {
            Ok(RoutingResult::Drop {
                reason: format!(
                    "unknown connection ID {connection_id:?} for {:?} packet",
                    routing_info.kind
                ),
            })
        }
    }

    /// Create a new connection and add it to the routing table.
    pub async fn create_connection(
        &mut self,
        cx: &Cx,
        connection_id: ConnectionId,
        peer_addr: SocketAddr,
        is_server: bool,
    ) -> Result<(), ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        if self.connections.contains_key(&connection_id) {
            return Err(ConnectionRouterError::ConnectionCreationFailed(format!(
                "connection ID collision: {connection_id:?}"
            )));
        }
        if self.connections.len() >= self.max_connections {
            return Err(ConnectionRouterError::ConnectionCreationFailed(format!(
                "connection limit reached: active={}, max={}",
                self.connections.len(),
                self.max_connections
            )));
        }

        let mut config = self.config_template;
        config.role = if is_server {
            crate::net::quic_native::StreamRole::Server
        } else {
            crate::net::quic_native::StreamRole::Client
        };

        // Create the QUIC connection state machine
        let connection = NativeQuicConnection::new(config);

        let handle = ConnectionHandle {
            connection,
            packet_protection: None,
            peer_addr,
            last_activity: Instant::now(),
            established_at: None,
            next_timer_deadline: None,
        };

        self.connections.insert(connection_id, handle);

        cx.trace(&format!(
            "Created new connection {connection_id:?} for peer {peer_addr}"
        ));

        Ok(())
    }

    /// Install packet protection for a managed connection's 1-RTT
    /// UDP handoff path.
    ///
    /// Once installed, application-data frames drained by this router are
    /// assembled as short-header packets and protected before being passed to
    /// the UDP endpoint.
    pub fn install_packet_protection(
        &mut self,
        cx: &Cx,
        connection_id: ConnectionId,
        protection: AtpPacketProtection,
    ) -> Result<(), ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let handle = self
            .connections
            .get_mut(&connection_id)
            .ok_or(ConnectionRouterError::ConnectionNotFound(connection_id))?;
        handle.packet_protection = Some(ConnectionPacketProtection { protection });
        Ok(())
    }

    /// Borrow a managed connection for focused integration tests.
    #[cfg(any(test, feature = "test-internals"))]
    pub fn connection_mut_for_testing(
        &mut self,
        cx: &Cx,
        connection_id: ConnectionId,
    ) -> Result<&mut NativeQuicConnection, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        self.connections
            .get_mut(&connection_id)
            .map(|handle| &mut handle.connection)
            .ok_or(ConnectionRouterError::ConnectionNotFound(connection_id))
    }

    /// Drain one batch of protected application-data packets for focused
    /// integration tests that need to cross the real UDP endpoint boundary.
    #[cfg(any(test, feature = "test-internals"))]
    pub async fn drain_application_data_for_testing(
        &mut self,
        cx: &Cx,
        connection_id: ConnectionId,
        dst_addr: SocketAddr,
        now: Instant,
    ) -> Result<Vec<OutgoingPacket>, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let now_micros = self.instant_micros(now);
        let handle = self
            .connections
            .get_mut(&connection_id)
            .ok_or(ConnectionRouterError::ConnectionNotFound(connection_id))?;
        drain_connection_frames(
            cx,
            connection_id,
            handle,
            PacketNumberSpace::ApplicationData,
            dst_addr,
            now,
            now_micros,
        )
        .await
    }

    /// Remove a connection from the routing table.
    pub fn remove_connection(
        &mut self,
        cx: &Cx,
        connection_id: ConnectionId,
    ) -> Result<(), ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        if self.connections.remove(&connection_id).is_some() {
            cx.trace(&format!("Removed connection {connection_id:?}"));
            Ok(())
        } else {
            Err(ConnectionRouterError::ConnectionNotFound(connection_id))
        }
    }

    /// Remove a connection from the router and hand ownership to the caller.
    pub fn take_connection(
        &mut self,
        cx: &Cx,
        connection_id: ConnectionId,
    ) -> Result<AcceptedNativeQuicConnection, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let handle = self
            .connections
            .remove(&connection_id)
            .ok_or(ConnectionRouterError::ConnectionNotFound(connection_id))?;

        cx.trace(&format!(
            "Accepted native QUIC connection {connection_id:?}"
        ));
        Ok(AcceptedNativeQuicConnection {
            connection_id,
            connection: handle.connection,
            peer_addr: handle.peer_addr,
        })
    }

    /// Remove the next routed connection using deterministic connection-ID order.
    pub fn take_next_connection(
        &mut self,
        cx: &Cx,
    ) -> Result<Option<AcceptedNativeQuicConnection>, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let Some(connection_id) = self
            .connections
            .keys()
            .min_by(|left, right| left.as_bytes().cmp(right.as_bytes()))
            .copied()
        else {
            return Ok(None);
        };

        self.take_connection(cx, connection_id).map(Some)
    }

    /// Close and remove every active connection.
    pub fn close_all(
        &mut self,
        cx: &Cx,
        now: Instant,
        app_error_code: u64,
    ) -> Result<usize, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let now_micros = self.instant_micros(now);
        for (connection_id, handle) in &mut self.connections {
            handle
                .connection
                .begin_close(cx, now_micros, app_error_code)
                .or_else(|_| handle.connection.close_immediately(cx, app_error_code))
                .map_err(|err| ConnectionRouterError::PacketProcessingFailed {
                    connection_id: *connection_id,
                    reason: err.to_string(),
                })?;
        }
        let closed = self.connections.len();
        self.connections.clear();
        Ok(closed)
    }

    /// Refresh a connection's PTO deadline from its transport state.
    fn refresh_connection_timer(
        cx: &Cx,
        connection_id: ConnectionId,
        handle: &mut ConnectionHandle,
        origin: Instant,
        now_micros: u64,
        now_instant: Instant,
    ) -> Result<(), ConnectionRouterError> {
        handle.next_timer_deadline = handle
            .connection
            .pto_deadline_micros(cx, now_micros)
            .map_err(|err| ConnectionRouterError::PacketProcessingFailed {
                connection_id,
                reason: err.to_string(),
            })?
            .and_then(|deadline| {
                let delta = deadline.saturating_sub(now_micros);
                origin
                    .checked_add(Duration::from_micros(deadline))
                    .or_else(|| now_instant.checked_add(Duration::from_micros(delta)))
            });
        Ok(())
    }

    /// Get the next timer deadline across all connections.
    pub fn next_timer_deadline(&self) -> Option<Instant> {
        self.connections
            .values()
            .filter_map(|handle| handle.next_timer_deadline)
            .min()
    }

    /// Process timer events for connections.
    pub async fn process_timer_events(
        &mut self,
        cx: &Cx,
        current_time: Instant,
    ) -> Result<Vec<OutgoingPacket>, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let mut outgoing_packets = Vec::new();
        let origin = self.clock_origin;

        for (connection_id, handle) in &mut self.connections {
            if let Some(deadline) = handle.next_timer_deadline {
                if current_time >= deadline {
                    cx.trace(&format!(
                        "Timer fired for connection {connection_id:?} at {current_time:?}"
                    ));

                    handle.next_timer_deadline = None;
                    handle.connection.on_probe_timeout(cx).map_err(|err| {
                        ConnectionRouterError::PacketProcessingFailed {
                            connection_id: *connection_id,
                            reason: err.to_string(),
                        }
                    })?;
                    let peer_addr = handle.peer_addr;
                    outgoing_packets.extend(
                        drain_connection_frames(
                            cx,
                            *connection_id,
                            handle,
                            PacketNumberSpace::ApplicationData,
                            peer_addr,
                            current_time,
                            instant_micros_from(origin, current_time),
                        )
                        .await?,
                    );
                    Self::refresh_connection_timer(
                        cx,
                        *connection_id,
                        handle,
                        origin,
                        instant_micros_from(origin, current_time),
                        current_time,
                    )?;
                }
            }
        }

        Ok(outgoing_packets)
    }

    /// Get connection statistics for observability.
    pub fn connection_stats(&self) -> ConnectionRouterStats {
        let active_connections = self.connections.len();
        let established_connections = self
            .connections
            .values()
            .filter(|h| h.established_at.is_some())
            .count();

        ConnectionRouterStats {
            active_connections,
            established_connections,
            pending_connections: active_connections - established_connections,
        }
    }

    fn decode_routing_info(
        &self,
        packet: &ReceivedPacket,
    ) -> Result<PacketRoutingInfo, QuicCoreError> {
        if packet.data.first().is_some_and(|first| first & 0x80 != 0) {
            let (header, header_len) = PacketHeader::decode(&packet.data, 0)?;
            return PacketRoutingInfo::from_header(header, header_len);
        }

        for cid_len in self.known_connection_id_lengths() {
            if let Ok((header, header_len)) = PacketHeader::decode(&packet.data, cid_len) {
                let info = PacketRoutingInfo::from_header(header, header_len)?;
                if self.connections.contains_key(&info.destination_cid) {
                    return Ok(info);
                }
            }
        }

        let (header, header_len) = PacketHeader::decode(&packet.data, 0)?;
        PacketRoutingInfo::from_header(header, header_len)
    }

    fn known_connection_id_lengths(&self) -> Vec<usize> {
        let mut lengths = self
            .connections
            .keys()
            .map(ConnectionId::len)
            .collect::<Vec<_>>();
        lengths.sort_unstable_by(|a, b| b.cmp(a));
        lengths.dedup();
        if !lengths.contains(&0) {
            lengths.push(0);
        }
        lengths
    }

    fn instant_micros(&self, instant: Instant) -> u64 {
        instant_micros_from(self.clock_origin, instant)
    }

    /// Allocate a new connection ID.
    pub(crate) fn allocate_connection_id(&mut self) -> ConnectionId {
        let id = self.next_connection_id;
        self.next_connection_id += 1;

        // Create connection ID from counter
        let id_bytes = id.to_be_bytes();
        ConnectionId::new(&id_bytes).expect("Connection ID from counter should always be valid")
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PacketRoutingKind {
    Initial,
    Handshake,
    ZeroRtt,
    OneRtt,
    Retry,
}

#[derive(Debug, Clone)]
struct PacketRoutingInfo {
    destination_cid: ConnectionId,
    kind: PacketRoutingKind,
    space: PacketNumberSpace,
    packet_number: u64,
    key_phase: bool,
    header_len: usize,
}

impl PacketRoutingInfo {
    fn from_header(header: PacketHeader, header_len: usize) -> Result<Self, QuicCoreError> {
        match header {
            PacketHeader::Long(header) => {
                let (kind, space) = match header.packet_type {
                    LongPacketType::Initial => {
                        (PacketRoutingKind::Initial, PacketNumberSpace::Initial)
                    }
                    LongPacketType::ZeroRtt => (
                        PacketRoutingKind::ZeroRtt,
                        PacketNumberSpace::ApplicationData,
                    ),
                    LongPacketType::Handshake => {
                        (PacketRoutingKind::Handshake, PacketNumberSpace::Handshake)
                    }
                    LongPacketType::Retry => (PacketRoutingKind::Retry, PacketNumberSpace::Initial),
                };
                Ok(Self {
                    destination_cid: header.dst_cid,
                    kind,
                    space,
                    packet_number: header.packet_number,
                    key_phase: false,
                    header_len,
                })
            }
            PacketHeader::Retry(header) => Ok(Self {
                destination_cid: header.dst_cid,
                kind: PacketRoutingKind::Retry,
                space: PacketNumberSpace::Initial,
                packet_number: 0,
                key_phase: false,
                header_len,
            }),
            PacketHeader::Short(header) => Ok(Self {
                destination_cid: header.dst_cid,
                kind: PacketRoutingKind::OneRtt,
                space: PacketNumberSpace::ApplicationData,
                packet_number: header.packet_number,
                key_phase: header.key_phase,
                header_len,
            }),
        }
    }
}

fn instant_micros_from(origin: Instant, instant: Instant) -> u64 {
    instant
        .checked_duration_since(origin)
        .unwrap_or(Duration::ZERO)
        .as_micros()
        .min(u128::from(u64::MAX)) as u64
}

async fn drain_connection_frames(
    cx: &Cx,
    connection_id: ConnectionId,
    handle: &mut ConnectionHandle,
    space: PacketNumberSpace,
    dst_addr: SocketAddr,
    now: Instant,
    now_micros: u64,
) -> Result<Vec<OutgoingPacket>, ConnectionRouterError> {
    let max_frame_bytes = if space == PacketNumberSpace::ApplicationData {
        if handle.packet_protection.is_none() {
            return Err(ConnectionRouterError::PacketProtectionUnavailable { connection_id });
        }
        PROTECTED_1RTT_MAX_PACKET_BYTES.saturating_sub(protected_1rtt_packet_len(connection_id, 0))
    } else {
        PROTECTED_1RTT_MAX_PACKET_BYTES
    };
    let frames = handle
        .connection
        .generate_frames(cx, space, max_frame_bytes)
        .map_err(|err| ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: err.to_string(),
        })?;
    if frames.is_empty() {
        return Ok(Vec::new());
    }

    let mut payload = crate::bytes::BytesMut::new();
    NativeQuicConnection::encode_frames(&frames, &mut payload).map_err(
        |err: NativeQuicConnectionError| ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: err.to_string(),
        },
    )?;
    let data = if space == PacketNumberSpace::ApplicationData {
        match handle.packet_protection.as_mut() {
            Some(packet_protection) => {
                assemble_protected_1rtt_packet(
                    cx,
                    connection_id,
                    &mut handle.connection,
                    packet_protection,
                    payload.as_ref(),
                    now_micros,
                    frames.iter().any(is_ack_eliciting),
                )
                .await?
            }
            None => {
                return Err(ConnectionRouterError::PacketProtectionUnavailable { connection_id });
            }
        }
    } else {
        payload.to_vec()
    };

    Ok(vec![OutgoingPacket {
        dst_addr,
        data,
        send_time: Some(now),
    }])
}

async fn assemble_protected_1rtt_packet(
    cx: &Cx,
    connection_id: ConnectionId,
    connection: &mut NativeQuicConnection,
    packet_protection: &mut ConnectionPacketProtection,
    payload: &[u8],
    now_micros: u64,
    ack_eliciting: bool,
) -> Result<Vec<u8>, ConnectionRouterError> {
    let packet_len = protected_1rtt_packet_len(connection_id, payload.len());
    if packet_len > PROTECTED_1RTT_MAX_PACKET_BYTES {
        return Err(ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: format!(
                "protected 1-RTT packet length {packet_len} exceeds max {PROTECTED_1RTT_MAX_PACKET_BYTES}"
            ),
        });
    }
    let packet_number = connection
        .on_packet_sent(
            cx,
            PacketNumberSpace::ApplicationData,
            packet_len as u64,
            ack_eliciting,
            true,
            now_micros,
        )
        .map_err(|err| ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: err.to_string(),
        })?;
    let key_phase = connection.tls().local_key_phase();
    let header = PacketHeader::Short(ShortHeader {
        spin: false,
        key_phase,
        dst_cid: connection_id,
        packet_number,
        packet_number_len: PROTECTED_1RTT_PACKET_NUMBER_LEN,
    });
    let mut header_bytes = Vec::new();
    header.encode(&mut header_bytes).map_err(|err| {
        ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: err.to_string(),
        }
    })?;
    let protected = match packet_protection
        .protection
        .protect_packet(
            cx,
            PacketProtectionRequest {
                space: PacketProtectionSpace::OneRtt,
                key_phase,
                packet_number,
                associated_data: &header_bytes,
                payload,
            },
        )
        .await
    {
        Outcome::Ok(packet) => packet,
        Outcome::Err(err) => {
            return Err(ConnectionRouterError::PacketProcessingFailed {
                connection_id,
                reason: format!("1-RTT packet protection failed: {err:?}"),
            });
        }
        Outcome::Cancelled(_) => return Err(ConnectionRouterError::Cancelled),
        Outcome::Panicked(payload) => {
            return Err(ConnectionRouterError::PacketProcessingFailed {
                connection_id,
                reason: format!("1-RTT packet protection panicked: {payload:?}"),
            });
        }
    };

    let mut packet =
        Vec::with_capacity(header_bytes.len() + protected.ciphertext.len() + protected.tag.len());
    packet.extend_from_slice(&header_bytes);
    packet.extend_from_slice(&protected.ciphertext);
    packet.extend_from_slice(&protected.tag);
    Ok(packet)
}

async fn unprotect_1rtt_packet(
    cx: &Cx,
    connection_id: ConnectionId,
    handle: &mut ConnectionHandle,
    associated_data: &[u8],
    protected_payload: &[u8],
    packet_number: u64,
    key_phase: bool,
) -> Result<Vec<u8>, ConnectionRouterError> {
    let Some(packet_protection) = handle.packet_protection.as_mut() else {
        return Err(ConnectionRouterError::PacketProtectionUnavailable { connection_id });
    };
    if protected_payload.len() < PROTECTED_1RTT_TAG_LEN {
        return Err(ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: format!(
                "protected 1-RTT packet too short: payload_len={}, tag_len={PROTECTED_1RTT_TAG_LEN}",
                protected_payload.len()
            ),
        });
    }

    let tag_offset = protected_payload.len() - PROTECTED_1RTT_TAG_LEN;
    let tag: [u8; PROTECTED_1RTT_TAG_LEN] = protected_payload[tag_offset..]
        .try_into()
        .expect("tag length checked above");
    let protected = ProtectedPacket {
        space: PacketProtectionSpace::OneRtt,
        key_phase,
        packet_number,
        ciphertext: protected_payload[..tag_offset].to_vec(),
        tag,
        proof: ProtectionProof {
            provider_kind: packet_protection.protection.provider_kind(),
            space: PacketProtectionSpace::OneRtt,
            key_phase,
            generation: 0,
            transcript_hash: TranscriptHash::from_bytes([0; 32]),
            failure_code: None,
        },
    };

    match packet_protection
        .protection
        .unprotect_packet(cx, &protected, associated_data)
        .await
    {
        Outcome::Ok(packet) => Ok(packet.plaintext),
        Outcome::Err(err) => Err(ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: format!("1-RTT packet unprotection failed: {err:?}"),
        }),
        Outcome::Cancelled(_) => Err(ConnectionRouterError::Cancelled),
        Outcome::Panicked(payload) => Err(ConnectionRouterError::PacketProcessingFailed {
            connection_id,
            reason: format!("1-RTT packet unprotection panicked: {payload:?}"),
        }),
    }
}

const PROTECTED_1RTT_MAX_PACKET_BYTES: usize = 1_200;
const PROTECTED_1RTT_PACKET_NUMBER_LEN: u8 = 4;
const PROTECTED_1RTT_TAG_LEN: usize = 16;

fn protected_1rtt_packet_len(connection_id: ConnectionId, payload_len: usize) -> usize {
    1 + connection_id.len()
        + usize::from(PROTECTED_1RTT_PACKET_NUMBER_LEN)
        + payload_len
        + PROTECTED_1RTT_TAG_LEN
}

fn is_ack_eliciting(frame: &crate::net::atp::protocol::quic_frames::QuicFrame) -> bool {
    !matches!(
        frame,
        crate::net::atp::protocol::quic_frames::QuicFrame::Padding { .. }
            | crate::net::atp::protocol::quic_frames::QuicFrame::Ack { .. }
    )
}

/// Statistics about the connection router state.
#[derive(Debug, Clone)]
pub struct ConnectionRouterStats {
    /// Number of active connections in the routing table.
    pub active_connections: usize,
    /// Number of established connections.
    pub established_connections: usize,
    /// Number of connections still in handshake.
    pub pending_connections: usize,
}

/// Timer scheduler for QUIC connections that integrates with Asupersync runtime.
#[derive(Debug)]
pub struct QuicTimerScheduler {
    /// Currently scheduled timer sleep.
    current_sleep: Option<Sleep>,
    /// Next deadline we're sleeping until.
    current_deadline: Option<Instant>,
}

impl QuicTimerScheduler {
    /// Create a new timer scheduler.
    pub fn new() -> Self {
        Self {
            current_sleep: None,
            current_deadline: None,
        }
    }

    /// Schedule a timer to fire at the given deadline.
    ///
    /// If a timer is already scheduled for an earlier time, this is a no-op.
    /// If the new deadline is earlier, the current timer is cancelled and
    /// a new one is scheduled.
    pub async fn schedule_timer(
        &mut self,
        cx: &Cx,
        deadline: Instant,
    ) -> Result<(), ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        let now = Instant::now();

        // If deadline is in the past, fire immediately
        if deadline <= now {
            return Ok(());
        }

        // Check if we need to reschedule
        let should_reschedule = match self.current_deadline {
            Some(current) => deadline < current,
            None => true,
        };

        if should_reschedule {
            let duration = deadline.saturating_duration_since(now);
            let duration_from_now = deadline.saturating_duration_since(Instant::now());
            let time_deadline = crate::Time::from_nanos(duration_from_now.as_nanos() as u64);
            self.current_sleep = Some(Sleep::new(time_deadline));
            self.current_deadline = Some(deadline);

            cx.trace(&format!(
                "Scheduled QUIC timer for {deadline:?} (in {duration:?})"
            ));
        }

        Ok(())
    }

    /// Wait for the next timer to fire.
    ///
    /// Returns the deadline that was reached, or None if no timer was scheduled.
    pub async fn wait_for_timer(
        &mut self,
        cx: &Cx,
    ) -> Result<Option<Instant>, ConnectionRouterError> {
        if cx.checkpoint().is_err() {
            return Err(ConnectionRouterError::Cancelled);
        }

        if let Some(sleep) = self.current_sleep.take() {
            let deadline = self.current_deadline.take();

            // Wait for the timer to fire
            sleep.await;

            cx.trace(&format!("QUIC timer fired for {deadline:?}"));
            Ok(deadline)
        } else {
            Ok(None)
        }
    }

    /// Check if a timer is currently scheduled.
    pub fn has_pending_timer(&self) -> bool {
        self.current_sleep.is_some()
    }

    /// Get the current timer deadline if any.
    pub fn current_deadline(&self) -> Option<Instant> {
        self.current_deadline
    }

    /// Cancel the pending timer, if one is armed.
    pub fn cancel(&mut self) {
        self.current_sleep = None;
        self.current_deadline = None;
    }
}

impl Default for QuicTimerScheduler {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bytes::{Bytes, BytesMut};
    use crate::net::atp::protocol::quic_frames::QuicFrame;
    use crate::net::atp::quic::AtpPacketProtection;
    use crate::net::quic_core::{LongHeader, LongPacketType, PacketHeader};
    use crate::net::quic_native::QuicHandshakeTranscript;
    use crate::test_utils::run_test_with_cx;

    #[test]
    fn test_connection_router_creation() {
        let config = NativeQuicConnectionConfig::default();
        let router = ConnectionRouter::new(config);

        assert_eq!(router.connections.len(), 0);
        assert_eq!(router.next_connection_id, 1);
    }

    #[test]
    fn test_connection_id_allocation() {
        run_test_with_cx(|_cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);

            let id1 = router.allocate_connection_id();
            let id2 = router.allocate_connection_id();

            assert_ne!(id1, id2);
            assert!(router.next_connection_id > 2);
        });
    }

    #[test]
    fn test_connection_creation() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);

            let connection_id = router.allocate_connection_id();
            let peer_addr = "127.0.0.1:12345".parse().unwrap();

            router
                .create_connection(&cx, connection_id, peer_addr, false)
                .await
                .expect("connection creation should succeed");

            assert_eq!(router.connections.len(), 1);
            assert!(router.connections.contains_key(&connection_id));
        });
    }

    #[test]
    fn create_connection_rejects_duplicate_connection_id_without_overwrite() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let connection_id = ConnectionId::new(&[0x31, 0x71, 0x00, 0x01]).expect("cid");
            let original_peer: SocketAddr = "127.0.0.1:4401".parse().unwrap();
            let colliding_peer: SocketAddr = "127.0.0.1:4402".parse().unwrap();

            router
                .create_connection(&cx, connection_id, original_peer, true)
                .await
                .expect("first connection creation should succeed");

            let err = router
                .create_connection(&cx, connection_id, colliding_peer, false)
                .await
                .expect_err("duplicate destination CID must fail closed");
            assert!(matches!(
                err,
                ConnectionRouterError::ConnectionCreationFailed(ref msg)
                    if msg.contains("connection ID collision")
            ));
            assert_eq!(router.connections.len(), 1);
            assert_eq!(
                router
                    .connections
                    .get(&connection_id)
                    .expect("original connection remains")
                    .peer_addr,
                original_peer
            );
        });
    }

    #[test]
    fn connection_limit_rejects_create_and_drops_unknown_initials() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::with_max_connections(config, 0);
            assert_eq!(router.max_connections, 1);

            let peer_addr: SocketAddr = "127.0.0.1:4403".parse().unwrap();
            let first = ConnectionId::new(&[0x31, 0x71, 0x00, 0x02]).expect("first cid");
            router
                .create_connection(&cx, first, peer_addr, true)
                .await
                .expect("first connection within normalized limit should succeed");

            let second = ConnectionId::new(&[0x31, 0x71, 0x00, 0x03]).expect("second cid");
            let err = router
                .create_connection(&cx, second, peer_addr, true)
                .await
                .expect_err("second connection must hit the cap");
            assert!(matches!(
                err,
                ConnectionRouterError::ConnectionCreationFailed(ref msg)
                    if msg.contains("connection limit reached")
            ));

            let packet = ReceivedPacket {
                src_addr: peer_addr,
                data: encode_long_packet(second, LongPacketType::Initial, 0, QuicFrame::Ping),
                receive_time: Instant::now(),
                transmit_time: None,
            };
            match router.route_packet(&cx, packet).await.expect("route") {
                RoutingResult::Drop { reason } => {
                    assert!(reason.contains("connection limit reached"));
                }
                other => panic!("full router must not advertise a new connection: {other:?}"),
            }
        });
    }

    #[test]
    fn test_take_connection_removes_handle_and_preserves_peer() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let connection_id = ConnectionId::new(&[0x10, 0x00, 0x00, 0x01]).expect("cid");
            let peer_addr: SocketAddr = "127.0.0.1:5544".parse().unwrap();
            router
                .create_connection(&cx, connection_id, peer_addr, true)
                .await
                .expect("connection creation should succeed");

            let accepted = router
                .take_connection(&cx, connection_id)
                .expect("connection should be handed off");
            assert_eq!(accepted.connection_id, connection_id);
            assert_eq!(accepted.peer_addr, peer_addr);
            assert_eq!(accepted.connection.pending_outbound_datagram_count(), 0);
            assert!(!router.connections.contains_key(&connection_id));
            assert_eq!(router.connection_stats().active_connections, 0);

            let err = router
                .take_connection(&cx, connection_id)
                .expect_err("missing connection must fail closed");
            assert!(matches!(
                err,
                ConnectionRouterError::ConnectionNotFound(id) if id == connection_id
            ));
        });
    }

    #[test]
    fn test_take_next_connection_uses_deterministic_connection_id_order() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let low = ConnectionId::new(&[0x01, 0x00, 0x00, 0x00]).expect("low cid");
            let mid = ConnectionId::new(&[0x10, 0x00, 0x00, 0x00]).expect("mid cid");
            let high = ConnectionId::new(&[0xff, 0x00, 0x00, 0x00]).expect("high cid");
            let peer_addr: SocketAddr = "127.0.0.1:5545".parse().unwrap();

            for connection_id in [high, low, mid] {
                router
                    .create_connection(&cx, connection_id, peer_addr, true)
                    .await
                    .expect("connection creation should succeed");
            }

            let first = router
                .take_next_connection(&cx)
                .expect("take should succeed")
                .expect("connection should exist");
            assert_eq!(first.connection_id, low);

            let second = router
                .take_next_connection(&cx)
                .expect("take should succeed")
                .expect("connection should exist");
            assert_eq!(second.connection_id, mid);

            let third = router
                .take_next_connection(&cx)
                .expect("take should succeed")
                .expect("connection should exist");
            assert_eq!(third.connection_id, high);

            assert!(
                router
                    .take_next_connection(&cx)
                    .expect("empty take should succeed")
                    .is_none()
            );
        });
    }

    #[test]
    fn test_long_header_initial_routes_as_new_connection() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let dst_cid = ConnectionId::new(&[0xaa, 0xbb, 0xcc]).expect("cid");
            let src_addr: SocketAddr = "127.0.0.1:4433".parse().unwrap();
            let packet = ReceivedPacket {
                src_addr,
                data: encode_long_packet(dst_cid, LongPacketType::Initial, 0, QuicFrame::Ping),
                receive_time: Instant::now(),
                transmit_time: None,
            };

            match router.route_packet(&cx, packet).await.expect("route") {
                RoutingResult::NewConnection { peer_addr, .. } => assert_eq!(peer_addr, src_addr),
                other => panic!("expected new connection, got {other:?}"),
            }
        });
    }

    #[test]
    fn test_new_initial_reroutes_after_connection_creation() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let dst_cid = ConnectionId::new(&[0xda, 0x7a, 0x00, 0x01]).expect("cid");
            let src_addr: SocketAddr = "127.0.0.1:4436".parse().unwrap();
            let packet = ReceivedPacket {
                src_addr,
                data: encode_long_packet(dst_cid, LongPacketType::Initial, 7, QuicFrame::Ping),
                receive_time: Instant::now(),
                transmit_time: None,
            };

            let triggering_packet = match router.route_packet(&cx, packet).await.expect("route") {
                RoutingResult::NewConnection {
                    connection_id,
                    peer_addr,
                    triggering_packet,
                    outgoing_packets,
                } => {
                    assert_eq!(connection_id, dst_cid);
                    assert_eq!(peer_addr, src_addr);
                    assert!(outgoing_packets.is_empty());
                    triggering_packet
                }
                other => panic!("expected new connection, got {other:?}"),
            };

            router
                .create_connection(&cx, dst_cid, src_addr, true)
                .await
                .expect("connection creation should succeed");

            match router
                .route_packet(&cx, triggering_packet)
                .await
                .expect("reroute")
            {
                RoutingResult::Routed {
                    connection_id,
                    outgoing_packets,
                } => {
                    assert_eq!(connection_id, dst_cid);
                    assert_eq!(outgoing_packets.len(), 1);
                    assert_eq!(outgoing_packets[0].dst_addr, src_addr);
                    assert!(!outgoing_packets[0].data.is_empty());
                }
                other => panic!("expected triggering Initial to reroute, got {other:?}"),
            }
        });
    }

    #[test]
    fn test_existing_connection_processes_ping_and_emits_ack_frame() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let connection_id = router.allocate_connection_id();
            let peer_addr: SocketAddr = "127.0.0.1:4434".parse().unwrap();
            router
                .create_connection(&cx, connection_id, peer_addr, false)
                .await
                .expect("connection creation should succeed");

            let packet = ReceivedPacket {
                src_addr: peer_addr,
                data: encode_long_packet(
                    connection_id,
                    LongPacketType::Initial,
                    42,
                    QuicFrame::Ping,
                ),
                receive_time: Instant::now(),
                transmit_time: None,
            };

            match router.route_packet(&cx, packet).await.expect("route") {
                RoutingResult::Routed {
                    outgoing_packets, ..
                } => {
                    assert_eq!(outgoing_packets.len(), 1);
                    assert_eq!(outgoing_packets[0].dst_addr, peer_addr);
                    assert!(!outgoing_packets[0].data.is_empty());
                }
                other => panic!("expected routed packet, got {other:?}"),
            }
        });
    }

    #[test]
    fn test_application_datagram_handoff_uses_protected_short_header_packet() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let connection_id = ConnectionId::new(&[0xa1, 0x01, 0x00, 0x01]).expect("cid");
            let peer_addr: SocketAddr = "127.0.0.1:4435".parse().unwrap();
            router
                .create_connection(&cx, connection_id, peer_addr, false)
                .await
                .expect("connection creation should succeed");
            router
                .install_packet_protection(
                    &cx,
                    connection_id,
                    deterministic_one_rtt_protection(&cx).await,
                )
                .expect("install packet protection");
            let mut receiver = ConnectionRouter::new(config);
            receiver
                .create_connection(&cx, connection_id, peer_addr, false)
                .await
                .expect("receiver connection creation should succeed");
            receiver
                .install_packet_protection(
                    &cx,
                    connection_id,
                    deterministic_one_rtt_protection(&cx).await,
                )
                .expect("install receiver packet protection");

            let datagram = Bytes::from_static(b"a1 protected udp symbol");
            {
                let handle = router
                    .connections
                    .get_mut(&connection_id)
                    .expect("connection handle");
                establish_for_application_data(&cx, &mut handle.connection);
                handle
                    .connection
                    .send_datagram(&cx, datagram.clone())
                    .expect("queue datagram");
            }

            let now = Instant::now();
            let packets = {
                let handle = router
                    .connections
                    .get_mut(&connection_id)
                    .expect("connection handle");
                drain_connection_frames(
                    &cx,
                    connection_id,
                    handle,
                    PacketNumberSpace::ApplicationData,
                    peer_addr,
                    now,
                    42_000,
                )
                .await
                .expect("drain protected packet")
            };
            assert_eq!(packets.len(), 1);
            assert_eq!(packets[0].dst_addr, peer_addr);
            assert_eq!(packets[0].send_time, Some(now));
            assert!(packets[0].data.len() <= PROTECTED_1RTT_MAX_PACKET_BYTES);

            let mut raw_frame_payload = BytesMut::new();
            QuicFrame::Datagram { data: datagram }
                .encode(&mut raw_frame_payload)
                .expect("encode raw DATAGRAM frame");
            let packet = &packets[0].data;
            assert_ne!(packet.as_slice(), raw_frame_payload.as_ref());

            let (decoded, header_len) =
                PacketHeader::decode(packet, connection_id.len()).expect("decode short header");
            let PacketHeader::Short(header) = decoded else {
                panic!("expected a protected 1-RTT short header packet");
            };
            assert!(!header.spin);
            assert!(!header.key_phase);
            assert_eq!(header.dst_cid, connection_id);
            assert_eq!(header.packet_number, 0);
            assert_eq!(header.packet_number_len, PROTECTED_1RTT_PACKET_NUMBER_LEN);

            let protected_payload = &packet[header_len..];
            assert_eq!(
                protected_payload.len(),
                raw_frame_payload.len() + PROTECTED_1RTT_TAG_LEN
            );
            assert_ne!(
                &protected_payload[..raw_frame_payload.len()],
                raw_frame_payload.as_ref()
            );

            {
                let handle = receiver
                    .connections
                    .get_mut(&connection_id)
                    .expect("receiver connection handle");
                establish_for_application_data(&cx, &mut handle.connection);
            }
            let received = ReceivedPacket {
                src_addr: peer_addr,
                data: packet.clone(),
                receive_time: Instant::now(),
                transmit_time: None,
            };
            match receiver
                .route_packet(&cx, received)
                .await
                .expect("route protected packet")
            {
                RoutingResult::Routed {
                    connection_id: routed_id,
                    ..
                } => assert_eq!(routed_id, connection_id),
                other => panic!("expected routed protected packet, got {other:?}"),
            }
            let received_datagram = receiver
                .connections
                .get_mut(&connection_id)
                .expect("receiver connection handle")
                .connection
                .recv_datagram()
                .expect("datagram delivered after unprotect");
            assert_eq!(received_datagram.as_ref(), b"a1 protected udp symbol");
        });
    }

    #[test]
    fn test_application_datagram_handoff_without_protection_fails_closed() {
        run_test_with_cx(|cx| async move {
            let config = NativeQuicConnectionConfig::default();
            let mut router = ConnectionRouter::new(config);
            let connection_id = ConnectionId::new(&[0xa1, 0x01, 0x00, 0x02]).expect("cid");
            let peer_addr: SocketAddr = "127.0.0.1:4436".parse().unwrap();
            router
                .create_connection(&cx, connection_id, peer_addr, false)
                .await
                .expect("connection creation should succeed");

            let handle = router
                .connections
                .get_mut(&connection_id)
                .expect("connection handle");
            establish_for_application_data(&cx, &mut handle.connection);
            handle
                .connection
                .send_datagram(&cx, Bytes::from_static(b"must not leak raw"))
                .expect("queue datagram");

            let err = drain_connection_frames(
                &cx,
                connection_id,
                handle,
                PacketNumberSpace::ApplicationData,
                peer_addr,
                Instant::now(),
                42_000,
            )
            .await
            .expect_err("missing 1-RTT packet protection must fail closed");
            assert!(matches!(
                err,
                ConnectionRouterError::PacketProtectionUnavailable { connection_id: id }
                    if id == connection_id
            ));
            assert_eq!(handle.connection.pending_outbound_datagram_count(), 1);
        });
    }

    #[test]
    fn test_timer_scheduler_basic() {
        run_test_with_cx(|cx| async move {
            let mut scheduler = QuicTimerScheduler::new();

            assert!(!scheduler.has_pending_timer());
            assert_eq!(scheduler.current_deadline(), None);

            let deadline = Instant::now() + std::time::Duration::from_millis(10);
            scheduler
                .schedule_timer(&cx, deadline)
                .await
                .expect("timer scheduling should succeed");

            assert!(scheduler.has_pending_timer());
            assert_eq!(scheduler.current_deadline(), Some(deadline));
        });
    }

    fn encode_long_packet(
        dst_cid: ConnectionId,
        packet_type: LongPacketType,
        packet_number: u64,
        frame: QuicFrame,
    ) -> Vec<u8> {
        let mut payload = BytesMut::new();
        frame.encode(&mut payload).expect("frame encode");
        let header = PacketHeader::Long(LongHeader {
            packet_type,
            version: 1,
            dst_cid,
            src_cid: ConnectionId::new(&[0x01, 0x02, 0x03, 0x04]).expect("src cid"),
            token: Vec::new(),
            payload_length: payload.len() as u64 + 1,
            packet_number,
            packet_number_len: 1,
        });
        let mut out = Vec::new();
        header.encode(&mut out).expect("header encode");
        out.extend_from_slice(&payload);
        out
    }

    async fn deterministic_one_rtt_protection(cx: &Cx) -> AtpPacketProtection {
        let mut transcript = QuicHandshakeTranscript::new();
        transcript.record("client_initial", b"a1 client hello");
        transcript.record("server_handshake", b"a1 server hello");
        let mut protection =
            AtpPacketProtection::new_client(true).expect("deterministic ATP packet protection");
        protection
            .derive_keys(
                cx,
                PacketProtectionSpace::OneRtt,
                &transcript,
                b"asupersync a1 protected udp handoff",
            )
            .await
            .expect("derive 1-RTT keys");
        protection
    }

    fn establish_for_application_data(cx: &Cx, connection: &mut NativeQuicConnection) {
        connection.begin_handshake(cx).expect("begin");
        connection
            .on_handshake_keys_available(cx)
            .expect("handshake keys");
        connection.on_1rtt_keys_available(cx).expect("1rtt keys");
        connection.record_verified_server_identity();
        connection
            .on_handshake_confirmed(cx)
            .expect("handshake confirmed");
    }
}