nxtquic-api 0.1.2

//! Endpoint configuration and socket ownership.
//!
//! This module owns the UDP socket used by an endpoint.  Packet driving is
//! intentionally kept separate from socket binding so a server cannot report
//! that it is listening before the operating system has accepted the bind.

use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};

use crate::connection::Connection;
use nxtquic_crypto::{
    CryptoProvider, DirectionalKeys, HandshakeState, HeaderProtectionKey, Keys, PacketKey,
    RustlsCryptoProvider,
};
use nxtquic_proto::{
    ConnectionId, StreamId, Version,
    frame::{AckFrame, CryptoFrame, Frame, StreamFrame},
    packet::{
        PacketType, decode_header,
        initial::parse_initial,
        protection::{apply_header_protection, remove_header_protection},
    },
    varint::VarInt,
};

/// Configuration for an endpoint.
#[derive(Default, Clone, Debug)]
pub struct EndpointConfig;

/// Configuration for a server endpoint.
#[derive(Clone)]
pub struct ServerConfig {
    rustls: Arc<rustls::ServerConfig>,
    transport: TransportConfig,
}

impl std::fmt::Debug for ServerConfig {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("ServerConfig")
            .field("transport", &self.transport)
            .finish_non_exhaustive()
    }
}

impl ServerConfig {
    /// Starts building a server configuration.
    pub fn builder() -> ServerConfigBuilder {
        ServerConfigBuilder::default()
    }

    /// Returns the Rustls configuration used by this endpoint.
    pub fn rustls(&self) -> &Arc<rustls::ServerConfig> {
        &self.rustls
    }

    /// Returns the transport configuration.
    pub fn transport(&self) -> &TransportConfig {
        &self.transport
    }
}

/// Builder for [`ServerConfig`].
#[derive(Default)]
pub struct ServerConfigBuilder {
    rustls: Option<Arc<rustls::ServerConfig>>,
    transport: TransportConfig,
}

impl ServerConfigBuilder {
    /// Uses an already configured Rustls server configuration.
    pub fn with_rustls(mut self, config: Arc<rustls::ServerConfig>) -> io::Result<Self> {
        if config
            .alpn_protocols
            .iter()
            .all(|protocol| protocol.as_slice() != b"h3")
        {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "a QUIC HTTP/3 server configuration must advertise ALPN h3",
            ));
        }
        self.rustls = Some(config);
        Ok(self)
    }

    /// Applies transport settings before building the server configuration.
    pub fn with_transport_config(mut self, configure: impl FnOnce(&mut TransportConfig)) -> Self {
        configure(&mut self.transport);
        self
    }

    /// Builds the server configuration.
    pub fn build(self) -> io::Result<ServerConfig> {
        let rustls = self.rustls.ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "missing Rustls server configuration",
            )
        })?;
        Ok(ServerConfig {
            rustls,
            transport: self.transport,
        })
    }
}

/// Transport settings shared by all server connections.
#[derive(Default, Clone, Debug)]
pub struct TransportConfig {
    max_idle_timeout: Option<Duration>,
}

impl TransportConfig {
    /// Sets the maximum time a connection may remain idle.
    pub fn max_idle_timeout(&mut self, timeout: Option<Duration>) -> &mut Self {
        self.max_idle_timeout = timeout;
        self
    }

    /// Returns the configured idle timeout.
    pub fn idle_timeout(&self) -> Option<Duration> {
        self.max_idle_timeout
    }
}

/// Configuration for a client endpoint.
#[derive(Default, Clone, Debug)]
pub struct ClientConfig;

/// A QUIC endpoint.
pub struct Endpoint {
    config: EndpointConfig,
    server_config: Option<ServerConfig>,
    socket: Arc<tokio::net::UdpSocket>,
    incoming: Arc<tokio::sync::Mutex<tokio::sync::mpsc::Receiver<Incoming>>>,
}

/// An ongoing connection attempt.
pub struct Connecting;

/// An incoming connection attempt.
pub struct Incoming {
    remote_addr: SocketAddr,
    packet: Vec<u8>,
    initial: DecryptedInitial,
    socket: Arc<tokio::net::UdpSocket>,
    server_config: Option<ServerConfig>,
    datagrams: Option<tokio::sync::mpsc::Receiver<Vec<u8>>>,
}

/// A decrypted client Initial packet.  This is kept internal while the
/// endpoint driver is completed, but it deliberately contains decoded frames
/// rather than a synthetic "accepted" connection so callers cannot mistake an
/// unprocessed UDP datagram for a completed handshake.
#[derive(Debug)]
struct DecryptedInitial {
    version: Version,
    dcid: ConnectionId,
    scid: ConnectionId,
    packet_number: u64,
    frames: Vec<Frame>,
}

fn decrypt_client_initial(packet: &[u8]) -> io::Result<DecryptedInitial> {
    let parsed =
        parse_initial(packet).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    let version = match parsed.version {
        1 => Version::V1,
        0x6b33_43cf => Version::V2,
        version => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("unsupported QUIC version {version:#010x}"),
            ));
        }
    };

    let provider = RustlsCryptoProvider::new(None, None);
    let keys = provider.initial_keys(version, &parsed.dcid);
    let mut protected = packet[..parsed.packet_end].to_vec();

    let sample_offset = parsed.packet_number_offset + 4;
    let sample: [u8; 16] = protected[sample_offset..sample_offset + 16]
        .try_into()
        .expect("Initial parser validates the header-protection sample");
    let mask = keys
        .client
        .header_key
        .protection_mask(&sample)
        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    let packet_number_len =
        remove_header_protection(&mut protected, &mask, parsed.packet_number_offset);
    if packet_number_len == 0 || parsed.packet_number_offset + packet_number_len > parsed.packet_end
    {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "invalid QUIC Initial packet number",
        ));
    }

    let mut packet_number = 0_u64;
    for byte in
        &protected[parsed.packet_number_offset..parsed.packet_number_offset + packet_number_len]
    {
        packet_number = (packet_number << 8) | u64::from(*byte);
    }
    let header_end = parsed.packet_number_offset + packet_number_len;
    let mut payload = protected[header_end..parsed.packet_end].to_vec();
    keys.client
        .packet_key
        .decrypt(packet_number, &protected[..header_end], &mut payload)
        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;

    let mut frames = Vec::new();
    let mut payload = &payload[..];
    while !payload.is_empty() {
        let before = payload.len();
        let frame = Frame::decode(&mut payload)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        frames.push(frame);
        if payload.len() == before {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "QUIC frame decoder made no progress",
            ));
        }
    }

    Ok(DecryptedInitial {
        version,
        dcid: parsed.dcid,
        scid: parsed.scid,
        packet_number,
        frames,
    })
}

fn encode_server_initial(
    version: Version,
    client_scid: ConnectionId,
    server_scid: ConnectionId,
    packet_number: u64,
    crypto_data: Vec<u8>,
    keys: nxtquic_crypto::InitialKeys,
) -> io::Result<Vec<u8>> {
    let mut plaintext = Vec::new();
    Frame::Crypto(CryptoFrame {
        offset: VarInt::ZERO,
        data: crypto_data.into(),
    })
    .encode(&mut plaintext);

    let mut header_prefix = Vec::new();
    header_prefix.push(0xc0); // Initial, fixed bit, one-byte packet number
    header_prefix.extend_from_slice(&version.as_u32().to_be_bytes());
    header_prefix.push(client_scid.len() as u8);
    header_prefix.extend_from_slice(client_scid.as_bytes());
    header_prefix.push(server_scid.len() as u8);
    header_prefix.extend_from_slice(server_scid.as_bytes());
    header_prefix.push(0); // empty Retry token

    // Every datagram containing an Initial packet must be at least 1200 bytes
    // (RFC 9000 §14.1). Padding uses PADDING frames, which are zero bytes.
    loop {
        let encrypted_len = plaintext.len() + keys.server.packet_key.tag_len();
        let mut candidate = header_prefix.clone();
        VarInt::from_u64((1 + encrypted_len) as u64)
            .expect("QUIC packet lengths fit in a VarInt")
            .encode(&mut candidate);
        if candidate.len() + 1 + encrypted_len >= 1200 {
            header_prefix = candidate;
            break;
        }
        plaintext.push(0);
    }

    let packet_number_len = 1;
    let packet_number_offset = header_prefix.len();
    header_prefix.push(packet_number as u8);
    keys.server
        .packet_key
        .encrypt(packet_number, &header_prefix, &mut plaintext)
        .map_err(|error| io::Error::new(io::ErrorKind::Other, error))?;
    header_prefix.extend_from_slice(&plaintext);

    let sample: [u8; 16] = header_prefix[packet_number_offset + 4..packet_number_offset + 20]
        .try_into()
        .expect("1200-byte Initial always contains a header-protection sample");
    let mask = keys
        .server
        .header_key
        .protection_mask(&sample)
        .map_err(|error| io::Error::new(io::ErrorKind::Other, error))?;
    apply_header_protection(
        &mut header_prefix,
        &mask,
        packet_number_offset,
        packet_number_len,
    );
    Ok(header_prefix)
}

impl std::fmt::Debug for Incoming {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("Incoming")
            .field("remote_addr", &self.remote_addr)
            .field("packet_len", &self.packet.len())
            .finish()
    }
}

impl Incoming {
    /// Returns the peer address that sent the Initial packet.
    pub fn remote_address(&self) -> SocketAddr {
        self.remote_addr
    }

    /// Returns the first packet received for this connection attempt.
    pub(crate) fn packet(&self) -> &[u8] {
        &self.packet
    }
}

impl Endpoint {
    /// Creates a new endpoint.
    /// Binds the endpoint to a local socket address.
    pub async fn bind(addr: SocketAddr) -> io::Result<Self> {
        Self::bind_inner(EndpointConfig::default(), None, addr).await
    }

    /// Binds a server endpoint using the supplied TLS and transport settings.
    pub async fn server(server_config: ServerConfig, addr: SocketAddr) -> io::Result<Self> {
        Self::bind_inner(EndpointConfig::default(), Some(server_config), addr).await
    }

    async fn bind_inner(
        config: EndpointConfig,
        server_config: Option<ServerConfig>,
        addr: SocketAddr,
    ) -> io::Result<Self> {
        let socket = tokio::net::UdpSocket::bind(addr).await?;
        let socket = Arc::new(socket);
        let (incoming_tx, incoming_rx) = tokio::sync::mpsc::channel(128);
        let receive_socket = Arc::clone(&socket);
        let receive_server_config = server_config.clone();
        tokio::spawn(async move {
            let mut packet = vec![0_u8; 65_535];
            let mut peers: HashMap<SocketAddr, tokio::sync::mpsc::Sender<Vec<u8>>> = HashMap::new();
            loop {
                let (length, remote_addr) = match receive_socket.recv_from(&mut packet).await {
                    Ok(received) => received,
                    Err(_) => break,
                };

                let datagram = &packet[..length];
                if let Some(peer_tx) = peers.get(&remote_addr) {
                    if peer_tx.send(datagram.to_vec()).await.is_err() {
                        peers.remove(&remote_addr);
                    }
                    continue;
                }
                let Ok(header) = decode_header(datagram) else {
                    continue;
                };
                if header.packet_type != PacketType::Initial {
                    continue;
                }
                let Ok(initial) = decrypt_client_initial(datagram) else {
                    continue;
                };

                let (peer_tx, peer_rx) = tokio::sync::mpsc::channel(256);
                peers.insert(remote_addr, peer_tx);

                if incoming_tx
                    .send(Incoming {
                        remote_addr,
                        packet: datagram.to_vec(),
                        initial,
                        socket: Arc::clone(&receive_socket),
                        server_config: receive_server_config.clone(),
                        datagrams: Some(peer_rx),
                    })
                    .await
                    .is_err()
                {
                    break;
                }
            }
        });

        Ok(Self {
            config,
            server_config,
            socket,
            incoming: Arc::new(tokio::sync::Mutex::new(incoming_rx)),
        })
    }

    /// Returns the address the operating system assigned to this endpoint.
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.socket.local_addr()
    }

    /// Returns this endpoint's server configuration, if it is a server.
    pub fn server_config(&self) -> Option<&ServerConfig> {
        self.server_config.as_ref()
    }

    /// Connects to a remote endpoint.
    pub async fn connect(
        &self,
        _addr: SocketAddr,
        _server_name: &str,
    ) -> std::io::Result<Connecting> {
        Ok(Connecting)
    }

    /// Accepts an incoming connection.
    pub async fn accept(&self) -> Option<Incoming> {
        self.incoming.lock().await.recv().await
    }
}

impl Connecting {
    /// Waits for the connection attempt to complete.
    pub async fn await_connection(self) -> std::io::Result<Connection> {
        Ok(Connection::new())
    }
}

impl Incoming {
    /// Accepts the incoming connection.
    pub async fn accept(self) -> std::io::Result<Connection> {
        let config = self.server_config.ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "cannot accept a connection on an endpoint without ServerConfig",
            )
        })?;
        let mut crypto = Vec::new();
        for frame in &self.initial.frames {
            if let Frame::Crypto(frame) = frame {
                if frame.offset.into_inner() != crypto.len() as u64 {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "Initial CRYPTO frames must be contiguous",
                    ));
                }
                crypto.extend_from_slice(&frame.data);
            }
        }
        if crypto.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Initial packet contains no CRYPTO frame",
            ));
        }

        let provider = RustlsCryptoProvider::new(None, Some(Arc::clone(config.rustls())));
        
        let mut transport_params = nxtquic_proto::TransportParameters::default();
        transport_params.initial_max_data = VarInt::from_u64(10 * 1024 * 1024).unwrap();
        transport_params.initial_max_stream_data_bidi_local = VarInt::from_u64(1024 * 1024).unwrap();
        transport_params.initial_max_stream_data_bidi_remote = VarInt::from_u64(1024 * 1024).unwrap();
        transport_params.initial_max_stream_data_uni = VarInt::from_u64(1024 * 1024).unwrap();
        transport_params.initial_max_streams_bidi = VarInt::from_u64(100).unwrap();
        transport_params.initial_max_streams_uni = VarInt::from_u64(100).unwrap();
        if let Some(timeout) = config.transport().idle_timeout() {
            transport_params.max_idle_timeout = VarInt::from_u64(timeout.as_millis() as u64).unwrap_or(VarInt::from_u32(0));
        }
        transport_params.max_udp_payload_size = VarInt::from_u64(1350).unwrap();
        transport_params.active_connection_id_limit = VarInt::from_u64(2).unwrap();
        
        let server_scid = ConnectionId::from_slice(&rand::random::<[u8; 16]>());
        transport_params.initial_source_connection_id = Some(server_scid.clone());
        transport_params.original_destination_connection_id = Some(self.initial.dcid.clone());

        let mut tp_encoded = Vec::new();
        transport_params.encode(&mut tp_encoded);

        let mut handshake = provider
            .new_server_handshake(tp_encoded)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        let output = handshake
            .process(&crypto)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        if output.data.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "TLS produced no Initial handshake data",
            ));
        }

        let keys = provider.initial_keys(self.initial.version, &self.initial.dcid);
        let response = encode_server_initial(
            self.initial.version,
            self.initial.scid,
            server_scid,
            0,
            output.data,
            keys,
        )?;
        self.socket.send_to(&response, self.remote_addr).await?;
        let (outgoing_tx, outgoing_rx) = tokio::sync::mpsc::unbounded_channel();
        let connection = Connection::with_outgoing(outgoing_tx);
        if let Some(datagrams) = self.datagrams {
            let driver_connection = connection.clone();
            let driver_initial = self.initial;
            let driver_socket = Arc::clone(&self.socket);
            let driver_remote = self.remote_addr;
            let driver_initial_keys = provider.initial_keys(driver_initial.version, &driver_initial.dcid);
            let mut driver_handshake_keys = None;
            let mut driver_one_rtt_keys = None;
            while let Some((level, keys)) = handshake.next_keys() {
                if level == nxtquic_crypto::Level::OneRtt {
                    driver_one_rtt_keys = Some(keys);
                } else {
                    driver_handshake_keys = Some((level, keys));
                }
            }
            tokio::spawn(async move {
                drive_server_connection(
                    driver_connection,
                    driver_socket,
                    driver_remote,
                    datagrams,
                    driver_initial,
                    server_scid,
                    driver_initial_keys,
                    driver_handshake_keys,
                    driver_one_rtt_keys,
                    handshake,
                    outgoing_rx,
                )
                .await;
            });
        }
        Ok(connection)
    }
}

struct LongPacket {
    packet_type: PacketType,
    dcid: ConnectionId,
    scid: ConnectionId,
    pn_offset: usize,
    packet_end: usize,
}

fn parse_long_packet(datagram: &[u8]) -> io::Result<LongPacket> {
    if datagram.len() < 7 || datagram[0] & 0xc0 != 0xc0 {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid QUIC long header"));
    }
    let version = u32::from_be_bytes(datagram[1..5].try_into().unwrap());
    if version == 0 {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "version negotiation is not a protected packet"));
    }
    let packet_type = match (datagram[0] >> 4) & 0x03 {
        0 => PacketType::Initial,
        1 => PacketType::ZeroRtt,
        2 => PacketType::Handshake,
        _ => PacketType::Retry,
    };
    let mut offset = 5;
    let dcid_len = datagram[offset] as usize;
    offset += 1;
    let dcid_end = offset.checked_add(dcid_len).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid DCID"))?;
    let dcid = ConnectionId::from_slice(datagram.get(offset..dcid_end).ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "truncated DCID"))?);
    offset = dcid_end;
    let scid_len = *datagram.get(offset).ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "missing SCID length"))? as usize;
    offset += 1;
    let scid_end = offset.checked_add(scid_len).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid SCID"))?;
    let scid = ConnectionId::from_slice(datagram.get(offset..scid_end).ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "truncated SCID"))?);
    offset = scid_end;
    if packet_type == PacketType::Initial {
        let mut rest = &datagram[offset..];
        let token_len = VarInt::decode(&mut rest).map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid Initial token length"))?.into_inner() as usize;
        let consumed = datagram[offset..].len() - rest.len();
        offset += consumed.checked_add(token_len).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid token"))?;
    }
    let mut rest = &datagram[offset..];
    let payload_len = VarInt::decode(&mut rest).map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid QUIC payload length"))?.into_inner() as usize;
    offset += datagram[offset..].len() - rest.len();
    let packet_end = offset.checked_add(payload_len).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid packet length"))?;
    if packet_end > datagram.len() || packet_end < offset + 4 + 16 {
        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "truncated protected QUIC packet"));
    }
    Ok(LongPacket { packet_type, dcid, scid, pn_offset: offset, packet_end })
}

fn decrypt_payload(mut packet: Vec<u8>, parsed: &LongPacket, keys: &DirectionalKeys) -> io::Result<(u64, Vec<u8>)> {
    let sample: [u8; 16] = packet[parsed.pn_offset + 4..parsed.pn_offset + 20].try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "missing header-protection sample"))?;
    let mask = keys.header_key.protection_mask(&sample).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    let pn_len = remove_header_protection(&mut packet, &mask, parsed.pn_offset);
    if pn_len == 0 || parsed.pn_offset + pn_len > parsed.packet_end {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid packet number"));
    }
    let mut packet_number = 0;
    for byte in &packet[parsed.pn_offset..parsed.pn_offset + pn_len] {
        packet_number = (packet_number << 8) | u64::from(*byte);
    }
    let header_end = parsed.pn_offset + pn_len;
    let mut payload = packet[header_end..parsed.packet_end].to_vec();
    keys.packet_key.decrypt(packet_number, &packet[..header_end], &mut payload).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    Ok((packet_number, payload))
}

async fn drive_server_connection(
    connection: Connection,
    socket: Arc<tokio::net::UdpSocket>,
    remote_addr: SocketAddr,
    mut datagrams: tokio::sync::mpsc::Receiver<Vec<u8>>,
    initial: DecryptedInitial,
    server_scid: ConnectionId,
    initial_keys: nxtquic_crypto::InitialKeys,
    mut handshake_keys: Option<(nxtquic_crypto::Level, Keys)>,
    mut one_rtt_keys: Option<Keys>,
    mut handshake: impl HandshakeState,
    mut outgoing: tokio::sync::mpsc::UnboundedReceiver<crate::stream::WriteCommand>,
) {
    let mut largest_handshake = 0;
    let mut application_packet_number = 0;
    loop {
        let Some(event) = (tokio::select! {
            datagram = datagrams.recv() => datagram.map(Ok),
            command = outgoing.recv() => command.map(Err),
            else => None,
        }) else { break };
        let datagram = match event {
            Err(command) => {
                if let Some(keys) = one_rtt_keys.as_ref() {
                    let crate::stream::WriteCommand::Data { stream_id, offset, data, fin } = command;
                    let mut encoded = Vec::new();
                    Frame::Stream(StreamFrame {
                        stream_id: StreamId::from_u64(stream_id),
                        offset: VarInt::from_u64(offset).unwrap(),
                        length: Some(VarInt::from_u64(data.len() as u64).unwrap()),
                        fin,
                        data: bytes::Bytes::from(data),
                    }).encode(&mut encoded);
                    let _ = send_protected_short(&socket, remote_addr, initial.scid, application_packet_number, &encoded, &keys.local).await;
                    application_packet_number += 1;
                }
                continue;
            }
            Ok(datagram) => datagram,
        };
        let Ok(header) = decode_header(&datagram) else { continue };
        let (payload, packet_number, level) = match header.packet_type {
            PacketType::Initial => {
                let Ok(parsed) = parse_long_packet(&datagram) else { continue };
                let Ok((pn, payload)) = decrypt_payload(datagram, &parsed, &initial_keys.client) else { continue };
                (payload, pn, nxtquic_crypto::Level::Initial)
            }
            PacketType::Handshake => {
                let Ok(parsed) = parse_long_packet(&datagram) else { continue };
                let Some((_, keys)) = handshake_keys.as_ref() else { continue };
                let Ok((pn, payload)) = decrypt_payload(datagram, &parsed, &keys.remote) else { continue };
                largest_handshake = largest_handshake.max(pn);
                (payload, pn, nxtquic_crypto::Level::Handshake)
            }
            PacketType::Short => {
                let Some(keys) = one_rtt_keys.as_ref() else { continue };
                if datagram.len() < 1 + server_scid.len() + 4 + 16 { continue; }
                let parsed = LongPacket { packet_type: PacketType::Short, dcid: server_scid, scid: ConnectionId::from_slice(&[]), pn_offset: 1 + server_scid.len(), packet_end: datagram.len() };
                let Ok((pn, payload)) = decrypt_payload(datagram, &parsed, &keys.remote) else { continue };
                (payload, pn, nxtquic_crypto::Level::OneRtt)
            }
            _ => continue,
        };
        let mut frames = &payload[..];
        while !frames.is_empty() {
            let before = frames.len();
            let Ok(frame) = Frame::decode(&mut frames) else { break };
            match frame {
                Frame::Crypto(crypto) if matches!(level, nxtquic_crypto::Level::Initial | nxtquic_crypto::Level::Handshake) => {
                    let Ok(output) = handshake.process(&crypto.data) else { continue };
                    while let Some((next_level, keys)) = handshake.next_keys() {
                        if next_level == nxtquic_crypto::Level::OneRtt { one_rtt_keys = Some(keys); } else { handshake_keys = Some((next_level, keys)); }
                    }
                    if !output.data.is_empty() {
                        let (packet_type, packet_number, send_keys) = match output.level {
                            nxtquic_crypto::Level::Initial => (PacketType::Initial, 1, Some(&initial_keys.server)),
                            nxtquic_crypto::Level::Handshake => (PacketType::Handshake, largest_handshake + 1, handshake_keys.as_ref().map(|(_, k)| &k.local)),
                            _ => (PacketType::Handshake, largest_handshake + 1, None),
                        };
                        let _ = send_protected_long(&socket, remote_addr, packet_type, initial.scid, server_scid, packet_number, &output.data, send_keys).await;
                    }
                    if handshake.is_complete() {
                        if let Some(keys) = one_rtt_keys.as_ref() {
                            let mut handshake_done = Vec::new();
                            Frame::HandshakeDone.encode(&mut handshake_done);
                            let _ = send_protected_short(&socket, remote_addr, initial.scid, 0, &handshake_done, &keys.local).await;
                        }
                    }
                }
                Frame::Stream(stream) if level == nxtquic_crypto::Level::OneRtt => {
                    let mut encoded = Vec::new();
                    Frame::Stream(stream).encode(&mut encoded);
                    let _ = connection.ingest_frames(&encoded).await;
                }
                _ => {}
            }
            if frames.len() >= before { break; }
        }
        let mut ack = Vec::new();
        Frame::Ack(AckFrame {
            largest_acknowledged: VarInt::from_u64(packet_number).unwrap(),
            ack_delay: VarInt::ZERO,
            first_ack_range: VarInt::ZERO,
            ack_ranges: Vec::new(),
            ecn_counts: None,
        }).encode(&mut ack);
        match level {
            nxtquic_crypto::Level::Initial => {
                let _ = send_protected_long(&socket, remote_addr, PacketType::Initial, initial.scid, server_scid, 2, &ack, Some(&initial_keys.server)).await;
            }
            nxtquic_crypto::Level::Handshake => {
                if let Some((_, keys)) = handshake_keys.as_ref() {
                    let _ = send_protected_long(&socket, remote_addr, PacketType::Handshake, initial.scid, server_scid, largest_handshake + 1, &ack, Some(&keys.local)).await;
                }
            }
            nxtquic_crypto::Level::OneRtt => {
                if let Some(keys) = one_rtt_keys.as_ref() {
                    let _ = send_protected_short(&socket, remote_addr, initial.scid, application_packet_number, &ack, &keys.local).await;
                    application_packet_number += 1;
                }
            }
            _ => {}
        }
        let _ = packet_number;
    }
    connection.close().await;
}

async fn send_protected_long(
    socket: &tokio::net::UdpSocket,
    remote_addr: SocketAddr,
    packet_type: PacketType,
    dcid: ConnectionId,
    scid: ConnectionId,
    packet_number: u64,
    payload: &[u8],
    keys: Option<&DirectionalKeys>,
) -> io::Result<()> {
    let Some(keys) = keys else { return Ok(()); };
    let type_bits = match packet_type { PacketType::Handshake => 0x20, PacketType::Initial => 0x00, _ => return Ok(()) };
    let mut header = vec![0xC0 | type_bits, 0, 0, 0, 1, dcid.len() as u8];
    header.extend_from_slice(dcid.as_bytes());
    header.push(scid.len() as u8);
    header.extend_from_slice(scid.as_bytes());
    if packet_type == PacketType::Initial { VarInt::ZERO.encode(&mut header); }
    let mut body = payload.to_vec();
    while body.len() < 20 {
        body.push(0);
    }
    let packet_len = 1 + body.len() + keys.packet_key.tag_len();
    VarInt::from_u64(packet_len as u64).map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "packet too large"))?.encode(&mut header);
    let pn_offset = header.len();
    header.push(packet_number as u8);
    keys.packet_key.encrypt(packet_number, &header, &mut body).map_err(|error| io::Error::new(io::ErrorKind::Other, error))?;
    header.extend_from_slice(&body);
    let sample: [u8; 16] = header[pn_offset + 4..pn_offset + 20].try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "packet too short for sample"))?;
    let mask = keys.header_key.protection_mask(&sample).map_err(|error| io::Error::new(io::ErrorKind::Other, error))?;
    apply_header_protection(&mut header, &mask, pn_offset, 1);
    socket.send_to(&header, remote_addr).await.map(|_| ())
}

async fn send_protected_short(
    socket: &tokio::net::UdpSocket,
    remote_addr: SocketAddr,
    dcid: ConnectionId,
    packet_number: u64,
    payload: &[u8],
    keys: &DirectionalKeys,
) -> io::Result<()> {
    let mut header = vec![0x40];
    header.extend_from_slice(dcid.as_bytes());
    let pn_offset = header.len();
    header.push(packet_number as u8);
    let mut body = payload.to_vec();
    while body.len() < 20 {
        body.push(0);
    }
    keys.packet_key.encrypt(packet_number, &header, &mut body).map_err(|error| io::Error::new(io::ErrorKind::Other, error))?;
    header.extend_from_slice(&body);
    let sample: [u8; 16] = header[pn_offset + 4..pn_offset + 20].try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "packet too short for sample"))?;
    let mask = keys.header_key.protection_mask(&sample).map_err(|error| io::Error::new(io::ErrorKind::Other, error))?;
    apply_header_protection(&mut header, &mask, pn_offset, 1);
    socket.send_to(&header, remote_addr).await.map(|_| ())
}

#[cfg(test)]
mod packet_tests {
    use super::*;
    use bytes::Bytes;
    use nxtquic_crypto::{CryptoProvider, RustlsCryptoProvider};
    use nxtquic_proto::{
        ConnectionId, Version,
        frame::{CryptoFrame, Frame},
        packet::protection::apply_header_protection,
        varint::VarInt,
    };

    fn protected_client_initial() -> Vec<u8> {
        let dcid = ConnectionId::from_slice(&[0x10; 8]);
        let scid = [0x20; 8];
        let provider = RustlsCryptoProvider::new(None, None);
        let keys = provider.initial_keys(Version::V1, &dcid);

        let mut plaintext = Vec::new();
        Frame::Ping.encode(&mut plaintext);
        Frame::Crypto(CryptoFrame {
            offset: VarInt::ZERO,
            data: Bytes::from_static(b"client hello bytes"),
        })
        .encode(&mut plaintext);

        // The Initial Length field includes the packet number and AEAD tag.
        let ciphertext_len = plaintext.len() + keys.client.packet_key.tag_len();
        let mut header = vec![0xc0, 0, 0, 0, 1, dcid.len() as u8];
        header.extend_from_slice(dcid.as_bytes());
        header.push(scid.len() as u8);
        header.extend_from_slice(&scid);
        header.push(0); // empty token
        VarInt::from_u64((1 + ciphertext_len) as u64)
            .unwrap()
            .encode(&mut header);
        let packet_number_offset = header.len();
        header.push(0); // packet number 0, one byte

        keys.client
            .packet_key
            .encrypt(0, &header, &mut plaintext)
            .unwrap();
        header.extend_from_slice(&plaintext);
        let sample: [u8; 16] = header[packet_number_offset + 4..packet_number_offset + 20]
            .try_into()
            .unwrap();
        let mask = keys.client.header_key.protection_mask(&sample).unwrap();
        apply_header_protection(&mut header, &mask, packet_number_offset, 1);
        header
    }

    #[test]
    fn decrypts_client_initial_and_decodes_crypto_frames() {
        let packet = protected_client_initial();
        let initial = decrypt_client_initial(&packet).unwrap();
        assert_eq!(initial.packet_number, 0);
        assert!(matches!(initial.frames.first(), Some(Frame::Ping)));
        assert!(
            matches!(initial.frames.get(1), Some(Frame::Crypto(frame)) if frame.data == Bytes::from_static(b"client hello bytes"))
        );
    }

    #[test]
    fn rejects_tampered_client_initial() {
        let mut packet = protected_client_initial();
        let last = packet.len() - 1;
        packet[last] ^= 0x01;
        assert!(decrypt_client_initial(&packet).is_err());
    }
}