wlan_common/
buffer_writer.rs

Help
1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::append::{Append, BufferTooSmall, TrackedAppend};
6use zerocopy::SplitByteSliceMut;
7
8pub struct BufferWriter<B> {
9    buffer: B,
10    written: usize,
11}
12
13impl<B: SplitByteSliceMut> BufferWriter<B> {
14    pub fn new(buffer: B) -> Self {
15        Self { buffer, written: 0 }
16    }
17
18    pub fn remaining(&self) -> usize {
19        self.buffer.len() - self.written
20    }
21
22    pub fn written(&self) -> usize {
23        self.written
24    }
25
26    pub fn into_written(self) -> B {
27        let (written, _remaining) = self.buffer.split_at(self.written).ok().unwrap();
28        written
29    }
30
31    fn next_mut_slice(&mut self, len: usize) -> Result<&mut [u8], BufferTooSmall> {
32        if self.written + len > self.buffer.len() {
33            return Err(BufferTooSmall);
34        }
35        let ret = &mut self.buffer[self.written..self.written + len];
36        self.written += len;
37        Ok(ret)
38    }
39}
40
41impl<B: SplitByteSliceMut> Append for BufferWriter<B> {
42    fn append_bytes(&mut self, bytes: &[u8]) -> Result<(), BufferTooSmall> {
43        self.next_mut_slice(bytes.len())?.copy_from_slice(bytes);
44        Ok(())
45    }
46
47    fn append_bytes_zeroed(&mut self, len: usize) -> Result<&mut [u8], BufferTooSmall> {
48        let ret = self.next_mut_slice(len)?;
49        for x in &mut ret[..] {
50            *x = 0;
51        }
52        Ok(ret)
53    }
54
55    fn can_append(&self, bytes: usize) -> bool {
56        self.remaining() >= bytes
57    }
58}
59
60impl<B: SplitByteSliceMut> TrackedAppend for BufferWriter<B> {
61    fn bytes_appended(&self) -> usize {
62        self.written
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn buffer_too_short() {
72        assert!(BufferWriter::new(&mut [0; 0][..]).append_byte(1).is_err());
73        assert!(BufferWriter::new(&mut [0; 5][..]).append_bytes(&[0; 6]).is_err());
74        assert!(BufferWriter::new(&mut [0; 5][..]).append_value_zeroed::<[u8; 6]>().is_err());
75    }
76
77    #[test]
78    fn append_value_zeroed() {
79        let mut buffer = [1u8; 5];
80        let mut w = BufferWriter::new(&mut buffer[..]);
81        let mut data = w.append_value_zeroed::<[u8; 3]>().expect("failed writing buffer");
82        data[0] = 42;
83        // Don't write `data[1]`: BufferWriter should zero this byte.
84        data[2] = 43;
85
86        assert_eq!(3, w.written());
87        assert_eq!(2, w.remaining());
88        assert_eq!([42, 0, 43, 1, 1], buffer);
89    }
90
91    #[test]
92    fn append_bytes() {
93        let mut buffer = [1u8; 5];
94        let mut w = BufferWriter::new(&mut buffer[..]);
95        w.append_byte(42).expect("failed writing buffer");
96        w.append_byte(43).expect("failed writing buffer");
97        w.append_bytes(&[2, 3]).expect("failed writing buffer");
98
99        assert_eq!(4, w.written());
100        assert_eq!(1, w.remaining());
101        assert_eq!(&[42, 43, 2, 3], w.into_written());
102        assert_eq!([42, 43, 2, 3, 1], buffer);
103    }
104
105    #[test]
106    fn can_append() {
107        let mut buffer = [1u8; 5];
108        let mut w = BufferWriter::new(&mut buffer[..]);
109        assert!(w.can_append(0));
110        assert!(w.can_append(4));
111        assert!(w.can_append(5));
112        assert!(!w.can_append(6));
113
114        w.append_byte(42).unwrap();
115        assert!(w.can_append(4));
116        assert!(!w.can_append(5));
117    }
118}
wlan_common/buffer_writer.rs

wlan_common/
buffer_writer.rs
