compio_quic/
recv_stream.rs

1use std::{
2    io,
3    mem::MaybeUninit,
4    task::{Context, Poll},
5};
6
7use compio_buf::{BufResult, IntoInner, IoBufExt, IoBufMut, IoBufMutExt, SetLenExt, bytes::Bytes};
8use compio_io::AsyncRead;
9use futures_util::future::poll_fn;
10use quinn_proto::{Chunk, Chunks, ClosedStream, ReadableError, StreamId, VarInt};
11use thiserror::Error;
12
13use crate::{ConnectionError, ConnectionInner, sync::shared::Shared};
14
15/// A stream that can only be used to receive data
16///
17/// `stop(0)` is implicitly called on drop unless:
18/// - A variant of [`ReadError`] has been yielded by a read call
19/// - [`stop()`] was called explicitly
20///
21/// # Cancellation
22///
23/// A `read` method is said to be *cancel-safe* when dropping its future before
24/// the future becomes ready cannot lead to loss of stream data. This is true of
25/// methods which succeed immediately when any progress is made, and is not true
26/// of methods which might need to perform multiple reads internally before
27/// succeeding. Each `read` method documents whether it is cancel-safe.
28///
29/// # Common issues
30///
31/// ## Data never received on a locally-opened stream
32///
33/// Peers are not notified of streams until they or a later-numbered stream are
34/// used to send data. If a bidirectional stream is locally opened but never
35/// used to send, then the peer may never see it. Application protocols should
36/// always arrange for the endpoint which will first transmit on a stream to be
37/// the endpoint responsible for opening it.
38///
39/// ## Data never received on a remotely-opened stream
40///
41/// Verify that the stream you are receiving is the same one that the server is
42/// sending on, e.g. by logging the [`id`] of each. Streams are always accepted
43/// in the same order as they are created, i.e. ascending order by [`StreamId`].
44/// For example, even if a sender first transmits on bidirectional stream 1, the
45/// first stream yielded by [`Connection::accept_bi`] on the receiver
46/// will be bidirectional stream 0.
47///
48/// [`stop()`]: RecvStream::stop
49/// [`id`]: RecvStream::id
50/// [`Connection::accept_bi`]: crate::Connection::accept_bi
51#[derive(Debug)]
52pub struct RecvStream {
53    conn: Shared<ConnectionInner>,
54    stream: StreamId,
55    is_0rtt: bool,
56    all_data_read: bool,
57    reset: Option<VarInt>,
58}
59
60impl RecvStream {
61    pub(crate) fn new(conn: Shared<ConnectionInner>, stream: StreamId, is_0rtt: bool) -> Self {
62        Self {
63            conn,
64            stream,
65            is_0rtt,
66            all_data_read: false,
67            reset: None,
68        }
69    }
70
71    /// Get the identity of this stream
72    pub fn id(&self) -> StreamId {
73        self.stream
74    }
75
76    /// Check if this stream has been opened during 0-RTT.
77    ///
78    /// In which case any non-idempotent request should be considered dangerous
79    /// at the application level. Because read data is subject to replay
80    /// attacks.
81    pub fn is_0rtt(&self) -> bool {
82        self.is_0rtt
83    }
84
85    /// Stop accepting data
86    ///
87    /// Discards unread data and notifies the peer to stop transmitting. Once
88    /// stopped, further attempts to operate on a stream will yield
89    /// `ClosedStream` errors.
90    pub fn stop(&mut self, error_code: VarInt) -> Result<(), ClosedStream> {
91        let mut state = self.conn.state();
92        if self.is_0rtt && !state.check_0rtt() {
93            return Ok(());
94        }
95        state.conn.recv_stream(self.stream).stop(error_code)?;
96        state.wake();
97        self.all_data_read = true;
98        Ok(())
99    }
100
101    /// Completes when the stream has been reset by the peer or otherwise closed
102    ///
103    /// Yields `Some` with the reset error code when the stream is reset by the
104    /// peer. Yields `None` when the stream was previously
105    /// [`stop()`](Self::stop)ed, or when the stream was
106    /// [`finish()`](crate::SendStream::finish)ed by the peer and all data has
107    /// been received, after which it is no longer meaningful for the stream
108    /// to be reset.
109    ///
110    /// This operation is cancel-safe.
111    pub async fn received_reset(&mut self) -> Result<Option<VarInt>, ResetError> {
112        poll_fn(|cx| {
113            let mut state = self.conn.state();
114
115            if self.is_0rtt && !state.check_0rtt() {
116                return Poll::Ready(Err(ResetError::ZeroRttRejected));
117            }
118            if let Some(code) = self.reset {
119                return Poll::Ready(Ok(Some(code)));
120            }
121
122            match state.conn.recv_stream(self.stream).received_reset() {
123                Err(_) => Poll::Ready(Ok(None)),
124                Ok(Some(error_code)) => {
125                    // Stream state has just now been freed, so the connection
126                    // may need to issue new stream ID flow
127                    // control credit
128                    state.wake();
129                    Poll::Ready(Ok(Some(error_code)))
130                }
131                Ok(None) => {
132                    if let Some(e) = &state.error {
133                        return Poll::Ready(Err(e.clone().into()));
134                    }
135                    // Resets always notify readers, since a reset is an
136                    // immediate read error. We
137                    // could introduce a dedicated channel to reduce the risk of
138                    // spurious wakeups, but that increased
139                    // complexity is probably not justified, as an application
140                    // that is expecting a reset is not likely to receive large
141                    // amounts of data.
142                    state.readable.insert(self.stream, cx.waker().clone());
143                    Poll::Pending
144                }
145            }
146        })
147        .await
148    }
149
150    /// Handle common logic related to reading out of a receive stream.
151    ///
152    /// This takes an `FnMut` closure that takes care of the actual reading
153    /// process, matching the detailed read semantics for the calling
154    /// function with a particular return type. The closure can read from
155    /// the passed `&mut Chunks` and has to return the status after reading:
156    /// the amount of data read, and the status after the final read call.
157    fn execute_poll_read<F, T>(
158        &mut self,
159        cx: &mut Context,
160        ordered: bool,
161        mut read_fn: F,
162    ) -> Poll<Result<Option<T>, ReadError>>
163    where
164        F: FnMut(&mut Chunks) -> ReadStatus<T>,
165    {
166        use quinn_proto::ReadError::*;
167
168        if self.all_data_read {
169            return Poll::Ready(Ok(None));
170        }
171
172        let mut state = self.conn.state();
173        if self.is_0rtt && !state.check_0rtt() {
174            return Poll::Ready(Err(ReadError::ZeroRttRejected));
175        }
176
177        // If we stored an error during a previous call, return it now. This can
178        // happen if a `read_fn` both wants to return data and also
179        // returns an error in its final stream status.
180        let status = match self.reset {
181            Some(code) => ReadStatus::Failed(None, Reset(code)),
182            None => {
183                let mut recv = state.conn.recv_stream(self.stream);
184                let mut chunks = recv.read(ordered)?;
185                let status = read_fn(&mut chunks);
186                if chunks.finalize().should_transmit() {
187                    state.wake();
188                }
189                status
190            }
191        };
192
193        match status {
194            ReadStatus::Readable(read) => Poll::Ready(Ok(Some(read))),
195            ReadStatus::Finished(read) => {
196                self.all_data_read = true;
197                Poll::Ready(Ok(read))
198            }
199            ReadStatus::Failed(read, Blocked) => match read {
200                Some(val) => Poll::Ready(Ok(Some(val))),
201                None => {
202                    if let Some(error) = &state.error {
203                        return Poll::Ready(Err(error.clone().into()));
204                    }
205                    state.readable.insert(self.stream, cx.waker().clone());
206                    Poll::Pending
207                }
208            },
209            ReadStatus::Failed(read, Reset(error_code)) => match read {
210                None => {
211                    self.all_data_read = true;
212                    self.reset = Some(error_code);
213                    Poll::Ready(Err(ReadError::Reset(error_code)))
214                }
215                done => {
216                    self.reset = Some(error_code);
217                    Poll::Ready(Ok(done))
218                }
219            },
220        }
221    }
222
223    pub(crate) fn poll_read_impl(
224        &mut self,
225        cx: &mut Context,
226        buf: &mut [MaybeUninit<u8>],
227    ) -> Poll<Result<Option<usize>, ReadError>> {
228        if buf.is_empty() {
229            return Poll::Ready(Ok(Some(0)));
230        }
231
232        self.execute_poll_read(cx, true, |chunks| {
233            let mut read = 0;
234            loop {
235                if read >= buf.len() {
236                    // We know `read > 0` because `buf` cannot be empty here
237                    return ReadStatus::Readable(read);
238                }
239
240                match chunks.next(buf.len() - read) {
241                    Ok(Some(chunk)) => {
242                        let bytes = chunk.bytes;
243                        let len = bytes.len();
244                        buf[read..read + len].copy_from_slice(unsafe {
245                            std::slice::from_raw_parts(bytes.as_ptr().cast(), len)
246                        });
247                        read += len;
248                    }
249                    res => {
250                        return (if read == 0 { None } else { Some(read) }, res.err()).into();
251                    }
252                }
253            }
254        })
255    }
256
257    /// Attempts to read from the stream into the provided buffer
258    ///
259    /// On success, returns `Poll::Ready(Ok(num_bytes_read))` and places data
260    /// into `buf`. If the buffer passed in has non-zero length and a 0 is
261    /// returned, that indicates that the remote side has [`finish`]ed the
262    /// stream and the local side has already read all bytes.
263    ///
264    /// If no data is available for reading, this returns `Poll::Pending` and
265    /// arranges for the current task (via `cx.waker()`) to be notified when
266    /// the stream becomes readable or is closed.
267    ///
268    /// [`finish`]: crate::SendStream::finish
269    pub fn poll_read_uninit(
270        &mut self,
271        cx: &mut Context,
272        buf: &mut [MaybeUninit<u8>],
273    ) -> Poll<Result<usize, ReadError>> {
274        self.poll_read_impl(cx, buf)
275            .map(|res| res.map(|n| n.unwrap_or_default()))
276    }
277
278    /// Read the next segment of data.
279    ///
280    /// Yields `None` if the stream was finished. Otherwise, yields a segment of
281    /// data and its offset in the stream. If `ordered` is `true`, the chunk's
282    /// offset will be immediately after the last data yielded by
283    /// [`read()`](Self::read) or [`read_chunk()`](Self::read_chunk). If
284    /// `ordered` is `false`, segments may be received in any order, and the
285    /// `Chunk`'s `offset` field can be used to determine ordering in the
286    /// caller. Unordered reads are less prone to head-of-line blocking within a
287    /// stream, but require the application to manage reassembling the original
288    /// data.
289    ///
290    /// Slightly more efficient than `read` due to not copying. Chunk boundaries
291    /// do not correspond to peer writes, and hence cannot be used as framing.
292    ///
293    /// This operation is cancel-safe.
294    pub async fn read_chunk(
295        &mut self,
296        max_length: usize,
297        ordered: bool,
298    ) -> Result<Option<Chunk>, ReadError> {
299        poll_fn(|cx| {
300            self.execute_poll_read(cx, ordered, |chunks| match chunks.next(max_length) {
301                Ok(Some(chunk)) => ReadStatus::Readable(chunk),
302                res => (None, res.err()).into(),
303            })
304        })
305        .await
306    }
307
308    /// Read the next segments of data.
309    ///
310    /// Fills `bufs` with the segments of data beginning immediately after the
311    /// last data yielded by `read` or `read_chunk`, or `None` if the stream was
312    /// finished.
313    ///
314    /// Slightly more efficient than `read` due to not copying. Chunk boundaries
315    /// do not correspond to peer writes, and hence cannot be used as framing.
316    ///
317    /// This operation is cancel-safe.
318    pub async fn read_chunks(&mut self, bufs: &mut [Bytes]) -> Result<Option<usize>, ReadError> {
319        if bufs.is_empty() {
320            return Ok(Some(0));
321        }
322
323        poll_fn(|cx| {
324            self.execute_poll_read(cx, true, |chunks| {
325                let mut read = 0;
326                loop {
327                    if read >= bufs.len() {
328                        // We know `read > 0` because `bufs` cannot be empty
329                        // here
330                        return ReadStatus::Readable(read);
331                    }
332
333                    match chunks.next(usize::MAX) {
334                        Ok(Some(chunk)) => {
335                            bufs[read] = chunk.bytes;
336                            read += 1;
337                        }
338                        res => {
339                            return (if read == 0 { None } else { Some(read) }, res.err()).into();
340                        }
341                    }
342                }
343            })
344        })
345        .await
346    }
347
348    /// Convenience method to read all remaining data into a buffer.
349    ///
350    /// If unordered reads have already been made, the resulting buffer may have
351    /// gaps containing zeros.
352    ///
353    /// This operation is *not* cancel-safe.
354    pub async fn read_to_end<B: IoBufMut>(&mut self, mut buf: B) -> BufResult<usize, B> {
355        let mut start = u64::MAX;
356        let mut end = 0;
357        let mut chunks = vec![];
358        loop {
359            let chunk = match self.read_chunk(usize::MAX, false).await {
360                Ok(Some(chunk)) => chunk,
361                Ok(None) => break,
362                Err(e) => return BufResult(Err(e.into()), buf),
363            };
364            start = start.min(chunk.offset);
365            end = end.max(chunk.offset + chunk.bytes.len() as u64);
366            chunks.push((chunk.offset, chunk.bytes));
367        }
368        if start == u64::MAX || start >= end {
369            // no data read
370            return BufResult(Ok(0), buf);
371        }
372        let len = (end - start) as usize;
373        let cap = buf.buf_capacity();
374        let needed = len.saturating_sub(cap);
375        if needed > 0
376            && let Err(e) = buf.reserve(needed)
377        {
378            return BufResult(Err(io::Error::new(io::ErrorKind::OutOfMemory, e)), buf);
379        }
380        let mut buf = buf.slice(..len);
381        let slice = buf.ensure_init();
382        for (offset, bytes) in chunks {
383            let offset = (offset - start) as usize;
384            let buf_len = bytes.len();
385            slice[offset..offset + buf_len].copy_from_slice(&bytes);
386        }
387        let mut buf = buf.into_inner();
388        unsafe { buf.advance_to(len) }
389        BufResult(Ok(len), buf)
390    }
391
392    /// Convert into an [`futures_util`] compatible stream.
393    #[cfg(feature = "io-compat")]
394    pub fn into_compat(self) -> CompatRecvStream {
395        CompatRecvStream(self)
396    }
397}
398
399impl Drop for RecvStream {
400    fn drop(&mut self) {
401        let mut state = self.conn.state();
402
403        // clean up any previously registered wakers
404        state.readable.remove(&self.stream);
405
406        if state.error.is_some() || (self.is_0rtt && !state.check_0rtt()) {
407            return;
408        }
409        if !self.all_data_read {
410            // Ignore ClosedStream errors
411            let _ = state.conn.recv_stream(self.stream).stop(0u32.into());
412            state.wake();
413        }
414    }
415}
416
417enum ReadStatus<T> {
418    Readable(T),
419    Finished(Option<T>),
420    Failed(Option<T>, quinn_proto::ReadError),
421}
422
423impl<T> From<(Option<T>, Option<quinn_proto::ReadError>)> for ReadStatus<T> {
424    fn from(status: (Option<T>, Option<quinn_proto::ReadError>)) -> Self {
425        match status {
426            (read, None) => Self::Finished(read),
427            (read, Some(e)) => Self::Failed(read, e),
428        }
429    }
430}
431
432/// Errors that arise from reading from a stream.
433#[derive(Debug, Error, Clone, PartialEq, Eq)]
434pub enum ReadError {
435    /// The peer abandoned transmitting data on this stream.
436    ///
437    /// Carries an application-defined error code.
438    #[error("stream reset by peer: error {0}")]
439    Reset(VarInt),
440    /// The connection was lost.
441    #[error("connection lost")]
442    ConnectionLost(#[from] ConnectionError),
443    /// The stream has already been stopped, finished, or reset.
444    #[error("closed stream")]
445    ClosedStream,
446    /// Attempted an ordered read following an unordered read.
447    ///
448    /// Performing an unordered read allows discontinuities to arise in the
449    /// receive buffer of a stream which cannot be recovered, making further
450    /// ordered reads impossible.
451    #[error("ordered read after unordered read")]
452    IllegalOrderedRead,
453    /// This was a 0-RTT stream and the server rejected it.
454    ///
455    /// Can only occur on clients for 0-RTT streams, which can be opened using
456    /// [`Connecting::into_0rtt()`].
457    ///
458    /// [`Connecting::into_0rtt()`]: crate::Connecting::into_0rtt()
459    #[error("0-RTT rejected")]
460    ZeroRttRejected,
461}
462
463impl From<ReadableError> for ReadError {
464    fn from(e: ReadableError) -> Self {
465        match e {
466            ReadableError::ClosedStream => Self::ClosedStream,
467            ReadableError::IllegalOrderedRead => Self::IllegalOrderedRead,
468        }
469    }
470}
471
472impl From<ResetError> for ReadError {
473    fn from(e: ResetError) -> Self {
474        match e {
475            ResetError::ConnectionLost(e) => Self::ConnectionLost(e),
476            ResetError::ZeroRttRejected => Self::ZeroRttRejected,
477        }
478    }
479}
480
481impl From<ReadError> for io::Error {
482    fn from(x: ReadError) -> Self {
483        use self::ReadError::*;
484        let kind = match x {
485            Reset { .. } | ZeroRttRejected => io::ErrorKind::ConnectionReset,
486            ConnectionLost(_) | ClosedStream => io::ErrorKind::NotConnected,
487            IllegalOrderedRead => io::ErrorKind::InvalidInput,
488        };
489        Self::new(kind, x)
490    }
491}
492
493/// Errors that arise from reading from a stream.
494#[derive(Debug, Error, Clone, PartialEq, Eq)]
495pub enum ReadExactError {
496    /// The stream finished before all bytes were read
497    #[error("stream finished early (expected {0} bytes more)")]
498    FinishedEarly(usize),
499    /// A read error occurred
500    #[error(transparent)]
501    ReadError(#[from] ReadError),
502}
503
504/// Errors that arise while waiting for a stream to be reset
505#[derive(Debug, Error, Clone, PartialEq, Eq)]
506pub enum ResetError {
507    /// The connection was lost
508    #[error("connection lost")]
509    ConnectionLost(#[from] ConnectionError),
510    /// This was a 0-RTT stream and the server rejected it
511    ///
512    /// Can only occur on clients for 0-RTT streams, which can be opened using
513    /// [`Connecting::into_0rtt()`].
514    ///
515    /// [`Connecting::into_0rtt()`]: crate::Connecting::into_0rtt()
516    #[error("0-RTT rejected")]
517    ZeroRttRejected,
518}
519
520impl From<ResetError> for io::Error {
521    fn from(x: ResetError) -> Self {
522        use ResetError::*;
523        let kind = match x {
524            ZeroRttRejected => io::ErrorKind::ConnectionReset,
525            ConnectionLost(_) => io::ErrorKind::NotConnected,
526        };
527        Self::new(kind, x)
528    }
529}
530
531impl AsyncRead for RecvStream {
532    async fn read<B: IoBufMut>(&mut self, mut buf: B) -> BufResult<usize, B> {
533        let res = poll_fn(|cx| self.poll_read_uninit(cx, buf.as_uninit()))
534            .await
535            .inspect(|&n| unsafe { buf.advance_to(n) })
536            .map_err(Into::into);
537        BufResult(res, buf)
538    }
539}
540
541#[cfg(feature = "io-compat")]
542mod compat {
543    use std::{
544        ops::{Deref, DerefMut},
545        pin::Pin,
546        task::ready,
547    };
548
549    use compio_buf::{IntoInner, bytes::BufMut};
550
551    use super::*;
552
553    /// A [`futures_util`] compatible receive stream.
554    pub struct CompatRecvStream(pub(super) RecvStream);
555
556    impl CompatRecvStream {
557        fn poll_read(
558            &mut self,
559            cx: &mut Context,
560            mut buf: impl BufMut,
561        ) -> Poll<Result<Option<usize>, ReadError>> {
562            self.poll_read_impl(cx, unsafe { buf.chunk_mut().as_uninit_slice_mut() })
563                .map(|res| {
564                    if let Ok(Some(n)) = &res {
565                        unsafe { buf.advance_mut(*n) }
566                    }
567                    res
568                })
569        }
570
571        /// Read data contiguously from the stream.
572        ///
573        /// Yields the number of bytes read into `buf` on success, or `None` if
574        /// the stream was finished.
575        ///
576        /// This operation is cancel-safe.
577        pub async fn read(&mut self, mut buf: impl BufMut) -> Result<Option<usize>, ReadError> {
578            poll_fn(|cx| self.poll_read(cx, &mut buf)).await
579        }
580
581        /// Read an exact number of bytes contiguously from the stream.
582        ///
583        /// See [`read()`] for details. This operation is *not* cancel-safe.
584        ///
585        /// [`read()`]: CompatRecvStream::read
586        pub async fn read_exact(&mut self, mut buf: impl BufMut) -> Result<(), ReadExactError> {
587            poll_fn(|cx| {
588                while buf.has_remaining_mut() {
589                    if ready!(self.poll_read(cx, &mut buf))?.is_none() {
590                        return Poll::Ready(Err(ReadExactError::FinishedEarly(
591                            buf.remaining_mut(),
592                        )));
593                    }
594                }
595                Poll::Ready(Ok(()))
596            })
597            .await
598        }
599    }
600
601    impl IntoInner for CompatRecvStream {
602        type Inner = RecvStream;
603
604        fn into_inner(self) -> Self::Inner {
605            self.0
606        }
607    }
608
609    impl Deref for CompatRecvStream {
610        type Target = RecvStream;
611
612        fn deref(&self) -> &Self::Target {
613            &self.0
614        }
615    }
616
617    impl DerefMut for CompatRecvStream {
618        fn deref_mut(&mut self) -> &mut Self::Target {
619            &mut self.0
620        }
621    }
622
623    impl futures_util::AsyncRead for CompatRecvStream {
624        fn poll_read(
625            self: Pin<&mut Self>,
626            cx: &mut Context<'_>,
627            buf: &mut [u8],
628        ) -> Poll<io::Result<usize>> {
629            // SAFETY: buf is valid
630            self.get_mut()
631                .poll_read_uninit(cx, unsafe {
632                    std::slice::from_raw_parts_mut(buf.as_mut_ptr().cast(), buf.len())
633                })
634                .map_err(Into::into)
635        }
636    }
637}
638
639#[cfg(feature = "io-compat")]
640pub use compat::CompatRecvStream;
641
642#[cfg(feature = "h3")]
643pub(crate) mod h3_impl {
644    use h3::quic::{self, StreamErrorIncoming};
645
646    use super::*;
647
648    impl From<ReadError> for StreamErrorIncoming {
649        fn from(e: ReadError) -> Self {
650            use ReadError::*;
651            match e {
652                Reset(code) => Self::StreamTerminated {
653                    error_code: code.into_inner(),
654                },
655                ConnectionLost(e) => Self::ConnectionErrorIncoming {
656                    connection_error: e.into(),
657                },
658                IllegalOrderedRead => unreachable!("illegal ordered read"),
659                e => Self::Unknown(Box::new(e)),
660            }
661        }
662    }
663
664    impl quic::RecvStream for RecvStream {
665        type Buf = Bytes;
666
667        fn poll_data(
668            &mut self,
669            cx: &mut Context<'_>,
670        ) -> Poll<Result<Option<Self::Buf>, StreamErrorIncoming>> {
671            self.execute_poll_read(cx, true, |chunks| match chunks.next(usize::MAX) {
672                Ok(Some(chunk)) => ReadStatus::Readable(chunk.bytes),
673                res => (None, res.err()).into(),
674            })
675            .map_err(Into::into)
676        }
677
678        fn stop_sending(&mut self, error_code: u64) {
679            self.stop(error_code.try_into().expect("invalid error_code"))
680                .ok();
681        }
682
683        fn recv_id(&self) -> quic::StreamId {
684            u64::from(self.stream).try_into().unwrap()
685        }
686    }
687}
compio_quic/recv_stream.rs

compio_quic/
recv_stream.rs
