nxtquic_api/
connection.rs

1//! QUIC connection management.
2
3use crate::stream::{RecvStream, SendStream, WriteCommand};
4use bytes::Bytes;
5use std::collections::VecDeque;
6use std::collections::HashMap;
7use std::net::SocketAddr;
8use std::sync::Arc;
9use tokio::sync::{Mutex, Notify};
10use nxtquic_proto::frame::Frame;
11
12struct StreamInput {
13    tx: tokio::sync::mpsc::UnboundedSender<Option<Vec<u8>>>,
14    next_offset: u64,
15}
16
17/// A QUIC connection.
18#[derive(Clone)]
19pub struct Connection {
20    incoming_bi: Arc<Mutex<VecDeque<(SendStream, RecvStream)>>>,
21    incoming_uni: Arc<Mutex<VecDeque<RecvStream>>>,
22    incoming_bi_notify: Arc<Notify>,
23    incoming_uni_notify: Arc<Notify>,
24    closed: Arc<std::sync::atomic::AtomicBool>,
25    socket: Option<Arc<tokio::net::UdpSocket>>,
26    remote_addr: Option<SocketAddr>,
27    datagrams: Arc<Mutex<tokio::sync::mpsc::UnboundedReceiver<Bytes>>>,
28    datagram_tx: tokio::sync::mpsc::UnboundedSender<Bytes>,
29    stream_inputs: Arc<Mutex<HashMap<u64, StreamInput>>>,
30    outgoing_tx: Option<tokio::sync::mpsc::UnboundedSender<WriteCommand>>,
31    next_server_bi: Arc<std::sync::atomic::AtomicU64>,
32    next_server_uni: Arc<std::sync::atomic::AtomicU64>,
33}
34
35impl Connection {
36    pub(crate) fn new() -> Self {
37        let (datagram_tx, datagram_rx) = tokio::sync::mpsc::unbounded_channel();
38        Self {
39            incoming_bi: Arc::new(Mutex::new(VecDeque::new())),
40            incoming_uni: Arc::new(Mutex::new(VecDeque::new())),
41            incoming_bi_notify: Arc::new(Notify::new()),
42            incoming_uni_notify: Arc::new(Notify::new()),
43            closed: Arc::new(std::sync::atomic::AtomicBool::new(false)),
44            socket: None,
45            remote_addr: None,
46            datagrams: Arc::new(Mutex::new(datagram_rx)),
47            datagram_tx,
48            stream_inputs: Arc::new(Mutex::new(HashMap::new())),
49            outgoing_tx: None,
50            next_server_bi: Arc::new(std::sync::atomic::AtomicU64::new(1)),
51            next_server_uni: Arc::new(std::sync::atomic::AtomicU64::new(3)),
52        }
53    }
54
55    pub(crate) fn with_outgoing(tx: tokio::sync::mpsc::UnboundedSender<WriteCommand>) -> Self {
56        let mut connection = Self::new();
57        connection.outgoing_tx = Some(tx);
58        connection
59    }
60
61    pub(crate) fn from_udp(socket: Arc<tokio::net::UdpSocket>, remote_addr: SocketAddr) -> Self {
62        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
63        let receive_socket = Arc::clone(&socket);
64        let receive_tx = tx.clone();
65        tokio::spawn(async move {
66            let mut packet = vec![0_u8; 65_535];
67            while let Ok((length, peer)) = receive_socket.recv_from(&mut packet).await {
68                if peer == remote_addr && receive_tx.send(Bytes::copy_from_slice(&packet[..length])).is_err() { break; }
69            }
70        });
71        Self::from_datagrams(socket, remote_addr, rx)
72    }
73
74    pub(crate) fn from_datagrams(
75        socket: Arc<tokio::net::UdpSocket>,
76        remote_addr: SocketAddr,
77        rx: tokio::sync::mpsc::UnboundedReceiver<Bytes>,
78    ) -> Self {
79        let connection = Self::new();
80        Self {
81            socket: Some(socket),
82            remote_addr: Some(remote_addr),
83            datagrams: Arc::new(Mutex::new(rx)),
84            ..connection
85        }
86    }
87
88    /// Returns the peer address for a network-backed connection.
89    pub fn remote_address(&self) -> Option<SocketAddr> {
90        self.remote_addr
91    }
92
93    /// Returns whether the connection has been closed locally.
94    pub fn is_closed(&self) -> bool {
95        self.closed.load(std::sync::atomic::Ordering::Acquire)
96    }
97
98    pub async fn open_bi(&self) -> std::io::Result<(SendStream, RecvStream)> {
99        self.ensure_open()?;
100        if let Some(tx) = self.outgoing_tx.as_ref() {
101            let id = self.next_server_bi.fetch_add(4, std::sync::atomic::Ordering::Relaxed);
102            let (_, recv) = SendStream::pair();
103            return Ok((SendStream::network(tx.clone(), id), recv));
104        }
105        Ok(SendStream::pair())
106    }
107
108    pub async fn open_uni(&self) -> std::io::Result<SendStream> {
109        self.ensure_open()?;
110        if let Some(tx) = self.outgoing_tx.as_ref() {
111            let id = self.next_server_uni.fetch_add(4, std::sync::atomic::Ordering::Relaxed);
112            return Ok(SendStream::network(tx.clone(), id));
113        }
114        let (send, _recv) = SendStream::pair();
115        Ok(send)
116    }
117
118    pub async fn accept_bi(&self) -> std::io::Result<(SendStream, RecvStream)> {
119        loop {
120            self.ensure_open()?;
121            if let Some(stream) = self.incoming_bi.lock().await.pop_front() {
122                return Ok(stream);
123            }
124            self.incoming_bi_notify.notified().await;
125        }
126    }
127
128    pub async fn accept_uni(&self) -> std::io::Result<RecvStream> {
129        loop {
130            self.ensure_open()?;
131            if let Some(stream) = self.incoming_uni.lock().await.pop_front() {
132                return Ok(stream);
133            }
134            self.incoming_uni_notify.notified().await;
135        }
136    }
137
138    pub async fn send_datagram(&self, data: bytes::Bytes) -> std::io::Result<()> {
139        self.ensure_open()?;
140        match (&self.socket, self.remote_addr) {
141            (Some(socket), Some(remote_addr)) => {
142                socket.send_to(&data, remote_addr).await.map(|_| ())
143            }
144            _ => Err(std::io::Error::new(
145                std::io::ErrorKind::NotConnected,
146                "connection has no UDP path",
147            )),
148        }
149    }
150
151    /// Receives the next UDP datagram from this connection's peer.
152    pub async fn recv_datagram(&self) -> std::io::Result<Bytes> {
153        self.ensure_open()?;
154        self.datagrams.lock().await.recv().await.ok_or_else(|| {
155            std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "connection receive queue closed")
156        })
157    }
158
159    /// Decodes a QUIC frame payload and queues DATAGRAM frames for the API.
160    /// STREAM and control frames are rejected until their connection-level
161    /// state machines are attached, rather than being silently discarded.
162    pub(crate) async fn ingest_frames(&self, mut payload: &[u8]) -> std::io::Result<()> {
163        while !payload.is_empty() {
164            let before = payload.len();
165            let frame = Frame::decode(&mut payload)
166                .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?;
167            match frame {
168                Frame::Datagram(frame) => self.datagram_tx.send(frame.data).map_err(|_| {
169                    std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection receive queue closed")
170                })?,
171                Frame::Padding | Frame::Ping | Frame::Ack(_) => {}
172                Frame::Stream(frame) => {
173                    let stream_id = frame.stream_id.into_inner().into_inner();
174                    let mut streams = self.stream_inputs.lock().await;
175                    if let Some(input) = streams.get_mut(&stream_id) {
176                        if frame.offset.into_inner() != input.next_offset { return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "out-of-order STREAM frame")); }
177                        input.next_offset += frame.data.len() as u64;
178                        input.tx.send(Some(frame.data.to_vec())).map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed"))?;
179                        if frame.fin { let _ = input.tx.send(None); }
180                    } else {
181                        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
182                        tx.send(Some(frame.data.to_vec())).map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed"))?;
183                        let recv = RecvStream::from_receiver(rx);
184                        let uni = stream_id & 0b10 != 0;
185                        if uni {
186                            self.incoming_uni.lock().await.push_back(recv);
187                            self.incoming_uni_notify.notify_one();
188                        } else {
189                            let send = if let Some(tx) = self.outgoing_tx.as_ref() {
190                                SendStream::network(tx.clone(), stream_id)
191                            } else {
192                                SendStream::pair().0
193                            };
194                            self.incoming_bi.lock().await.push_back((send, recv));
195                            self.incoming_bi_notify.notify_one();
196                        }
197                        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); }
198                    }
199                }
200                _ => return Err(std::io::Error::new(std::io::ErrorKind::Unsupported, "QUIC control frame handling is not enabled")),
201            }
202            if payload.len() == before { return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "frame decoder made no progress")); }
203        }
204        Ok(())
205    }
206
207    pub async fn close(&self) {
208        self.closed
209            .store(true, std::sync::atomic::Ordering::Release);
210        self.incoming_bi_notify.notify_waiters();
211        self.incoming_uni_notify.notify_waiters();
212    }
213
214    fn ensure_open(&self) -> std::io::Result<()> {
215        if self.closed.load(std::sync::atomic::Ordering::Acquire) {
216            Err(std::io::Error::new(
217                std::io::ErrorKind::NotConnected,
218                "connection is closed",
219            ))
220        } else {
221            Ok(())
222        }
223    }
224}
nxtquic_api/connection.rs

nxtquic_api/
connection.rs
