wlan_common/
append.rs

Help
1// Copyright 2020 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 std::mem::size_of;
6use std::ops::Deref;
7use thiserror::Error;
8use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref, Unaligned};
9
10#[derive(Error, Debug)]
11#[error("Buffer is too small for the written data")]
12pub struct BufferTooSmall;
13
14pub trait Append {
15    fn append_bytes(&mut self, bytes: &[u8]) -> Result<(), BufferTooSmall>;
16
17    fn append_bytes_zeroed(&mut self, len: usize) -> Result<&mut [u8], BufferTooSmall>;
18
19    fn can_append(&self, bytes: usize) -> bool;
20
21    fn append_value<T>(&mut self, value: &T) -> Result<(), BufferTooSmall>
22    where
23        T: IntoBytes + Immutable + ?Sized,
24    {
25        self.append_bytes(value.as_bytes())
26    }
27
28    fn append_byte(&mut self, byte: u8) -> Result<(), BufferTooSmall> {
29        self.append_bytes(&[byte])
30    }
31
32    fn append_value_zeroed<T>(&mut self) -> Result<Ref<&mut [u8], T>, BufferTooSmall>
33    where
34        T: FromBytes + KnownLayout + Immutable + Unaligned,
35    {
36        let bytes = self.append_bytes_zeroed(size_of::<T>())?;
37        Ok(Ref::from_bytes(bytes).unwrap())
38    }
39
40    fn append_array_zeroed<T>(
41        &mut self,
42        num_elems: usize,
43    ) -> Result<Ref<&mut [u8], [T]>, BufferTooSmall>
44    where
45        T: FromBytes + Immutable + Unaligned,
46    {
47        let bytes = self.append_bytes_zeroed(size_of::<T>() * num_elems)?;
48        Ok(Ref::from_bytes(bytes).unwrap())
49    }
50}
51
52impl Append for Vec<u8> {
53    fn append_bytes(&mut self, bytes: &[u8]) -> Result<(), BufferTooSmall> {
54        self.extend_from_slice(bytes);
55        Ok(())
56    }
57
58    fn append_bytes_zeroed(&mut self, len: usize) -> Result<&mut [u8], BufferTooSmall> {
59        let old_len = self.len();
60        self.resize(old_len + len, 0);
61        Ok(&mut self[old_len..])
62    }
63
64    fn can_append(&self, _bytes: usize) -> bool {
65        true
66    }
67}
68
69pub trait TrackedAppend: Append {
70    fn bytes_appended(&self) -> usize;
71}
72
73#[derive(Debug)]
74pub struct VecCursor {
75    written: usize,
76    inner: Vec<u8>,
77}
78
79impl VecCursor {
80    pub fn new() -> Self {
81        Self { written: 0, inner: vec![] }
82    }
83
84    pub fn with_capacity(capacity: usize) -> Self {
85        Self { written: 0, inner: Vec::with_capacity(capacity) }
86    }
87
88    pub fn len(&self) -> usize {
89        self.inner.len()
90    }
91
92    pub fn capacity(&self) -> usize {
93        self.inner.capacity()
94    }
95
96    pub fn into_vec(self) -> Vec<u8> {
97        self.inner
98    }
99}
100
101impl From<VecCursor> for Vec<u8> {
102    fn from(other: VecCursor) -> Self {
103        other.inner
104    }
105}
106
107impl Deref for VecCursor {
108    type Target = [u8];
109
110    fn deref(&self) -> &Self::Target {
111        &self.inner[..]
112    }
113}
114
115impl Append for VecCursor {
116    fn append_bytes(&mut self, bytes: &[u8]) -> Result<(), BufferTooSmall> {
117        self.inner.extend_from_slice(bytes);
118        self.written += bytes.len();
119        Ok(())
120    }
121
122    fn append_bytes_zeroed(&mut self, len: usize) -> Result<&mut [u8], BufferTooSmall> {
123        let old_len = self.len();
124        self.inner.resize(old_len + len, 0);
125        self.written += len;
126        Ok(&mut self.inner[old_len..])
127    }
128
129    fn can_append(&self, _bytes: usize) -> bool {
130        true
131    }
132}
133
134impl TrackedAppend for VecCursor {
135    fn bytes_appended(&self) -> usize {
136        self.written
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    pub fn append_to_vec() {
146        let mut data = VecCursor::new();
147        assert_eq!(0, data.bytes_appended());
148
149        data.append_bytes(&[1, 2, 3]).unwrap();
150        assert_eq!(3, data.bytes_appended());
151
152        let bytes = data.append_bytes_zeroed(2).unwrap();
153        bytes[0] = 4;
154        assert_eq!(5, data.bytes_appended());
155
156        data.append_value(&0x0706_u16).unwrap();
157        assert_eq!(7, data.bytes_appended());
158
159        data.append_byte(8).unwrap();
160        assert_eq!(8, data.bytes_appended());
161
162        let mut bytes = data.append_value_zeroed::<[u8; 2]>().unwrap();
163        bytes[0] = 9;
164        assert_eq!(10, data.bytes_appended());
165
166        data.append_value(&[0x0c0b_u16, 0x0e0d_u16]).unwrap();
167        assert_eq!(14, data.bytes_appended());
168
169        let mut arr = data.append_array_zeroed::<[u8; 2]>(2).unwrap();
170        arr[0] = [15, 16];
171        assert_eq!(18, data.bytes_appended());
172
173        #[rustfmt::skip]
174        assert_eq!(
175            &[
176                1, 2, 3,
177                4, 0,
178                6, 7,
179                8,
180                9, 0,
181                11, 12, 13, 14,
182                15, 16, 0, 0
183            ],
184            &data[..]
185        );
186    }
187}
wlan_common/append.rs

wlan_common/
append.rs
