s2n_quic_core/buffer/reader/
complete.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, Error,
8    },
9    varint::VarInt,
10};
11
12/// Wraps a single [`Storage`] instance as a [`Reader`].
13///
14/// This can be used for scenarios where the entire stream is buffered and known up-front.
15#[derive(Debug)]
16pub struct Complete<'a, S> {
17    storage: &'a mut S,
18    current_offset: VarInt,
19    final_offset: VarInt,
20}
21
22impl<'a, S> Complete<'a, S>
23where
24    S: Storage,
25{
26    #[inline]
27    pub fn new(storage: &'a mut S) -> Result<Self, Error> {
28        let final_offset = VarInt::try_from(storage.buffered_len())
29            .ok()
30            .ok_or(Error::OutOfRange)?;
31        Ok(Self {
32            storage,
33            current_offset: VarInt::ZERO,
34            final_offset,
35        })
36    }
37}
38
39impl<S> Storage for Complete<'_, S>
40where
41    S: Storage,
42{
43    type Error = S::Error;
44
45    #[inline]
46    fn buffered_len(&self) -> usize {
47        self.storage.buffered_len()
48    }
49
50    #[inline]
51    fn buffer_is_empty(&self) -> bool {
52        self.storage.buffer_is_empty()
53    }
54
55    #[inline]
56    fn read_chunk(&mut self, watermark: usize) -> Result<Chunk, Self::Error> {
57        let chunk = self.storage.read_chunk(watermark)?;
58        self.current_offset += chunk.len();
59        Ok(chunk)
60    }
61
62    #[inline]
63    fn partial_copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<Chunk, Self::Error>
64    where
65        Dest: writer::Storage + ?Sized,
66    {
67        let mut dest = dest.track_write();
68        let chunk = self.storage.partial_copy_into(&mut dest)?;
69        self.current_offset += chunk.len();
70        self.current_offset += dest.written_len();
71        Ok(chunk)
72    }
73
74    #[inline]
75    fn copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<(), Self::Error>
76    where
77        Dest: writer::Storage + ?Sized,
78    {
79        let mut dest = dest.track_write();
80        self.storage.copy_into(&mut dest)?;
81        self.current_offset += dest.written_len();
82        Ok(())
83    }
84}
85
86impl<C> Reader for Complete<'_, C>
87where
88    C: Storage,
89{
90    #[inline]
91    fn current_offset(&self) -> VarInt {
92        self.current_offset
93    }
94
95    #[inline]
96    fn final_offset(&self) -> Option<VarInt> {
97        Some(self.final_offset)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn read_chunk_test() {
107        let mut storage: &[u8] = &[1, 2, 3, 4];
108        let mut reader = Complete::new(&mut storage).unwrap();
109        let mut reader = reader.with_checks();
110
111        assert_eq!(reader.current_offset(), VarInt::ZERO);
112        assert_eq!(reader.final_offset(), Some(VarInt::from_u8(4)));
113
114        let chunk = reader.read_chunk(usize::MAX).unwrap();
115        assert_eq!(&*chunk, &[1, 2, 3, 4]);
116
117        assert_eq!(reader.current_offset(), VarInt::from_u8(4));
118        assert!(reader.buffer_is_empty());
119    }
120
121    #[test]
122    fn partial_copy_test() {
123        let mut storage: &[u8] = &[1, 2, 3, 4];
124        let mut reader = Complete::new(&mut storage).unwrap();
125        let mut reader = reader.with_checks();
126
127        assert_eq!(reader.current_offset(), VarInt::ZERO);
128        assert_eq!(reader.final_offset(), Some(VarInt::from_u8(4)));
129
130        let mut dest: &mut [u8] = &mut [0; 4];
131        let chunk = reader.partial_copy_into(&mut dest).unwrap();
132        assert_eq!(&*chunk, &[1, 2, 3, 4]);
133
134        assert_eq!(reader.current_offset(), VarInt::from_u8(4));
135        assert!(reader.buffer_is_empty());
136    }
137
138    #[test]
139    fn copy_test() {
140        let mut storage: &[u8] = &[1, 2, 3, 4];
141        let mut reader = Complete::new(&mut storage).unwrap();
142        let mut reader = reader.with_checks();
143
144        assert_eq!(reader.current_offset(), VarInt::ZERO);
145        assert_eq!(reader.final_offset(), Some(VarInt::from_u8(4)));
146
147        let mut dest = [0; 4];
148        {
149            let mut dest = &mut dest[..];
150            reader.copy_into(&mut dest).unwrap();
151        }
152        assert_eq!(&dest[..], &[1, 2, 3, 4]);
153
154        assert_eq!(reader.current_offset(), VarInt::from_u8(4));
155        assert!(reader.buffer_is_empty());
156    }
157}
s2n_quic_core/buffer/reader/complete.rs

s2n_quic_core/buffer/reader/
complete.rs
