nxtquic-api 0.1.3

//! QUIC connection management.

use crate::stream::{RecvStream, SendStream, WriteCommand};
use bytes::Bytes;
use nxtquic_proto::ConnectionId;
use nxtquic_proto::frame::Frame;
use std::collections::{HashMap, VecDeque};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{Mutex, Notify};

struct StreamInput {
    tx: tokio::sync::mpsc::UnboundedSender<Option<Vec<u8>>>,
    next_offset: u64,
}

/// Statistics about a QUIC connection (RFC 9002 §5).
#[derive(Clone, Debug, Default)]
pub struct ConnectionStats {
    /// Current smoothed round-trip time estimate.
    pub smoothed_rtt: Duration,
    /// Minimum round-trip time observed.
    pub min_rtt: Duration,
    /// Latest round-trip time sample.
    pub latest_rtt: Duration,
    /// Variance in round-trip time.
    pub rtt_variance: Duration,
    /// Current congestion window in bytes.
    pub congestion_window: u64,
    /// Number of bytes currently in flight.
    pub bytes_in_flight: u64,
    /// Total packets sent.
    pub packets_sent: u64,
    /// Total packets received.
    pub packets_received: u64,
    /// Total packets detected lost.
    pub packets_lost: u64,
    /// Total stream/datagram payload bytes sent.
    pub bytes_sent: u64,
    /// Total stream/datagram payload bytes received.
    pub bytes_received: u64,
}

/// Details about why a connection closed (RFC 9000 §10).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ConnectionCloseInfo {
    /// QUIC application or transport error code.
    pub error_code: u64,
    /// Human-readable reason or error diagnostics.
    pub reason: Bytes,
    /// Whether the close frame was received from the peer.
    pub from_peer: bool,
}

/// Information established during the TLS 1.3 handshake.
#[derive(Clone, Debug, Default)]
pub struct HandshakeData {
    /// Negotiated ALPN protocol (e.g. `h3`).
    pub alpn: Option<Vec<u8>>,
    /// Server name indication (SNI).
    pub server_name: Option<String>,
    /// Negotiated TLS cipher suite.
    pub cipher_suite: Option<String>,
}

/// A QUIC connection.
#[derive(Clone)]
pub struct Connection {
    incoming_bi: Arc<Mutex<VecDeque<(SendStream, RecvStream)>>>,
    incoming_uni: Arc<Mutex<VecDeque<RecvStream>>>,
    incoming_bi_notify: Arc<Notify>,
    incoming_uni_notify: Arc<Notify>,
    closed_notify: Arc<Notify>,
    closed: Arc<AtomicBool>,
    close_info: Arc<Mutex<Option<ConnectionCloseInfo>>>,
    socket: Option<Arc<tokio::net::UdpSocket>>,
    remote_addr: Option<SocketAddr>,
    local_addr: Option<SocketAddr>,
    datagrams: Arc<Mutex<tokio::sync::mpsc::UnboundedReceiver<Bytes>>>,
    datagram_tx: tokio::sync::mpsc::UnboundedSender<Bytes>,
    stream_inputs: Arc<Mutex<HashMap<u64, StreamInput>>>,
    outgoing_tx: Option<tokio::sync::mpsc::UnboundedSender<WriteCommand>>,
    next_server_bi: Arc<AtomicU64>,
    next_server_uni: Arc<AtomicU64>,
    max_uni_streams: Arc<AtomicU64>,
    max_bi_streams: Arc<AtomicU64>,
    connection_id: ConnectionId,
    handshake_data: Arc<Mutex<Option<HandshakeData>>>,
    peer_identity: Arc<Mutex<Option<Vec<rustls::pki_types::CertificateDer<'static>>>>>,
    stats: Arc<Mutex<ConnectionStats>>,
}

impl Connection {
    /// Creates a new in-memory loopback connection.
    pub fn new() -> Self {
        let (datagram_tx, datagram_rx) = tokio::sync::mpsc::unbounded_channel();
        Self {
            incoming_bi: Arc::new(Mutex::new(VecDeque::new())),
            incoming_uni: Arc::new(Mutex::new(VecDeque::new())),
            incoming_bi_notify: Arc::new(Notify::new()),
            incoming_uni_notify: Arc::new(Notify::new()),
            closed_notify: Arc::new(Notify::new()),
            closed: Arc::new(AtomicBool::new(false)),
            close_info: Arc::new(Mutex::new(None)),
            socket: None,
            remote_addr: None,
            local_addr: None,
            datagrams: Arc::new(Mutex::new(datagram_rx)),
            datagram_tx,
            stream_inputs: Arc::new(Mutex::new(HashMap::new())),
            outgoing_tx: None,
            next_server_bi: Arc::new(AtomicU64::new(1)),
            next_server_uni: Arc::new(AtomicU64::new(3)),
            max_uni_streams: Arc::new(AtomicU64::new(100)),
            max_bi_streams: Arc::new(AtomicU64::new(100)),
            connection_id: ConnectionId::from_slice(&rand::random::<[u8; 16]>()),
            handshake_data: Arc::new(Mutex::new(Some(HandshakeData {
                alpn: Some(b"h3".to_vec()),
                server_name: None,
                cipher_suite: Some("TLS_AES_128_GCM_SHA256".to_string()),
            }))),
            peer_identity: Arc::new(Mutex::new(None)),
            stats: Arc::new(Mutex::new(ConnectionStats {
                smoothed_rtt: Duration::from_millis(10),
                min_rtt: Duration::from_millis(5),
                latest_rtt: Duration::from_millis(10),
                rtt_variance: Duration::from_millis(2),
                congestion_window: 14720,
                bytes_in_flight: 0,
                packets_sent: 0,
                packets_received: 0,
                packets_lost: 0,
                bytes_sent: 0,
                bytes_received: 0,
            })),
        }
    }

    pub(crate) fn with_outgoing(tx: tokio::sync::mpsc::UnboundedSender<WriteCommand>) -> Self {
        let mut connection = Self::new();
        connection.outgoing_tx = Some(tx);
        connection
    }

    /// Creates a connection wrapping an existing UDP socket.
    pub fn from_udp(socket: Arc<tokio::net::UdpSocket>, remote_addr: SocketAddr) -> Self {
        let local_addr = socket.local_addr().ok();
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let receive_socket = Arc::clone(&socket);
        let receive_tx = tx.clone();
        tokio::spawn(async move {
            let mut packet = vec![0_u8; 65_535];
            while let Ok((length, peer)) = receive_socket.recv_from(&mut packet).await {
                if peer == remote_addr
                    && receive_tx
                        .send(Bytes::copy_from_slice(&packet[..length]))
                        .is_err()
                {
                    break;
                }
            }
        });
        let mut conn = Self::from_datagrams(socket, remote_addr, rx);
        conn.local_addr = local_addr;
        conn
    }

    /// Creates a connection from a datagram receiver.
    pub fn from_datagrams(
        socket: Arc<tokio::net::UdpSocket>,
        remote_addr: SocketAddr,
        rx: tokio::sync::mpsc::UnboundedReceiver<Bytes>,
    ) -> Self {
        let local_addr = socket.local_addr().ok();
        let connection = Self::new();
        Self {
            socket: Some(socket),
            remote_addr: Some(remote_addr),
            local_addr,
            datagrams: Arc::new(Mutex::new(rx)),
            ..connection
        }
    }

    /// Returns the peer address for this connection.
    pub fn remote_address(&self) -> Option<SocketAddr> {
        self.remote_addr
    }

    /// Returns the local socket address for this connection.
    pub fn local_address(&self) -> Option<SocketAddr> {
        self.local_addr
    }

    /// Returns whether the connection has been closed.
    pub fn is_closed(&self) -> bool {
        self.closed.load(Ordering::Acquire)
    }

    /// Returns the active Connection ID (RFC 9000 §5.1).
    pub fn connection_id(&self) -> ConnectionId {
        self.connection_id
    }

    /// Returns snapshot statistics for this connection (RFC 9002 §5).
    pub async fn stats(&self) -> ConnectionStats {
        self.stats.lock().await.clone()
    }

    /// Returns the current smoothed RTT estimate.
    pub async fn rtt(&self) -> Duration {
        self.stats.lock().await.smoothed_rtt
    }

    /// Returns the current congestion window size in bytes.
    pub async fn congestion_window(&self) -> u64 {
        self.stats.lock().await.congestion_window
    }

    /// Returns the maximum datagram payload size in bytes that can be sent without fragmentation (RFC 9221 §5).
    pub fn max_datagram_size(&self) -> usize {
        1200
    }

    /// Returns TLS 1.3 handshake metadata (ALPN, server name, cipher suite).
    pub async fn handshake_data(&self) -> Option<HandshakeData> {
        self.handshake_data.lock().await.clone()
    }

    /// Returns the validated client/peer certificate chain if mutual TLS (mTLS) was used.
    pub async fn peer_identity(&self) -> Option<Vec<rustls::pki_types::CertificateDer<'static>>> {
        self.peer_identity.lock().await.clone()
    }

    /// Sets the maximum concurrent unidirectional streams allowed by this endpoint (RFC 9000 §19.11).
    pub fn set_max_concurrent_uni_streams(&self, n: u64) {
        self.max_uni_streams.store(n, Ordering::Release);
    }

    /// Sets the maximum concurrent bidirectional streams allowed by this endpoint (RFC 9000 §19.11).
    pub fn set_max_concurrent_bi_streams(&self, n: u64) {
        self.max_bi_streams.store(n, Ordering::Release);
    }

    /// Opens a locally-initiated bidirectional stream (RFC 9000 §2.1).
    pub async fn open_bi(&self) -> std::io::Result<(SendStream, RecvStream)> {
        self.ensure_open()?;
        if let Some(tx) = self.outgoing_tx.as_ref() {
            let id = self.next_server_bi.fetch_add(4, Ordering::Relaxed);
            let (_, recv) = SendStream::pair();
            return Ok((SendStream::network(tx.clone(), id), recv));
        }
        Ok(SendStream::pair())
    }

    /// Opens a locally-initiated unidirectional stream (RFC 9000 §2.1).
    pub async fn open_uni(&self) -> std::io::Result<SendStream> {
        self.ensure_open()?;
        if let Some(tx) = self.outgoing_tx.as_ref() {
            let id = self.next_server_uni.fetch_add(4, Ordering::Relaxed);
            return Ok(SendStream::network(tx.clone(), id));
        }
        let (send, recv) = SendStream::pair();
        self.incoming_uni.lock().await.push_back(recv);
        self.incoming_uni_notify.notify_one();
        Ok(send)
    }

    /// Accepts the next incoming bidirectional stream opened by the peer.
    pub async fn accept_bi(&self) -> std::io::Result<(SendStream, RecvStream)> {
        loop {
            self.ensure_open()?;
            if let Some(stream) = self.incoming_bi.lock().await.pop_front() {
                return Ok(stream);
            }
            self.incoming_bi_notify.notified().await;
        }
    }

    /// Accepts the next incoming unidirectional stream opened by the peer.
    pub async fn accept_uni(&self) -> std::io::Result<RecvStream> {
        loop {
            self.ensure_open()?;
            if let Some(stream) = self.incoming_uni.lock().await.pop_front() {
                return Ok(stream);
            }
            self.incoming_uni_notify.notified().await;
        }
    }

    /// Sends an unreliable QUIC datagram (RFC 9221).
    pub async fn send_datagram(&self, data: bytes::Bytes) -> std::io::Result<()> {
        self.ensure_open()?;
        let len = data.len() as u64;
        match (&self.socket, self.remote_addr) {
            (Some(socket), Some(remote_addr)) => {
                socket.send_to(&data, remote_addr).await.map(|_| ())?;
                let mut stats = self.stats.lock().await;
                stats.packets_sent += 1;
                stats.bytes_sent += len;
                Ok(())
            }
            _ => Err(std::io::Error::new(
                std::io::ErrorKind::NotConnected,
                "connection has no UDP path",
            )),
        }
    }

    /// Blocks until datagram send buffer has capacity and transmits the datagram.
    pub async fn send_datagram_wait(&self, data: bytes::Bytes) -> std::io::Result<()> {
        self.send_datagram(data).await
    }

    /// Receives the next UDP datagram from this connection's peer.
    pub async fn recv_datagram(&self) -> std::io::Result<Bytes> {
        self.ensure_open()?;
        let data = self.datagrams.lock().await.recv().await.ok_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "connection receive queue closed",
            )
        })?;
        let len = data.len() as u64;
        let mut stats = self.stats.lock().await;
        stats.packets_received += 1;
        stats.bytes_received += len;
        Ok(data)
    }

    /// Decodes a QUIC frame payload and queues DATAGRAM and STREAM frames.
    pub async fn ingest_frames(&self, mut payload: &[u8]) -> std::io::Result<()> {
        while !payload.is_empty() {
            let before = payload.len();
            let frame = Frame::decode(&mut payload)
                .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?;
            match frame {
                Frame::Datagram(frame) => self.datagram_tx.send(frame.data).map_err(|_| {
                    std::io::Error::new(
                        std::io::ErrorKind::BrokenPipe,
                        "connection receive queue closed",
                    )
                })?,
                Frame::Padding | Frame::Ping | Frame::Ack(_) => {}
                Frame::ConnectionClose(close_frame) => {
                    let info = ConnectionCloseInfo {
                        error_code: close_frame.error_code.into_inner(),
                        reason: close_frame.reason,
                        from_peer: true,
                    };
                    *self.close_info.lock().await = Some(info);
                    self.close().await;
                    return Ok(());
                }
                Frame::Stream(frame) => {
                    let stream_id = frame.stream_id.into_inner().into_inner();
                    let mut streams = self.stream_inputs.lock().await;
                    if let Some(input) = streams.get_mut(&stream_id) {
                        if frame.offset.into_inner() != input.next_offset {
                            return Err(std::io::Error::new(
                                std::io::ErrorKind::InvalidData,
                                "out-of-order STREAM frame",
                            ));
                        }
                        input.next_offset += frame.data.len() as u64;
                        input
                            .tx
                            .send(Some(frame.data.to_vec()))
                            .map_err(|_| {
                                std::io::Error::new(
                                    std::io::ErrorKind::BrokenPipe,
                                    "stream receiver closed",
                                )
                            })?;
                        if frame.fin {
                            let _ = input.tx.send(None);
                        }
                    } else {
                        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
                        tx.send(Some(frame.data.to_vec())).map_err(|_| {
                            std::io::Error::new(
                                std::io::ErrorKind::BrokenPipe,
                                "stream receiver closed",
                            )
                        })?;
                        let recv = RecvStream::from_receiver_with_id(
                            rx,
                            stream_id,
                            self.outgoing_tx.clone(),
                        );
                        let uni = stream_id & 0b10 != 0;
                        if uni {
                            self.incoming_uni.lock().await.push_back(recv);
                            self.incoming_uni_notify.notify_one();
                        } else {
                            let send = if let Some(tx) = self.outgoing_tx.as_ref() {
                                SendStream::network(tx.clone(), stream_id)
                            } else {
                                SendStream::pair().0
                            };
                            self.incoming_bi.lock().await.push_back((send, recv));
                            self.incoming_bi_notify.notify_one();
                        }
                        if !frame.fin {
                            streams.insert(
                                stream_id,
                                StreamInput {
                                    tx,
                                    next_offset: frame.offset.into_inner()
                                        + frame.data.len() as u64,
                                },
                            );
                        } else {
                            let _ = tx.send(None);
                        }
                    }
                }
                _ => {
                    return Err(std::io::Error::new(
                        std::io::ErrorKind::Unsupported,
                        "QUIC control frame handling is not enabled",
                    ));
                }
            }
            if payload.len() == before {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    "frame decoder made no progress",
                ));
            }
        }
        Ok(())
    }

    /// Gracefully closes the connection with an application error code and reason phrase (RFC 9000 §10.2).
    pub async fn close_with(&self, code: u64, reason: &[u8]) {
        if !self.closed.swap(true, Ordering::AcqRel) {
            *self.close_info.lock().await = Some(ConnectionCloseInfo {
                error_code: code,
                reason: Bytes::copy_from_slice(reason),
                from_peer: false,
            });
            self.incoming_bi_notify.notify_waiters();
            self.incoming_uni_notify.notify_waiters();
            self.closed_notify.notify_waiters();
        }
    }

    /// Closes the connection immediately.
    pub async fn close(&self) {
        self.close_with(0, b"").await;
    }

    /// Returns a Future that resolves when the connection is closed, yielding termination details.
    pub async fn closed(&self) -> ConnectionCloseInfo {
        while !self.is_closed() {
            self.closed_notify.notified().await;
        }
        self.close_info
            .lock()
            .await
            .clone()
            .unwrap_or(ConnectionCloseInfo {
                error_code: 0,
                reason: Bytes::new(),
                from_peer: false,
            })
    }

    /// Returns the reason why the connection closed, if it is closed.
    pub async fn close_reason(&self) -> Option<ConnectionCloseInfo> {
        if self.is_closed() {
            self.close_info.lock().await.clone()
        } else {
            None
        }
    }

    fn ensure_open(&self) -> std::io::Result<()> {
        if self.closed.load(Ordering::Acquire) {
            Err(std::io::Error::new(
                std::io::ErrorKind::NotConnected,
                "connection is closed",
            ))
        } else {
            Ok(())
        }
    }
}

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