flex_dns/
buffer.rs

1use crate::BufferError;
2
3pub trait Buffer {
4    /// The whole buffer as a slice.
5    fn bytes(&self) -> &[u8];
6
7    /// The length of the data in the buffer.
8    fn len(&self) -> usize;
9
10    /// Reads a given number of bytes at the given offset.
11    fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError>;
12}
13
14pub trait MutBuffer {
15    /// Truncates the buffer to the given length.
16    fn truncate(&mut self, new_len: usize) -> Result<(), BufferError>;
17
18    /// Writes the given data at the given offset.
19    fn write_array_at<const BYTES: usize>(&mut self, offset: usize, data: [u8; BYTES]) -> Result<(), BufferError>;
20
21    /// Writes the given data at the end of the buffer.
22    fn write_bytes(&mut self, data: &[u8]) -> Result<(), BufferError>;
23
24    /// Writes the given data at the end of the buffer.
25    fn write_array<const BYTES: usize>(&mut self, data: [u8; BYTES]) -> Result<(), BufferError>;
26
27    /// Reads a given number of bytes at the given offset.
28    fn read_bytes_at_mut(&mut self, offset: usize, length: usize) -> Result<&mut [u8], BufferError>;
29}
30
31impl<T: Buffer> Buffer for &'_ T {
32    fn bytes(&self) -> &[u8] {
33        (*self).bytes()
34    }
35
36    fn len(&self) -> usize {
37        (*self).len()
38    }
39
40    fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError> {
41        (*self).read_bytes_at(offset, length)
42    }
43}
44
45impl<T: Buffer> Buffer for &'_ mut T {
46    fn bytes(&self) -> &[u8] {
47        <T as Buffer>::bytes(*self)
48    }
49
50    fn len(&self) -> usize {
51        <T as Buffer>::len(*self)
52    }
53
54    fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError> {
55        <T as Buffer>::read_bytes_at(*self, offset, length)
56    }
57}
58
59impl<T: MutBuffer + Buffer> MutBuffer for &'_ mut T {
60    fn truncate(&mut self, new_len: usize) -> Result<(), BufferError> {
61        (*self).truncate(new_len)
62    }
63
64    fn write_array_at<const BYTES: usize>(&mut self, offset: usize, data: [u8; BYTES]) -> Result<(), BufferError> {
65        (*self).write_array_at(offset, data)
66    }
67
68    fn write_bytes(&mut self, data: &[u8]) -> Result<(), BufferError> {
69        (*self).write_bytes(data)
70    }
71
72    fn write_array<const BYTES: usize>(&mut self, data: [u8; BYTES]) -> Result<(), BufferError> {
73        (*self).write_array(data)
74    }
75
76    fn read_bytes_at_mut(&mut self, offset: usize, length: usize) -> Result<&mut [u8], BufferError> {
77        (*self).read_bytes_at_mut(offset, length)
78    }
79}
80
81#[cfg(any(feature = "heapless", feature = "arrayvec", feature = "vec"))]
82#[inline(always)]
83fn check_length<const SIZE: usize>(offset: usize, length: usize) -> Result<(), BufferError> {
84    if offset + length > SIZE {
85        if offset > SIZE {
86            return Err(BufferError::OffsetOutOfBounds);
87        }
88        return Err(BufferError::LengthOutOfBounds);
89    }
90
91    Ok(())
92}
93
94impl Buffer for &'_ [u8] {
95    #[inline(always)]
96    fn bytes(&self) -> &[u8] {
97        self
98    }
99
100    #[inline(always)]
101    fn len(&self) -> usize {
102        <[u8]>::len(self)
103    }
104
105    #[inline(always)]
106    fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError> {
107        if offset + length > self.len() {
108            if offset > self.len() {
109                return Err(BufferError::OffsetOutOfBounds);
110            }
111            return Err(BufferError::LengthOutOfBounds);
112        }
113
114        Ok(&self[offset..offset + length])
115    }
116}
117
118#[cfg(feature = "arrayvec")]
119mod impl_arrayvec {
120    use super::*;
121
122    impl<const SIZE: usize> Buffer for arrayvec::ArrayVec<u8, SIZE> {
123        #[inline(always)]
124        fn bytes(&self) -> &[u8] {
125            self.as_slice()
126        }
127
128        #[inline(always)]
129        fn len(&self) -> usize {
130            arrayvec::ArrayVec::len(self)
131        }
132
133        #[inline(always)]
134        fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError> {
135            if offset + length > self.len() {
136                if offset > self.len() {
137                    return Err(BufferError::OffsetOutOfBounds);
138                }
139                return Err(BufferError::LengthOutOfBounds);
140            }
141
142            Ok(&self.as_slice()[offset..offset + length])
143        }
144    }
145
146    impl<const SIZE: usize> MutBuffer for arrayvec::ArrayVec<u8, SIZE> {
147        #[inline(always)]
148        fn truncate(&mut self, new_len: usize) -> Result<(), BufferError> {
149            if new_len > self.len() {
150                return Err(BufferError::InvalidLength);
151            }
152
153            arrayvec::ArrayVec::truncate(self, new_len);
154            Ok(())
155        }
156
157        #[inline(always)]
158        fn write_array_at<const BYTES: usize>(&mut self, offset: usize, data: [u8; BYTES]) -> Result<(), BufferError> {
159            check_length::<SIZE>(offset, BYTES)?;
160            // Grow the buffer to the required size
161            unsafe { self.set_len(core::cmp::max(self.len(), offset + BYTES)); };
162            self.as_mut_slice()[offset..offset + BYTES].copy_from_slice(&data);
163
164            Ok(())
165        }
166
167        #[inline(always)]
168        fn write_bytes(&mut self, data: &[u8]) -> Result<(), BufferError> {
169            check_length::<SIZE>(self.len(), data.len())?;
170            self.try_extend_from_slice(data).map_err(|_| BufferError::OutOfMemory)?;
171
172            Ok(())
173        }
174
175        #[inline(always)]
176        fn write_array<const BYTES: usize>(&mut self, data: [u8; BYTES]) -> Result<(), BufferError> {
177            check_length::<SIZE>(self.len(), BYTES)?;
178            self.try_extend_from_slice(&data).map_err(|_| BufferError::OutOfMemory)?;
179
180            Ok(())
181        }
182
183        #[inline(always)]
184        fn read_bytes_at_mut(&mut self, offset: usize, length: usize) -> Result<&mut [u8], BufferError> {
185            check_length::<SIZE>(offset, length)?;
186            unsafe { self.set_len(core::cmp::max(self.len(), offset + length)); };
187            Ok(&mut self.as_mut_slice()[offset..offset + length])
188        }
189    }
190}
191
192#[cfg(feature = "heapless")]
193mod impl_heapless {
194    use super::*;
195
196    impl<const SIZE: usize> Buffer for heapless::Vec<u8, SIZE> {
197        #[inline(always)]
198        fn bytes(&self) -> &[u8] {
199            self.as_slice()
200        }
201
202        #[inline(always)]
203        fn len(&self) -> usize {
204            self.bytes().len()
205        }
206
207        #[inline(always)]
208        fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError> {
209            if offset + length > self.len() {
210                if offset > self.len() {
211                    return Err(BufferError::OffsetOutOfBounds);
212                }
213                return Err(BufferError::LengthOutOfBounds);
214            }
215
216            Ok(&self.as_slice()[offset..offset + length])
217        }
218    }
219
220    impl<const SIZE: usize> MutBuffer for heapless::Vec<u8, SIZE> {
221        #[inline(always)]
222        fn truncate(&mut self, new_len: usize) -> Result<(), BufferError> {
223            if new_len > self.len() {
224                return Err(BufferError::InvalidLength);
225            }
226
227            heapless::Vec::truncate(self, new_len);
228            Ok(())
229        }
230
231        #[inline(always)]
232        fn write_array_at<const BYTES: usize>(&mut self, offset: usize, data: [u8; BYTES]) -> Result<(), BufferError> {
233            check_length::<SIZE>(offset, BYTES)?;
234            // Grow the buffer to the required size
235            unsafe { self.set_len(core::cmp::max(self.len(), offset + BYTES)); };
236            self[offset..offset + BYTES].copy_from_slice(&data);
237
238            Ok(())
239        }
240
241        #[inline(always)]
242        fn write_bytes(&mut self, data: &[u8]) -> Result<(), BufferError> {
243            check_length::<SIZE>(self.len(), data.len())?;
244            self.extend_from_slice(data).map_err(|_| BufferError::OutOfMemory)?;
245
246            Ok(())
247        }
248
249        #[inline(always)]
250        fn write_array<const BYTES: usize>(&mut self, data: [u8; BYTES]) -> Result<(), BufferError> {
251            check_length::<SIZE>(self.len(), BYTES)?;
252            self.extend_from_slice(&data).map_err(|_| BufferError::OutOfMemory)?;
253
254            Ok(())
255        }
256
257        #[inline(always)]
258        fn read_bytes_at_mut(&mut self, offset: usize, length: usize) -> Result<&mut [u8], BufferError> {
259            check_length::<SIZE>(offset, length)?;
260            unsafe { self.set_len(core::cmp::max(self.len(), offset + length)); };
261            Ok(&mut self[offset..offset + length])
262        }
263    }
264}
265
266#[cfg(feature = "vec")]
267mod impl_vec {
268    extern crate alloc;
269
270    use super::*;
271
272    impl Buffer for alloc::vec::Vec<u8> {
273        #[inline(always)]
274        fn bytes(&self) -> &[u8] {
275            self.as_slice()
276        }
277
278        #[inline(always)]
279        fn len(&self) -> usize {
280            <alloc::vec::Vec<u8>>::len(self)
281        }
282
283        #[inline(always)]
284        fn read_bytes_at(&self, offset: usize, length: usize) -> Result<&[u8], BufferError> {
285            if offset + length > self.len() {
286                if offset > self.len() {
287                    return Err(BufferError::OffsetOutOfBounds);
288                }
289                return Err(BufferError::LengthOutOfBounds);
290            }
291
292            Ok(&self.as_slice()[offset..offset + length])
293        }
294    }
295
296    impl MutBuffer for alloc::vec::Vec<u8> {
297        #[inline(always)]
298        fn truncate(&mut self, new_len: usize) -> Result<(), BufferError> {
299            if new_len > self.len() {
300                return Err(BufferError::InvalidLength);
301            }
302
303            <alloc::vec::Vec<u8>>::truncate(self, new_len);
304            Ok(())
305        }
306
307        #[inline(always)]
308        fn write_array_at<const BYTES: usize>(&mut self, offset: usize, data: [u8; BYTES]) -> Result<(), BufferError> {
309            if offset + BYTES > self.len() {
310                if offset > self.len() {
311                    return Err(BufferError::OffsetOutOfBounds);
312                }
313                return Err(BufferError::LengthOutOfBounds);
314            }
315
316            self.as_mut_slice()[offset..offset + BYTES].copy_from_slice(&data);
317
318            Ok(())
319        }
320
321        #[inline(always)]
322        fn write_bytes(&mut self, data: &[u8]) -> Result<(), BufferError> {
323            self.extend_from_slice(data);
324            Ok(())
325        }
326
327        #[inline(always)]
328        fn write_array<const BYTES: usize>(&mut self, data: [u8; BYTES]) -> Result<(), BufferError> {
329            self.extend_from_slice(&data);
330            Ok(())
331        }
332
333        #[inline(always)]
334        fn read_bytes_at_mut(&mut self, offset: usize, length: usize) -> Result<&mut [u8], BufferError> {
335            if offset + length > self.len() {
336                if offset > self.len() {
337                    return Err(BufferError::OffsetOutOfBounds);
338                }
339                return Err(BufferError::LengthOutOfBounds);
340            }
341
342            Ok(&mut self.as_mut_slice()[offset..offset + length])
343        }
344    }
345}