nxtquic-api 0.1.2

//! QUIC connection management.

use crate::stream::{RecvStream, SendStream, WriteCommand};
use bytes::Bytes;
use std::collections::VecDeque;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::{Mutex, Notify};
use nxtquic_proto::frame::Frame;

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

/// 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: Arc<std::sync::atomic::AtomicBool>,
    socket: Option<Arc<tokio::net::UdpSocket>>,
    remote_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<std::sync::atomic::AtomicU64>,
    next_server_uni: Arc<std::sync::atomic::AtomicU64>,
}

impl Connection {
    pub(crate) 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: Arc::new(std::sync::atomic::AtomicBool::new(false)),
            socket: None,
            remote_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(std::sync::atomic::AtomicU64::new(1)),
            next_server_uni: Arc::new(std::sync::atomic::AtomicU64::new(3)),
        }
    }

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

    pub(crate) fn from_udp(socket: Arc<tokio::net::UdpSocket>, remote_addr: SocketAddr) -> Self {
        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; }
            }
        });
        Self::from_datagrams(socket, remote_addr, rx)
    }

    pub(crate) fn from_datagrams(
        socket: Arc<tokio::net::UdpSocket>,
        remote_addr: SocketAddr,
        rx: tokio::sync::mpsc::UnboundedReceiver<Bytes>,
    ) -> Self {
        let connection = Self::new();
        Self {
            socket: Some(socket),
            remote_addr: Some(remote_addr),
            datagrams: Arc::new(Mutex::new(rx)),
            ..connection
        }
    }

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

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

    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, std::sync::atomic::Ordering::Relaxed);
            let (_, recv) = SendStream::pair();
            return Ok((SendStream::network(tx.clone(), id), recv));
        }
        Ok(SendStream::pair())
    }

    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, std::sync::atomic::Ordering::Relaxed);
            return Ok(SendStream::network(tx.clone(), id));
        }
        let (send, _recv) = SendStream::pair();
        Ok(send)
    }

    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;
        }
    }

    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;
        }
    }

    pub async fn send_datagram(&self, data: bytes::Bytes) -> std::io::Result<()> {
        self.ensure_open()?;
        match (&self.socket, self.remote_addr) {
            (Some(socket), Some(remote_addr)) => {
                socket.send_to(&data, remote_addr).await.map(|_| ())
            }
            _ => Err(std::io::Error::new(
                std::io::ErrorKind::NotConnected,
                "connection has no UDP path",
            )),
        }
    }

    /// Receives the next UDP datagram from this connection's peer.
    pub async fn recv_datagram(&self) -> std::io::Result<Bytes> {
        self.ensure_open()?;
        self.datagrams.lock().await.recv().await.ok_or_else(|| {
            std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "connection receive queue closed")
        })
    }

    /// Decodes a QUIC frame payload and queues DATAGRAM frames for the API.
    /// STREAM and control frames are rejected until their connection-level
    /// state machines are attached, rather than being silently discarded.
    pub(crate) 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::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(rx);
                        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(())
    }

    pub async fn close(&self) {
        self.closed
            .store(true, std::sync::atomic::Ordering::Release);
        self.incoming_bi_notify.notify_waiters();
        self.incoming_uni_notify.notify_waiters();
    }

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