nxtquic_api/
stream.rs

1//! QUIC stream types implementing Tokio async I/O traits.
2
3use std::pin::Pin;
4use std::sync::atomic::{AtomicI32, AtomicU64, Ordering};
5use std::sync::Arc;
6use std::task::{Context, Poll};
7use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
8use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
9
10pub(crate) enum WriteCommand {
11    Data {
12        stream_id: u64,
13        offset: u64,
14        data: Vec<u8>,
15        fin: bool,
16    },
17    Reset {
18        stream_id: u64,
19        error_code: u64,
20    },
21    StopSending {
22        stream_id: u64,
23        error_code: u64,
24    },
25    Priority {
26        stream_id: u64,
27        priority: i32,
28    },
29}
30
31/// A stream that can be written to.
32pub struct SendStream {
33    tx: Option<UnboundedSender<Option<Vec<u8>>>>,
34    network_tx: Option<UnboundedSender<WriteCommand>>,
35    stream_id: u64,
36    offset: u64,
37    priority: Arc<AtomicI32>,
38    stopped_reason: Arc<AtomicU64>,
39}
40
41impl SendStream {
42    pub(crate) fn pair() -> (Self, RecvStream) {
43        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
44        (
45            Self {
46                tx: Some(tx),
47                network_tx: None,
48                stream_id: 0,
49                offset: 0,
50                priority: Arc::new(AtomicI32::new(0)),
51                stopped_reason: Arc::new(AtomicU64::new(u64::MAX)),
52            },
53            RecvStream {
54                rx,
55                pending: None,
56                offset: 0,
57                stream_id: 0,
58                priority: Arc::new(AtomicI32::new(0)),
59                reset_reason: Arc::new(AtomicU64::new(u64::MAX)),
60                network_tx: None,
61            },
62        )
63    }
64
65    pub(crate) fn network(tx: UnboundedSender<WriteCommand>, stream_id: u64) -> Self {
66        Self {
67            tx: None,
68            network_tx: Some(tx),
69            stream_id,
70            offset: 0,
71            priority: Arc::new(AtomicI32::new(0)),
72            stopped_reason: Arc::new(AtomicU64::new(u64::MAX)),
73        }
74    }
75
76    /// Returns the unique QUIC Stream ID (RFC 9000 §2.1).
77    pub fn stream_id(&self) -> u64 {
78        self.stream_id
79    }
80
81    /// Sets the stream priority for scheduling (RFC 9000 §2.3).
82    pub fn set_priority(&mut self, priority: i32) {
83        self.priority.store(priority, Ordering::Release);
84        if let Some(tx) = &self.network_tx {
85            let _ = tx.send(WriteCommand::Priority {
86                stream_id: self.stream_id,
87                priority,
88            });
89        }
90    }
91
92    /// Gets the current stream priority.
93    pub fn priority(&self) -> i32 {
94        self.priority.load(Ordering::Acquire)
95    }
96
97    /// Explicitly finishes the stream by sending a QUIC FIN (RFC 9000 §19.8).
98    pub async fn finish(&mut self) -> std::io::Result<()> {
99        self.shutdown().await
100    }
101
102    /// Abruptly terminates sending on this stream with an application error code (RFC 9000 §19.4).
103    pub async fn reset(&mut self, error_code: u64) -> std::io::Result<()> {
104        if let Some(tx) = self.network_tx.take() {
105            let _ = tx.send(WriteCommand::Reset {
106                stream_id: self.stream_id,
107                error_code,
108            });
109        }
110        if let Some(tx) = self.tx.take() {
111            let _ = tx.send(None);
112        }
113        Ok(())
114    }
115
116    /// Returns the error code if the peer sent a `STOP_SENDING` frame (RFC 9000 §19.5).
117    pub fn stopped(&self) -> Option<u64> {
118        let code = self.stopped_reason.load(Ordering::Acquire);
119        if code == u64::MAX {
120            None
121        } else {
122            Some(code)
123        }
124    }
125
126    /// Scatter-gather writes multiple contiguous byte chunks to the stream.
127    pub async fn write_chunks(&mut self, bufs: &[bytes::Bytes]) -> std::io::Result<usize> {
128        let mut total = 0;
129        for buf in bufs {
130            self.write_all(buf).await?;
131            total += buf.len();
132        }
133        Ok(total)
134    }
135}
136
137impl AsyncWrite for SendStream {
138    fn poll_write(
139        self: Pin<&mut Self>,
140        _cx: &mut Context<'_>,
141        buf: &[u8],
142    ) -> Poll<std::io::Result<usize>> {
143        let this = self.get_mut();
144        if let Some(tx) = this.network_tx.as_ref() {
145            let offset = this.offset;
146            tx.send(WriteCommand::Data {
147                stream_id: this.stream_id,
148                offset,
149                data: buf.to_vec(),
150                fin: false,
151            })
152            .map_err(|_| {
153                std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection closed")
154            })?;
155            this.offset += buf.len() as u64;
156            return Poll::Ready(Ok(buf.len()));
157        }
158        let tx = this.tx.as_ref().ok_or_else(|| {
159            std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream is finished")
160        });
161        let tx = match tx {
162            Ok(tx) => tx,
163            Err(err) => return Poll::Ready(Err(err)),
164        };
165        tx.send(Some(buf.to_vec()))
166            .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "peer closed"))?;
167        Poll::Ready(Ok(buf.len()))
168    }
169
170    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
171        Poll::Ready(Ok(()))
172    }
173
174    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
175        let this = self.get_mut();
176        if let Some(tx) = this.network_tx.take() {
177            let _ = tx.send(WriteCommand::Data {
178                stream_id: this.stream_id,
179                offset: this.offset,
180                data: Vec::new(),
181                fin: true,
182            });
183            return Poll::Ready(Ok(()));
184        }
185        if let Some(tx) = this.tx.take() {
186            let _ = tx.send(None);
187        }
188        Poll::Ready(Ok(()))
189    }
190}
191
192/// A stream that can be read from.
193pub struct RecvStream {
194    rx: UnboundedReceiver<Option<Vec<u8>>>,
195    pending: Option<Vec<u8>>,
196    offset: usize,
197    stream_id: u64,
198    priority: Arc<AtomicI32>,
199    reset_reason: Arc<AtomicU64>,
200    network_tx: Option<UnboundedSender<WriteCommand>>,
201}
202
203impl RecvStream {
204    pub(crate) fn from_receiver(rx: UnboundedReceiver<Option<Vec<u8>>>) -> Self {
205        Self {
206            rx,
207            pending: None,
208            offset: 0,
209            stream_id: 0,
210            priority: Arc::new(AtomicI32::new(0)),
211            reset_reason: Arc::new(AtomicU64::new(u64::MAX)),
212            network_tx: None,
213        }
214    }
215
216    pub(crate) fn from_receiver_with_id(
217        rx: UnboundedReceiver<Option<Vec<u8>>>,
218        stream_id: u64,
219        network_tx: Option<UnboundedSender<WriteCommand>>,
220    ) -> Self {
221        Self {
222            rx,
223            pending: None,
224            offset: 0,
225            stream_id,
226            priority: Arc::new(AtomicI32::new(0)),
227            reset_reason: Arc::new(AtomicU64::new(u64::MAX)),
228            network_tx,
229        }
230    }
231
232    /// Returns the unique QUIC Stream ID (RFC 9000 §2.1).
233    pub fn stream_id(&self) -> u64 {
234        self.stream_id
235    }
236
237    /// Sets the stream priority for scheduling (RFC 9000 §2.3).
238    pub fn set_priority(&mut self, priority: i32) {
239        self.priority.store(priority, Ordering::Release);
240        if let Some(tx) = &self.network_tx {
241            let _ = tx.send(WriteCommand::Priority {
242                stream_id: self.stream_id,
243                priority,
244            });
245        }
246    }
247
248    /// Gets the current stream priority.
249    pub fn priority(&self) -> i32 {
250        self.priority.load(Ordering::Acquire)
251    }
252
253    /// Signals the peer to stop sending on this stream with an application error code (RFC 9000 §19.5).
254    pub async fn stop_sending(&mut self, error_code: u64) -> std::io::Result<()> {
255        if let Some(tx) = &self.network_tx {
256            let _ = tx.send(WriteCommand::StopSending {
257                stream_id: self.stream_id,
258                error_code,
259            });
260        }
261        Ok(())
262    }
263
264    /// Returns the error code if the peer aborted this stream with a `RESET_STREAM` frame (RFC 9000 §19.4).
265    pub fn received_reset(&self) -> Option<u64> {
266        let code = self.reset_reason.load(Ordering::Acquire);
267        if code == u64::MAX {
268            None
269        } else {
270            Some(code)
271        }
272    }
273
274    /// Reads an individual contiguous chunk of data without copying into an intermediate buffer.
275    pub async fn read_chunk(&mut self, max: usize) -> std::io::Result<Option<bytes::Bytes>> {
276        if let Some(data) = self.pending.take() {
277            let remaining = &data[self.offset..];
278            if !remaining.is_empty() {
279                let chunk_size = remaining.len().min(max);
280                let chunk = bytes::Bytes::copy_from_slice(&remaining[..chunk_size]);
281                if chunk_size < remaining.len() {
282                    self.pending = Some(data);
283                    self.offset += chunk_size;
284                } else {
285                    self.offset = 0;
286                }
287                return Ok(Some(chunk));
288            }
289        }
290
291        match self.rx.recv().await {
292            Some(Some(data)) => {
293                let chunk_size = data.len().min(max);
294                let chunk = bytes::Bytes::copy_from_slice(&data[..chunk_size]);
295                if chunk_size < data.len() {
296                    self.pending = Some(data);
297                    self.offset = chunk_size;
298                }
299                Ok(Some(chunk))
300            }
301            Some(None) | None => Ok(None),
302        }
303    }
304
305    /// Reads into multiple buffers using vectored I/O.
306    pub async fn read_chunks(&mut self, bufs: &mut [bytes::Bytes]) -> std::io::Result<usize> {
307        let mut count = 0;
308        for slot in bufs.iter_mut() {
309            if let Some(chunk) = self.read_chunk(65536).await? {
310                *slot = chunk;
311                count += 1;
312            } else {
313                break;
314            }
315        }
316        Ok(count)
317    }
318}
319
320impl AsyncRead for RecvStream {
321    fn poll_read(
322        mut self: Pin<&mut Self>,
323        cx: &mut Context<'_>,
324        buf: &mut ReadBuf<'_>,
325    ) -> Poll<std::io::Result<()>> {
326        loop {
327            if let Some(data) = self.pending.as_ref() {
328                let remaining = &data[self.offset..];
329                if remaining.is_empty() {
330                    self.pending = None;
331                    self.offset = 0;
332                    continue;
333                }
334                let n = remaining.len().min(buf.remaining());
335                buf.put_slice(&remaining[..n]);
336                self.offset += n;
337                return Poll::Ready(Ok(()));
338            }
339            match Pin::new(&mut self.rx).poll_recv(cx) {
340                Poll::Ready(Some(Some(data))) => {
341                    self.pending = Some(data);
342                }
343                Poll::Ready(Some(None)) | Poll::Ready(None) => return Poll::Ready(Ok(())),
344                Poll::Pending => return Poll::Pending,
345            }
346        }
347    }
348}
nxtquic_api/stream.rs

nxtquic_api/
stream.rs
