embedded_tls/
read_buffer.rs

1/// A reference to consume bytes from the internal buffer.
2#[must_use]
3pub struct ReadBuffer<'a> {
4    data: &'a [u8],
5    consumed: usize,
6    used: bool,
7
8    decrypted_consumed: &'a mut usize,
9}
10
11impl<'a> ReadBuffer<'a> {
12    #[inline]
13    pub(crate) fn new(buffer: &'a [u8], decrypted_consumed: &'a mut usize) -> Self {
14        Self {
15            data: buffer,
16            consumed: 0,
17            used: false,
18            decrypted_consumed,
19        }
20    }
21
22    #[inline]
23    pub fn len(&self) -> usize {
24        self.data.len() - self.consumed
25    }
26
27    #[inline]
28    pub fn is_empty(&self) -> bool {
29        self.len() == 0
30    }
31
32    /// Consumes and returns a slice of at most `count` bytes.
33    #[inline]
34    pub fn peek(&mut self, count: usize) -> &'a [u8] {
35        let count = self.len().min(count);
36        let start = self.consumed;
37
38        // We mark the buffer used to prevent dropping unconsumed bytes.
39        self.used = true;
40
41        &self.data[start..start + count]
42    }
43
44    /// Consumes and returns a slice of at most `count` bytes.
45    #[inline]
46    pub fn peek_all(&mut self) -> &'a [u8] {
47        self.peek(self.len())
48    }
49
50    /// Consumes and returns a slice of at most `count` bytes.
51    #[inline]
52    pub fn pop(&mut self, count: usize) -> &'a [u8] {
53        let count = self.len().min(count);
54        let start = self.consumed;
55        self.consumed += count;
56        self.used = true;
57
58        &self.data[start..start + count]
59    }
60
61    /// Consumes and returns the internal buffer.
62    #[inline]
63    pub fn pop_all(&mut self) -> &'a [u8] {
64        self.pop(self.len())
65    }
66
67    /// Drops the reference and restores internal buffer.
68    #[inline]
69    pub fn revert(self) {
70        core::mem::forget(self);
71    }
72
73    /// Tries to fills the buffer by consuming and copying bytes into it.
74    #[inline]
75    pub fn pop_into(&mut self, buf: &mut [u8]) -> usize {
76        let to_copy = self.pop(buf.len());
77
78        buf[..to_copy.len()].copy_from_slice(to_copy);
79
80        to_copy.len()
81    }
82}
83
84impl Drop for ReadBuffer<'_> {
85    #[inline]
86    fn drop(&mut self) {
87        *self.decrypted_consumed += if self.used {
88            self.consumed
89        } else {
90            // Consume all if dropped unused
91            self.data.len()
92        };
93    }
94}
95
96#[cfg(test)]
97mod test {
98    use super::*;
99
100    #[test]
101    fn dropping_unused_buffer_consumes_all() {
102        let mut consumed = 1000;
103        let buffer = [0, 1, 2, 3];
104
105        _ = ReadBuffer::new(&buffer, &mut consumed);
106
107        assert_eq!(consumed, 1004);
108    }
109
110    #[test]
111    fn pop_moves_internal_cursor() {
112        let mut consumed = 0;
113
114        let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
115
116        assert_eq!(buffer.pop(1), &[0]);
117        assert_eq!(buffer.pop(1), &[1]);
118        assert_eq!(buffer.pop(1), &[2]);
119    }
120
121    #[test]
122    fn dropping_consumes_as_many_bytes_as_used() {
123        let mut consumed = 0;
124
125        let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
126
127        assert_eq!(buffer.pop(1), &[0]);
128        assert_eq!(buffer.pop(1), &[1]);
129        assert_eq!(buffer.pop(1), &[2]);
130
131        core::mem::drop(buffer);
132
133        assert_eq!(consumed, 3);
134    }
135
136    #[test]
137    fn pop_returns_fewer_bytes_if_requested_more_than_what_it_has() {
138        let mut consumed = 0;
139
140        let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
141
142        assert_eq!(buffer.pop(1), &[0]);
143        assert_eq!(buffer.pop(1), &[1]);
144        assert_eq!(buffer.pop(4), &[2, 3]);
145        assert_eq!(buffer.pop(1), &[]);
146
147        core::mem::drop(buffer);
148
149        assert_eq!(consumed, 4);
150    }
151
152    #[test]
153    fn peek_does_not_consume() {
154        let mut consumed = 0;
155
156        let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
157
158        assert_eq!(buffer.peek(1), &[0]);
159        assert_eq!(buffer.peek(1), &[0]);
160
161        core::mem::drop(buffer);
162
163        assert_eq!(consumed, 0);
164    }
165
166    #[test]
167    fn revert_undoes_pop() {
168        let mut consumed = 0;
169
170        let mut buffer = ReadBuffer::new(&[0, 1, 2, 3], &mut consumed);
171
172        assert_eq!(buffer.pop(4), &[0, 1, 2, 3]);
173
174        buffer.revert();
175
176        assert_eq!(consumed, 0);
177    }
178}
embedded_tls/read_buffer.rs

embedded_tls/
read_buffer.rs
