1 // Copyright 2024, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 //! ImageBuffer is helper to store uninitialized memory buffer. And work with it allowing to read
16 //! into the buffer and retrieve initialized part.
17 //!
18 //! Similar to [ReadBuf](https://docs.rs/tokio/latest/tokio/io/struct.ReadBuf.html) but works in
19 //! `no_std`.
20
21 use core::mem::MaybeUninit;
22 use liberror::{Error, Result};
23
24 /// Wrapper class for buffer received with [get_buffer] function.
25 ///
26 /// Helps to keep track of allocated/init memory and avoid getting same buffer more than once.
27 #[derive(Debug)]
28 pub struct ImageBuffer<'a> {
29 buffer: Option<&'a mut [MaybeUninit<u8>]>,
30 // number of initialized and filled bytes.
31 used_bytes: usize,
32 }
33
34 // Unstable MaybeUninit API
35 // feature = "maybe_uninit_slice", issue = "63569"
36
37 // Assuming all the elements are initialized, get a mutable slice to them.
38 //
39 // # Safety
40 //
41 // It is up to the caller to guarantee that the `MaybeUninit<T>` elements
42 // really are in an initialized state.
43 // Calling this when the content is not yet fully initialized causes undefined behavior.
44 #[inline(always)]
slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T]45 unsafe fn slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
46 // SAFETY: Caller must make sure provided data is initialized.
47 unsafe { &mut *(slice as *mut [MaybeUninit<T>] as *mut [T]) }
48 }
49
50 // Assuming all the elements are initialized, get a slice of them.
51 //
52 // # Safety
53 //
54 // It is up to the caller to guarantee that the `MaybeUninit<T>` elements
55 // really are in an initialized state.
56 // Calling this when the content is not yet fully initialized causes undefined behavior.
57 #[inline(always)]
slice_assume_init_ref<T>(slice: &[MaybeUninit<T>]) -> &[T]58 unsafe fn slice_assume_init_ref<T>(slice: &[MaybeUninit<T>]) -> &[T] {
59 // SAFETY: Caller must make sure provided data is initialized.
60 unsafe { &*(slice as *const [MaybeUninit<T>] as *const [T]) }
61 }
62
63 impl ImageBuffer<'_> {
64 /// Create new ImageBuffer from buffer and used_bytes
new(buffer: &mut [MaybeUninit<u8>]) -> ImageBuffer65 pub fn new(buffer: &mut [MaybeUninit<u8>]) -> ImageBuffer {
66 ImageBuffer { buffer: Some(buffer), used_bytes: 0 }
67 }
68
69 /// Total buffer capacity.
capacity(&self) -> usize70 pub fn capacity(&self) -> usize {
71 self.buffer.as_ref().unwrap().len()
72 }
73
74 /// Increase used part of the buffer by `len`
75 ///
76 /// Return:
77 /// Error() - if current used_bytes + len > capacity, or causes arithmetic overflow.
78 /// Ok(()) - on success
79 ///
80 /// SAFETY:
81 /// It is up to the user to guarantee that `len` bytes for tail was initialized and filled.
advance_used(&mut self, len: usize) -> Result<()>82 pub unsafe fn advance_used(&mut self, len: usize) -> Result<()> {
83 let Some(new_len) = self.used_bytes.checked_add(len) else {
84 return Err(Error::Other(Some("Used bytes overflow")));
85 };
86 if new_len > self.capacity() {
87 return Err(Error::BufferTooSmall(Some(new_len)));
88 }
89 self.used_bytes = new_len;
90 Ok(())
91 }
92
93 /// Return used and tail parts of the buffer
get_split(&self) -> (&[u8], &[MaybeUninit<u8>])94 pub fn get_split(&self) -> (&[u8], &[MaybeUninit<u8>]) {
95 let (used, tail) = self.buffer.as_ref().unwrap().split_at(self.used_bytes);
96 // SAFETY:
97 //
98 // ImageBuffer user guarantees that changing used elements means they were initialized.
99 // And object assumes initialized only for slice [..used_bytes]
100 let initialized = unsafe { slice_assume_init_ref(used) };
101 (initialized, tail)
102 }
103
104 /// Return used and tail parts of the buffer
get_split_mut(&mut self) -> (&mut [u8], &mut [MaybeUninit<u8>])105 pub fn get_split_mut(&mut self) -> (&mut [u8], &mut [MaybeUninit<u8>]) {
106 let (used, tail) = self.buffer.as_mut().unwrap().split_at_mut(self.used_bytes);
107 // SAFETY:
108 //
109 // ImageBuffer user guaranties that changing used elements means they were initialized.
110 // And object assumes initialized only for slice [..used_bytes]
111 let initialized = unsafe { slice_assume_init_mut(used) };
112 (initialized, tail)
113 }
114
115 /// Slice of the buffer that is used
used(&self) -> &[u8]116 pub fn used(&self) -> &[u8] {
117 let (used, _) = self.get_split();
118 used
119 }
120
121 /// Slice of the buffer that is used
used_mut(&mut self) -> &mut [u8]122 pub fn used_mut(&mut self) -> &mut [u8] {
123 let (used, _) = self.get_split_mut();
124 used
125 }
126
127 /// Return part of the buffer that is not used
tail(&mut self) -> &mut [MaybeUninit<u8>]128 pub fn tail(&mut self) -> &mut [MaybeUninit<u8>] {
129 let (_, tail) = self.get_split_mut();
130 tail
131 }
132 }
133
134 impl AsRef<[MaybeUninit<u8>]> for ImageBuffer<'_> {
as_ref(&self) -> &[MaybeUninit<u8>]135 fn as_ref(&self) -> &[MaybeUninit<u8>] {
136 self.buffer.as_ref().unwrap()
137 }
138 }
139
140 impl AsMut<[MaybeUninit<u8>]> for ImageBuffer<'_> {
as_mut(&mut self) -> &mut [MaybeUninit<u8>]141 fn as_mut(&mut self) -> &mut [MaybeUninit<u8>] {
142 self.buffer.as_mut().unwrap()
143 }
144 }
145
146 impl PartialEq for ImageBuffer<'_> {
eq(&self, other: &Self) -> bool147 fn eq(&self, other: &Self) -> bool {
148 self.used() == other.used()
149 }
150 }
151
152 #[cfg(test)]
153 /// Helper to create ImageBuffers from Vec<u8>
154 pub struct ImageBufferVec {
155 buf: Vec<u8>,
156 }
157
158 #[cfg(test)]
159 impl ImageBufferVec {
new(buf: Vec<u8>) -> Self160 pub fn new(buf: Vec<u8>) -> Self {
161 Self { buf }
162 }
163
get(&mut self) -> ImageBuffer164 pub fn get(&mut self) -> ImageBuffer {
165 ImageBuffer::new(Self::slice_assume_not_init_mut(self.buf.as_mut_slice()))
166 }
167
slice_assume_not_init_mut<T>(slice: &mut [T]) -> &mut [MaybeUninit<T>]168 fn slice_assume_not_init_mut<T>(slice: &mut [T]) -> &mut [MaybeUninit<T>] {
169 // SAFETY: similar to safety notes for `slice_get_ref`, but we have a
170 // mutable reference which is also guaranteed to be valid for writes.
171 unsafe { &mut *(slice as *mut [T] as *mut [MaybeUninit<T>]) }
172 }
173 }
174
175 #[cfg(test)]
176 mod test {
177 use super::*;
178
179 #[test]
test_image_buffer_capacity()180 fn test_image_buffer_capacity() {
181 assert_eq!(ImageBufferVec::new(vec![0u8; 0]).get().capacity(), 0);
182 assert_eq!(ImageBufferVec::new(vec![0u8; 0]).get().capacity(), 0);
183 assert_eq!(ImageBufferVec::new(vec![0u8; 1]).get().capacity(), 1);
184 assert_eq!(ImageBufferVec::new(vec![0u8; 100]).get().capacity(), 100);
185 assert_eq!(
186 ImageBufferVec::new(vec![0u8; 128 * 1024 * 1024]).get().capacity(),
187 128 * 1024 * 1024
188 );
189 }
190
191 #[test]
test_image_buffer_used()192 fn test_image_buffer_used() {
193 let mut img_buf_vec = ImageBufferVec::new(vec![0u8; 100]);
194 let mut img_buf = img_buf_vec.get();
195 assert_eq!(img_buf.used().len(), 0);
196 assert_eq!(img_buf.used_mut().len(), 0);
197 // SAFETY:
198 // All data in img_buf is initialized since it was created from vec
199 unsafe { img_buf.advance_used(1).unwrap() };
200 assert_eq!(img_buf.used().len(), 1);
201 assert_eq!(img_buf.used_mut().len(), 1);
202 // SAFETY:
203 // All data in img_buf is initialized since it was created from vec
204 unsafe { img_buf.advance_used(3).unwrap() };
205 assert_eq!(img_buf.used().len(), 4);
206 assert_eq!(img_buf.used_mut().len(), 4);
207 assert_eq!(
208 // SAFETY:
209 // All data in img_buf is initialized since it was created from vec
210 unsafe { img_buf.advance_used(1024) },
211 Err(Error::BufferTooSmall(Some(1028)))
212 );
213 assert_eq!(img_buf.used().len(), 4);
214 assert_eq!(img_buf.used_mut().len(), 4);
215 }
216
217 #[test]
test_image_buffer_get_split()218 fn test_image_buffer_get_split() {
219 let mut img_buf_vec = ImageBufferVec::new(vec![0u8, 1, 2, 3]);
220 let mut img_buf = img_buf_vec.get();
221
222 assert_eq!(img_buf.used(), [].as_slice());
223 assert_eq!(img_buf.used_mut(), [].as_mut_slice());
224 assert_eq!(img_buf.tail().len(), 4);
225 let (used, tail) = img_buf.get_split();
226 assert_eq!(used, [].as_slice());
227 assert_eq!(tail.len(), 4);
228
229 let (used, tail) = img_buf.get_split_mut();
230 assert_eq!(used, [].as_mut_slice());
231 assert_eq!(tail.len(), 4);
232
233 // SAFETY:
234 // All data in img_buf is initialized since it was created from vec
235 unsafe { img_buf.advance_used(2).unwrap() };
236 assert_eq!(img_buf.used(), [0, 1].as_slice());
237 assert_eq!(img_buf.used_mut(), [0, 1].as_mut_slice());
238 assert_eq!(img_buf.tail().len(), 2);
239 let (used, tail) = img_buf.get_split();
240 assert_eq!(used, [0, 1].as_slice());
241 assert_eq!(tail.len(), 2);
242 let (used, tail) = img_buf.get_split_mut();
243 assert_eq!(used, [0, 1].as_mut_slice());
244 assert_eq!(tail.len(), 2);
245
246 // SAFETY:
247 // All data in img_buf is initialized since it was created from vec
248 unsafe { img_buf.advance_used(2).unwrap() };
249 assert_eq!(img_buf.used(), [0, 1, 2, 3].as_slice());
250 assert_eq!(img_buf.used_mut(), [0, 1, 2, 3].as_mut_slice());
251 assert_eq!(img_buf.tail().len(), 0);
252 let (used, tail) = img_buf.get_split();
253 assert_eq!(used, [0, 1, 2, 3].as_slice());
254 assert_eq!(tail.len(), 0);
255 let (used, tail) = img_buf.get_split_mut();
256 assert_eq!(used, [0, 1, 2, 3].as_mut_slice());
257 assert_eq!(tail.len(), 0);
258 }
259
260 #[test]
test_image_buffer_eq_not_init()261 fn test_image_buffer_eq_not_init() {
262 assert_eq!(
263 ImageBufferVec::new(vec![0u8, 1, 2]).get(),
264 ImageBufferVec::new(vec![0u8, 1, 2]).get()
265 );
266 }
267
268 #[test]
test_image_buffer_eq_init_same()269 fn test_image_buffer_eq_init_same() {
270 let mut v1 = ImageBufferVec::new(vec![0u8, 1, 2]);
271 let mut v2 = ImageBufferVec::new(vec![0u8, 1, 2]);
272 let mut image_buffer_1 = v1.get();
273 let mut image_buffer_2 = v2.get();
274
275 // SAFETY:
276 // Buffers initialised on creation.
277 unsafe {
278 image_buffer_1.advance_used(3).unwrap();
279 image_buffer_2.advance_used(3).unwrap();
280 }
281
282 assert_eq!(image_buffer_1, image_buffer_2);
283 }
284
285 #[test]
test_image_buffer_eq_diff_capacity()286 fn test_image_buffer_eq_diff_capacity() {
287 let mut v1 = ImageBufferVec::new(vec![0u8, 1, 2]);
288 let mut v2 = ImageBufferVec::new(vec![0u8, 1, 2, 3]);
289 let mut image_buffer_1 = v1.get();
290 let mut image_buffer_2 = v2.get();
291
292 // SAFETY:
293 // Buffers initialised on creation.
294 unsafe {
295 image_buffer_1.advance_used(2).unwrap();
296 image_buffer_2.advance_used(2).unwrap();
297 }
298
299 assert_eq!(image_buffer_1, image_buffer_2);
300 }
301 }
302