giit_rbpf/
aligned_memory.rs

1//! Aligned memory
2
3/// Provides u8 slices at a specified alignment
4#[derive(Clone, Debug, PartialEq)]
5pub struct AlignedMemory {
6    max_len: usize,
7    align_offset: usize,
8    mem: Vec<u8>,
9}
10impl AlignedMemory {
11    fn get_mem(max_len: usize, align: usize) -> (Vec<u8>, usize) {
12        let mut mem: Vec<u8> = Vec::with_capacity(max_len + align);
13        mem.push(0);
14        let align_offset = mem.as_ptr().align_offset(align);
15        mem.resize(align_offset, 0);
16        (mem, align_offset)
17    }
18    /// Return a new AlignedMemory type
19    pub fn new(max_len: usize, align: usize) -> Self {
20        let (mem, align_offset) = Self::get_mem(max_len, align);
21        Self {
22            max_len,
23            align_offset,
24            mem,
25        }
26    }
27    /// Return a pre-filled AlignedMemory type
28    pub fn new_with_size(len: usize, align: usize) -> Self {
29        let (mut mem, align_offset) = Self::get_mem(len, align);
30        mem.resize(align_offset + len, 0);
31        Self {
32            max_len: len,
33            align_offset,
34            mem,
35        }
36    }
37    /// Return a pre-filled AlignedMemory type
38    pub fn new_with_data(data: &[u8], align: usize) -> Self {
39        let max_len = data.len();
40        let (mut mem, align_offset) = Self::get_mem(max_len, align);
41        mem.extend_from_slice(data);
42        Self {
43            max_len,
44            align_offset,
45            mem,
46        }
47    }
48    /// Get the length of the data
49    pub fn len(&self) -> usize {
50        self.mem.len() - self.align_offset
51    }
52    /// Is the memory empty
53    pub fn is_empty(&self) -> bool {
54        self.mem.len() - self.align_offset == 0
55    }
56    /// Get the current write index
57    pub fn write_index(&self) -> usize {
58        self.mem.len()
59    }
60    /// Get an aligned slice
61    pub fn as_slice(&self) -> &[u8] {
62        let start = self.align_offset;
63        let end = self.mem.len();
64        &self.mem[start..end]
65    }
66    /// Get an aligned mutable slice
67    pub fn as_slice_mut(&mut self) -> &mut [u8] {
68        let start = self.align_offset;
69        let end = self.mem.len();
70        &mut self.mem[start..end]
71    }
72    /// resize memory with value starting at the write_index
73    pub fn resize(&mut self, num: usize, value: u8) -> std::io::Result<()> {
74        if self.mem.len() + num > self.align_offset + self.max_len {
75            return Err(std::io::Error::new(
76                std::io::ErrorKind::InvalidInput,
77                "aligned memory resize failed",
78            ));
79        }
80        self.mem.resize(self.mem.len() + num, value);
81        Ok(())
82    }
83}
84impl std::io::Write for AlignedMemory {
85    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
86        if self.mem.len() + buf.len() > self.align_offset + self.max_len {
87            return Err(std::io::Error::new(
88                std::io::ErrorKind::InvalidInput,
89                "aligned memory write failed",
90            ));
91        }
92        self.mem.extend_from_slice(buf);
93        Ok(buf.len())
94    }
95    fn flush(&mut self) -> std::io::Result<()> {
96        Ok(())
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use std::io::Write;
104
105    fn do_test(align: usize) {
106        let mut aligned_memory = AlignedMemory::new(10, align);
107
108        assert_eq!(aligned_memory.write(&[42u8; 1]).unwrap(), 1);
109        assert_eq!(aligned_memory.write(&[42u8; 9]).unwrap(), 9);
110        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
111        assert_eq!(aligned_memory.write(&[42u8; 0]).unwrap(), 0);
112        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
113        aligned_memory.write(&[42u8; 1]).unwrap_err();
114        assert_eq!(aligned_memory.as_slice(), &[42u8; 10]);
115        aligned_memory.as_slice_mut().copy_from_slice(&[84u8; 10]);
116        assert_eq!(aligned_memory.as_slice(), &[84u8; 10]);
117
118        let mut aligned_memory = AlignedMemory::new(10, align);
119        aligned_memory.resize(5, 0).unwrap();
120        aligned_memory.resize(2, 1).unwrap();
121        assert_eq!(aligned_memory.write(&[2u8; 3]).unwrap(), 3);
122        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
123        aligned_memory.resize(1, 3).unwrap_err();
124        aligned_memory.write(&[4u8; 1]).unwrap_err();
125        assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
126    }
127
128    #[test]
129    fn test_aligned_memory() {
130        do_test(1);
131        do_test(32768);
132    }
133}
giit_rbpf/aligned_memory.rs

giit_rbpf/
aligned_memory.rs
