nxtquic_api/
stream.rs

1//! QUIC stream types implementing Tokio async I/O traits.
2
3use std::pin::Pin;
4use std::task::{Context, Poll};
5use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
6use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
7
8pub(crate) enum WriteCommand {
9    Data { stream_id: u64, offset: u64, data: Vec<u8>, fin: bool },
10}
11
12/// A stream that can be written to.
13pub struct SendStream {
14    tx: Option<UnboundedSender<Option<Vec<u8>>>>,
15    network_tx: Option<UnboundedSender<WriteCommand>>,
16    stream_id: u64,
17    offset: u64,
18}
19
20impl SendStream {
21    pub(crate) fn pair() -> (Self, RecvStream) {
22        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
23        (
24            Self { tx: Some(tx), network_tx: None, stream_id: 0, offset: 0 },
25            RecvStream {
26                rx,
27                pending: None,
28                offset: 0,
29            },
30        )
31    }
32
33    pub(crate) fn network(tx: UnboundedSender<WriteCommand>, stream_id: u64) -> Self {
34        Self { tx: None, network_tx: Some(tx), stream_id, offset: 0 }
35    }
36}
37
38impl AsyncWrite for SendStream {
39    fn poll_write(
40        self: Pin<&mut Self>,
41        _cx: &mut Context<'_>,
42        buf: &[u8],
43    ) -> Poll<std::io::Result<usize>> {
44        let this = self.get_mut();
45        if let Some(tx) = this.network_tx.as_ref() {
46            let offset = this.offset;
47            tx.send(WriteCommand::Data { stream_id: this.stream_id, offset, data: buf.to_vec(), fin: false })
48                .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection closed"))?;
49            this.offset += buf.len() as u64;
50            return Poll::Ready(Ok(buf.len()));
51        }
52        let tx = this.tx.as_ref().ok_or_else(|| {
53            std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream is finished")
54        });
55        let tx = match tx {
56            Ok(tx) => tx,
57            Err(err) => return Poll::Ready(Err(err)),
58        };
59        tx.send(Some(buf.to_vec()))
60            .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "peer closed"))?;
61        Poll::Ready(Ok(buf.len()))
62    }
63
64    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
65        Poll::Ready(Ok(()))
66    }
67
68    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
69        let this = self.get_mut();
70        if let Some(tx) = this.network_tx.take() {
71            let _ = tx.send(WriteCommand::Data { stream_id: this.stream_id, offset: this.offset, data: Vec::new(), fin: true });
72            return Poll::Ready(Ok(()));
73        }
74        if let Some(tx) = this.tx.take() {
75            let _ = tx.send(None);
76        }
77        Poll::Ready(Ok(()))
78    }
79}
80
81/// A stream that can be read from.
82pub struct RecvStream {
83    rx: UnboundedReceiver<Option<Vec<u8>>>,
84    pending: Option<Vec<u8>>,
85    offset: usize,
86}
87
88impl RecvStream {
89    pub(crate) fn from_receiver(rx: UnboundedReceiver<Option<Vec<u8>>>) -> Self {
90        Self { rx, pending: None, offset: 0 }
91    }
92}
93
94impl AsyncRead for RecvStream {
95    fn poll_read(
96        mut self: Pin<&mut Self>,
97        cx: &mut Context<'_>,
98        buf: &mut ReadBuf<'_>,
99    ) -> Poll<std::io::Result<()>> {
100        loop {
101            if let Some(data) = self.pending.as_ref() {
102                let remaining = &data[self.offset..];
103                if remaining.is_empty() {
104                    self.pending = None;
105                    self.offset = 0;
106                    continue;
107                }
108                let n = remaining.len().min(buf.remaining());
109                buf.put_slice(&remaining[..n]);
110                self.offset += n;
111                return Poll::Ready(Ok(()));
112            }
113            match Pin::new(&mut self.rx).poll_recv(cx) {
114                Poll::Ready(Some(Some(data))) => {
115                    self.pending = Some(data);
116                }
117                Poll::Ready(Some(None)) | Poll::Ready(None) => return Poll::Ready(Ok(())),
118                Poll::Pending => return Poll::Pending,
119            }
120        }
121    }
122}
nxtquic_api/stream.rs

nxtquic_api/
stream.rs
