s2n_quic_core/buffer/reader/
limit.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    buffer::{
6        reader::{storage::Chunk, Reader, Storage},
7        writer::{self, Storage as _},
8    },
9    varint::VarInt,
10};
11
12/// Wraps a reader and limits the amount of data that can be read from it
13///
14/// This can be used for applying back pressure to the reader with flow control.
15pub struct Limit<'a, R: Reader + ?Sized> {
16    len: usize,
17    reader: &'a mut R,
18}
19
20impl<'a, R: Reader + ?Sized> Limit<'a, R> {
21    #[inline]
22    pub fn new(reader: &'a mut R, max_buffered_len: usize) -> Self {
23        let len = max_buffered_len.min(reader.buffered_len());
24
25        Self { len, reader }
26    }
27}
28
29impl<R: Reader + ?Sized> Storage for Limit<'_, R> {
30    type Error = R::Error;
31
32    #[inline]
33    fn buffered_len(&self) -> usize {
34        self.len
35    }
36
37    #[inline]
38    fn read_chunk(&mut self, watermark: usize) -> Result<Chunk<'_>, Self::Error> {
39        let watermark = self.len.min(watermark);
40        let chunk = self.reader.read_chunk(watermark)?;
41        unsafe {
42            assume!(chunk.len() <= self.len);
43        }
44        self.len -= chunk.len();
45        Ok(chunk)
46    }
47
48    #[inline]
49    fn partial_copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<Chunk<'_>, Self::Error>
50    where
51        Dest: writer::Storage + ?Sized,
52    {
53        let mut dest = dest.with_write_limit(self.len);
54        let mut dest = dest.track_write();
55        let chunk = self.reader.partial_copy_into(&mut dest)?;
56        let len = dest.written_len() + chunk.len();
57        unsafe {
58            assume!(len <= self.len);
59        }
60        self.len -= len;
61        Ok(chunk)
62    }
63
64    #[inline]
65    fn copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<(), Self::Error>
66    where
67        Dest: writer::Storage + ?Sized,
68    {
69        let mut dest = dest.with_write_limit(self.len);
70        let mut dest = dest.track_write();
71        self.reader.copy_into(&mut dest)?;
72        let len = dest.written_len();
73        unsafe {
74            assume!(len <= self.len);
75        }
76        self.len -= len;
77        Ok(())
78    }
79}
80
81impl<R: Reader + ?Sized> Reader for Limit<'_, R> {
82    #[inline]
83    fn current_offset(&self) -> VarInt {
84        self.reader.current_offset()
85    }
86
87    #[inline]
88    fn final_offset(&self) -> Option<VarInt> {
89        self.reader.final_offset()
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::stream::testing::Data;
97
98    #[test]
99    fn max_data_test() {
100        let mut reader = Data::new(1000);
101        assert_eq!(reader.buffered_len(), 1000);
102        let mut reader = reader.with_checks();
103
104        let max_data = 32usize;
105
106        let mut reader = reader.with_max_data(VarInt::from_u8(max_data as _));
107        assert_eq!(reader.buffered_len(), max_data);
108
109        let mut dest = &mut [0u8; 16][..];
110        let chunk = reader.partial_copy_into(&mut dest).unwrap();
111        assert_eq!(chunk.len(), 16);
112
113        assert_eq!(reader.buffered_len(), max_data - 16);
114
115        let mut dest = &mut [0u8; 16][..];
116        let chunk = reader.partial_copy_into(&mut dest).unwrap();
117        assert_eq!(chunk.len(), 16);
118        assert!(reader.buffer_is_empty());
119    }
120}
s2n_quic_core/buffer/reader/limit.rs

s2n_quic_core/buffer/reader/
limit.rs
