s2n_quic_core/buffer/reader/storage/
tracked.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::buffer::{
5    reader::{Reader, Storage},
6    writer,
7};
8
9pub struct Tracked<'a, S: Storage + ?Sized> {
10    consumed: usize,
11    storage: &'a mut S,
12}
13
14impl<'a, S: Storage + ?Sized> Tracked<'a, S> {
15    #[inline]
16    pub fn new(storage: &'a mut S) -> Self {
17        Self {
18            consumed: 0,
19            storage,
20        }
21    }
22
23    #[inline]
24    pub fn consumed_len(&self) -> usize {
25        self.consumed
26    }
27}
28
29impl<S: Storage + ?Sized> Storage for Tracked<'_, S> {
30    type Error = S::Error;
31
32    #[inline]
33    fn buffered_len(&self) -> usize {
34        self.storage.buffered_len()
35    }
36
37    #[inline]
38    fn read_chunk(&mut self, watermark: usize) -> Result<super::Chunk<'_>, Self::Error> {
39        let chunk = self.storage.read_chunk(watermark)?;
40        self.consumed += chunk.len();
41        Ok(chunk)
42    }
43
44    #[inline]
45    fn partial_copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<super::Chunk<'_>, Self::Error>
46    where
47        Dest: writer::Storage + ?Sized,
48    {
49        let mut dest = dest.track_write();
50        let chunk = self.storage.partial_copy_into(&mut dest)?;
51        self.consumed += dest.written_len();
52        self.consumed += chunk.len();
53        Ok(chunk)
54    }
55
56    #[inline]
57    fn copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<(), Self::Error>
58    where
59        Dest: writer::Storage + ?Sized,
60    {
61        let mut dest = dest.track_write();
62        self.storage.copy_into(&mut dest)?;
63        self.consumed += dest.written_len();
64        Ok(())
65    }
66}
67
68impl<S: Reader + ?Sized> Reader for Tracked<'_, S> {
69    #[inline]
70    fn current_offset(&self) -> crate::varint::VarInt {
71        self.storage.current_offset()
72    }
73
74    #[inline]
75    fn final_offset(&self) -> Option<crate::varint::VarInt> {
76        self.storage.final_offset()
77    }
78
79    #[inline]
80    fn has_buffered_fin(&self) -> bool {
81        self.storage.has_buffered_fin()
82    }
83
84    #[inline]
85    fn is_consumed(&self) -> bool {
86        self.storage.is_consumed()
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use writer::Storage as _;
94
95    #[test]
96    fn tracked_test() {
97        let mut reader = &b"hello world"[..];
98        let mut writer: Vec<u8> = vec![];
99
100        {
101            let mut reader = reader.track_read();
102            let chunk = reader.read_chunk(1).unwrap();
103            assert_eq!(&chunk[..], b"h");
104            assert_eq!(reader.consumed_len(), 1);
105        }
106
107        {
108            let mut reader = reader.track_read();
109            let mut writer = writer.with_write_limit(5);
110
111            let chunk = reader.partial_copy_into(&mut writer).unwrap();
112            assert_eq!(&chunk[..], b"ello ");
113            assert_eq!(reader.consumed_len(), 5);
114        }
115
116        assert_eq!(writer.len(), 0);
117
118        {
119            let mut reader = reader.track_read();
120            reader.copy_into(&mut writer).unwrap();
121            assert_eq!(reader.consumed_len(), 5);
122            assert_eq!(&writer[..], b"world");
123        }
124
125        assert_eq!(reader.len(), 0);
126    }
127}
s2n_quic_core/buffer/reader/storage/tracked.rs

s2n_quic_core/buffer/reader/storage/
tracked.rs
