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

embedded_tls/
read_buffer.rs
