rs_matter/utils/storage/
writebuf.rs

1/*
2 *
3 *    Copyright (c) 2022-2026 Project CHIP Authors
4 *
5 *    Licensed under the Apache License, Version 2.0 (the "License");
6 *    you may not use this file except in compliance with the License.
7 *    You may obtain a copy of the License at
8 *
9 *        http://www.apache.org/licenses/LICENSE-2.0
10 *
11 *    Unless required by applicable law or agreed to in writing, software
12 *    distributed under the License is distributed on an "AS IS" BASIS,
13 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 *    See the License for the specific language governing permissions and
15 *    limitations under the License.
16 */
17
18use crate::error::*;
19
20#[derive(Debug)]
21#[cfg_attr(feature = "defmt", derive(defmt::Format))]
22pub struct WriteBuf<'a> {
23    pub(crate) buf: &'a mut [u8],
24    buf_size: usize,
25    start: usize,
26    end: usize,
27}
28
29impl<'a> WriteBuf<'a> {
30    pub fn new(buf: &'a mut [u8]) -> Self {
31        Self::new_with(buf, 0, 0)
32    }
33
34    pub fn new_with(buf: &'a mut [u8], start: usize, end: usize) -> Self {
35        let buf_size = buf.len();
36
37        Self {
38            buf,
39            buf_size,
40            start,
41            end,
42        }
43    }
44
45    pub fn get_start(&self) -> usize {
46        self.start
47    }
48
49    pub fn get_tail(&self) -> usize {
50        self.end
51    }
52
53    pub fn rewind_tail_to(&mut self, new_end: usize) {
54        self.end = new_end;
55    }
56
57    pub fn forward_tail_by(&mut self, new_offset: usize) {
58        self.end += new_offset
59    }
60
61    pub fn as_slice(&self) -> &[u8] {
62        &self.buf[self.start..self.end]
63    }
64
65    pub fn as_mut_slice(&mut self) -> &mut [u8] {
66        &mut self.buf[self.start..self.end]
67    }
68
69    pub fn empty_as_mut_slice(&mut self) -> &mut [u8] {
70        &mut self.buf[self.end..self.buf_size]
71    }
72
73    pub fn split(self) -> (&'a mut [u8], Self) {
74        let (head, tail) = self.buf.split_at_mut(self.end);
75        (head, Self::new(tail))
76    }
77
78    pub fn split_str(self) -> (&'a str, Self) {
79        let (bytes, wb) = self.split();
80
81        (core::str::from_utf8(bytes).unwrap(), wb)
82    }
83
84    pub fn into_buf(self) -> &'a mut [u8] {
85        self.buf
86    }
87
88    pub fn reset(&mut self) {
89        self.buf_size = self.buf.len();
90        self.start = 0;
91        self.end = 0;
92    }
93
94    pub fn load(&mut self, wb: &WriteBuf) -> Result<(), Error> {
95        if self.buf_size < wb.end {
96            Err(ErrorCode::NoSpace)?;
97        }
98
99        self.buf[0..wb.end].copy_from_slice(&wb.buf[..wb.end]);
100        self.start = wb.start;
101        self.end = wb.end;
102
103        Ok(())
104    }
105
106    pub fn reserve(&mut self, reserve: usize) -> Result<(), Error> {
107        if self.end != 0 || self.start != 0 || self.buf_size != self.buf.len() {
108            Err(ErrorCode::Invalid.into())
109        } else if reserve > self.buf_size {
110            Err(ErrorCode::NoSpace.into())
111        } else {
112            self.start = reserve;
113            self.end = reserve;
114            Ok(())
115        }
116    }
117
118    pub fn shrink(&mut self, with: usize) -> Result<(), Error> {
119        if self.end + with <= self.buf_size {
120            self.buf_size -= with;
121            Ok(())
122        } else {
123            Err(ErrorCode::NoSpace.into())
124        }
125    }
126
127    pub fn expand(&mut self, by: usize) -> Result<(), Error> {
128        if self.buf.len() - self.buf_size >= by {
129            self.buf_size += by;
130            Ok(())
131        } else {
132            Err(ErrorCode::NoSpace.into())
133        }
134    }
135
136    pub fn prepend_with<F>(&mut self, size: usize, f: F) -> Result<(), Error>
137    where
138        F: FnOnce(&mut Self),
139    {
140        if size <= self.start {
141            f(self);
142            self.start -= size;
143            return Ok(());
144        }
145        Err(ErrorCode::NoSpace.into())
146    }
147
148    pub fn prepend(&mut self, src: &[u8]) -> Result<(), Error> {
149        self.prepend_with(src.len(), |x| {
150            let dst_slice = &mut x.buf[(x.start - src.len())..x.start];
151            dst_slice.copy_from_slice(src);
152        })
153    }
154
155    pub fn append_with_buf<F>(&mut self, f: F) -> Result<usize, Error>
156    where
157        F: FnOnce(&mut [u8]) -> Result<usize, Error>,
158    {
159        let len = f(self.empty_as_mut_slice())?;
160        self.end += len;
161
162        Ok(len)
163    }
164
165    pub fn append_with<F>(&mut self, size: usize, f: F) -> Result<(), Error>
166    where
167        F: FnOnce(&mut Self),
168    {
169        if self.end + size <= self.buf_size {
170            f(self);
171            self.end += size;
172            return Ok(());
173        }
174        Err(ErrorCode::NoSpace.into())
175    }
176
177    pub fn append(&mut self, src: &[u8]) -> Result<(), Error> {
178        self.copy_from_slice(src)
179    }
180
181    pub fn copy_from_slice(&mut self, src: &[u8]) -> Result<(), Error> {
182        self.append_with(src.len(), |x| {
183            x.buf[x.end..(x.end + src.len())].copy_from_slice(src);
184        })
185    }
186
187    pub fn le_i8(&mut self, data: i8) -> Result<(), Error> {
188        self.le_u8(data as u8)
189    }
190
191    pub fn le_u8(&mut self, data: u8) -> Result<(), Error> {
192        self.append_with(1, |x| {
193            x.buf[x.end] = data;
194        })
195    }
196
197    pub fn le_u16(&mut self, data: u16) -> Result<(), Error> {
198        self.append(&data.to_le_bytes())
199    }
200
201    pub fn le_i16(&mut self, data: i16) -> Result<(), Error> {
202        self.append(&data.to_le_bytes())
203    }
204
205    pub fn le_u32(&mut self, data: u32) -> Result<(), Error> {
206        self.append(&data.to_le_bytes())
207    }
208
209    pub fn le_i32(&mut self, data: i32) -> Result<(), Error> {
210        self.append(&data.to_le_bytes())
211    }
212
213    pub fn le_u64(&mut self, data: u64) -> Result<(), Error> {
214        self.append(&data.to_le_bytes())
215    }
216
217    pub fn le_i64(&mut self, data: i64) -> Result<(), Error> {
218        self.append(&data.to_le_bytes())
219    }
220}
221
222impl core::fmt::Write for WriteBuf<'_> {
223    fn write_str(&mut self, s: &str) -> core::fmt::Result {
224        self.append(s.as_bytes()).map_err(|_| core::fmt::Error)?;
225
226        Ok(())
227    }
228}
229
230#[collapse_debuginfo(yes)]
231macro_rules! write_split {
232    ($f:expr, $s:literal $(, $x:expr)* $(,)?) => {
233        {
234            write!(&mut $f, $s $(, $x)*).map_err(|_| ErrorCode::BufferTooSmall)?;
235
236            Ok::<_, Error>($f.split_str())
237        }
238    };
239}
240
241pub(crate) use write_split;
242
243#[cfg(test)]
244mod tests {
245    use crate::utils::storage::WriteBuf;
246
247    #[test]
248    fn test_append_le_with_success() {
249        let mut test_slice = [0; 22];
250        let mut buf = WriteBuf::new(&mut test_slice);
251        unwrap!(buf.reserve(5));
252
253        unwrap!(buf.le_u8(1));
254        unwrap!(buf.le_u16(65));
255        unwrap!(buf.le_u32(0xcafebabe));
256        unwrap!(buf.le_u64(0xcafebabecafebabe));
257        unwrap!(buf.le_u16(64));
258        assert_eq!(
259            test_slice,
260            [
261                0, 0, 0, 0, 0, 1, 65, 0, 0xbe, 0xba, 0xfe, 0xca, 0xbe, 0xba, 0xfe, 0xca, 0xbe,
262                0xba, 0xfe, 0xca, 64, 0
263            ]
264        );
265    }
266
267    #[test]
268    fn test_len_param() {
269        let mut test_slice = [0; 20];
270        let mut buf = WriteBuf::new(&mut test_slice[..5]);
271        unwrap!(buf.reserve(5));
272
273        let _ = buf.le_u8(1);
274        let _ = buf.le_u16(65);
275        let _ = buf.le_u32(0xcafebabe);
276        let _ = buf.le_u64(0xcafebabecafebabe);
277        // All of the above must return error, and hence the slice shouldn't change
278        assert_eq!(test_slice, [0; 20]);
279    }
280
281    #[test]
282    fn test_overrun() {
283        let mut test_slice = [0; 20];
284        let mut buf = WriteBuf::new(&mut test_slice);
285        unwrap!(buf.reserve(4));
286        unwrap!(buf.le_u64(0xcafebabecafebabe));
287        unwrap!(buf.le_u64(0xcafebabecafebabe));
288        // Now the buffer is fully filled up, so no further puts will happen
289
290        if buf.le_u8(1).is_ok() {
291            panic!("Should return error")
292        }
293
294        if buf.le_u16(65).is_ok() {
295            panic!("Should return error")
296        }
297
298        if buf.le_u32(0xcafebabe).is_ok() {
299            panic!("Should return error")
300        }
301
302        if buf.le_u64(0xcafebabecafebabe).is_ok() {
303            panic!("Should return error")
304        }
305    }
306
307    #[test]
308    fn test_as_slice() {
309        let mut test_slice = [0; 20];
310        let mut buf = WriteBuf::new(&mut test_slice);
311        unwrap!(buf.reserve(5));
312
313        unwrap!(buf.le_u8(1));
314        unwrap!(buf.le_u16(65));
315        unwrap!(buf.le_u32(0xcafebabe));
316        unwrap!(buf.le_u64(0xcafebabecafebabe));
317
318        let new_slice: [u8; 3] = [0xa, 0xb, 0xc];
319        unwrap!(buf.prepend(&new_slice));
320
321        assert_eq!(
322            buf.as_slice(),
323            [
324                0xa, 0xb, 0xc, 1, 65, 0, 0xbe, 0xba, 0xfe, 0xca, 0xbe, 0xba, 0xfe, 0xca, 0xbe,
325                0xba, 0xfe, 0xca
326            ]
327        );
328    }
329
330    #[test]
331    fn test_copy_as_slice() {
332        let mut test_slice = [0; 20];
333        let mut buf = WriteBuf::new(&mut test_slice);
334        unwrap!(buf.reserve(5));
335
336        unwrap!(buf.le_u16(65));
337        let new_slice: [u8; 5] = [0xaa, 0xbb, 0xcc, 0xdd, 0xee];
338        unwrap!(buf.copy_from_slice(&new_slice));
339        unwrap!(buf.le_u32(65));
340        assert_eq!(
341            test_slice,
342            [0, 0, 0, 0, 0, 65, 0, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 65, 0, 0, 0, 0, 0, 0, 0]
343        );
344    }
345
346    #[test]
347    fn test_copy_as_slice_overrun() {
348        let mut test_slice = [0; 20];
349        let mut buf = WriteBuf::new(&mut test_slice[..7]);
350        unwrap!(buf.reserve(5));
351
352        unwrap!(buf.le_u16(65));
353        let new_slice: [u8; 5] = [0xaa, 0xbb, 0xcc, 0xdd, 0xee];
354        if buf.copy_from_slice(&new_slice).is_ok() {
355            panic!("This should have returned error")
356        }
357    }
358
359    #[test]
360    fn test_prepend() {
361        let mut test_slice = [0; 20];
362        let mut buf = WriteBuf::new(&mut test_slice);
363        unwrap!(buf.reserve(5));
364
365        unwrap!(buf.le_u16(65));
366        let new_slice: [u8; 5] = [0xaa, 0xbb, 0xcc, 0xdd, 0xee];
367        unwrap!(buf.prepend(&new_slice));
368        assert_eq!(
369            test_slice,
370            [0xaa, 0xbb, 0xcc, 0xdd, 0xee, 65, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
371        );
372    }
373
374    #[test]
375    fn test_prepend_overrun() {
376        let mut test_slice = [0; 20];
377        let mut buf = WriteBuf::new(&mut test_slice);
378        unwrap!(buf.reserve(5));
379
380        unwrap!(buf.le_u16(65));
381        let new_slice: [u8; 6] = [0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff];
382        if buf.prepend(&new_slice).is_ok() {
383            panic!("Prepend should return error")
384        }
385    }
386
387    #[test]
388    fn test_rewind_tail() {
389        let mut test_slice = [0; 20];
390        let mut buf = WriteBuf::new(&mut test_slice);
391        unwrap!(buf.reserve(5));
392
393        unwrap!(buf.le_u16(65));
394
395        let anchor = buf.get_tail();
396
397        let new_slice: [u8; 5] = [0xaa, 0xbb, 0xcc, 0xdd, 0xee];
398        unwrap!(buf.copy_from_slice(&new_slice));
399        assert_eq!(buf.as_slice(), [65, 0, 0xaa, 0xbb, 0xcc, 0xdd, 0xee,]);
400
401        buf.rewind_tail_to(anchor);
402        unwrap!(buf.le_u16(66));
403        assert_eq!(buf.as_slice(), [65, 0, 66, 0,]);
404    }
405}
rs_matter/utils/storage/writebuf.rs

rs_matter/utils/storage/
writebuf.rs
