1 // SPDX-License-Identifier: GPL-2.0
2 
3 //! Extensions to [`Box`] for fallible allocations.
4 
5 use super::{AllocError, Flags};
6 use alloc::boxed::Box;
7 use core::{mem::MaybeUninit, ptr, result::Result};
8 
9 /// Extensions to [`Box`].
10 pub trait BoxExt<T>: Sized {
11     /// Allocates a new box.
12     ///
13     /// The allocation may fail, in which case an error is returned.
new(x: T, flags: Flags) -> Result<Self, AllocError>14     fn new(x: T, flags: Flags) -> Result<Self, AllocError>;
15 
16     /// Allocates a new uninitialised box.
17     ///
18     /// The allocation may fail, in which case an error is returned.
new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>19     fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>;
20 
21     /// Drops the contents, but keeps the allocation.
22     ///
23     /// # Examples
24     ///
25     /// ```
26     /// use kernel::alloc::{flags, box_ext::BoxExt};
27     /// let value = Box::new([0; 32], flags::GFP_KERNEL)?;
28     /// assert_eq!(*value, [0; 32]);
29     /// let mut value = Box::drop_contents(value);
30     /// // Now we can re-use `value`:
31     /// value.write([1; 32]);
32     /// // SAFETY: We just wrote to it.
33     /// let value = unsafe { value.assume_init() };
34     /// assert_eq!(*value, [1; 32]);
35     /// # Ok::<(), Error>(())
36     /// ```
drop_contents(this: Self) -> Box<MaybeUninit<T>>37     fn drop_contents(this: Self) -> Box<MaybeUninit<T>>;
38 }
39 
40 impl<T> BoxExt<T> for Box<T> {
new(x: T, flags: Flags) -> Result<Self, AllocError>41     fn new(x: T, flags: Flags) -> Result<Self, AllocError> {
42         let mut b = <Self as BoxExt<_>>::new_uninit(flags)?;
43         b.write(x);
44         // SAFETY: We just wrote to it.
45         Ok(unsafe { b.assume_init() })
46     }
47 
48     #[cfg(any(test, testlib))]
new_uninit(_flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>49     fn new_uninit(_flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
50         Ok(Box::new_uninit())
51     }
52 
53     #[cfg(not(any(test, testlib)))]
new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>54     fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
55         let ptr = if core::mem::size_of::<MaybeUninit<T>>() == 0 {
56             core::ptr::NonNull::<_>::dangling().as_ptr()
57         } else {
58             let layout = core::alloc::Layout::new::<MaybeUninit<T>>();
59 
60             // SAFETY: Memory is being allocated (first arg is null). The only other source of
61             // safety issues is sleeping on atomic context, which is addressed by klint. Lastly,
62             // the type is not a SZT (checked above).
63             let ptr =
64                 unsafe { super::allocator::krealloc_aligned(core::ptr::null_mut(), layout, flags) };
65             if ptr.is_null() {
66                 return Err(AllocError);
67             }
68 
69             ptr.cast::<MaybeUninit<T>>()
70         };
71 
72         // SAFETY: For non-zero-sized types, we allocate above using the global allocator. For
73         // zero-sized types, we use `NonNull::dangling`.
74         Ok(unsafe { Box::from_raw(ptr) })
75     }
76 
drop_contents(this: Self) -> Box<MaybeUninit<T>>77     fn drop_contents(this: Self) -> Box<MaybeUninit<T>> {
78         let ptr = Box::into_raw(this);
79         // SAFETY: `ptr` is valid, because it came from `Box::into_raw`.
80         unsafe { ptr::drop_in_place(ptr) };
81 
82         // CAST: `MaybeUninit<T>` is a transparent wrapper of `T`.
83         let ptr = ptr.cast::<MaybeUninit<T>>();
84 
85         // SAFETY: `ptr` is valid for writes, because it came from `Box::into_raw` and it is valid for
86         // reads, since the pointer came from `Box::into_raw` and the type is `MaybeUninit<T>`.
87         unsafe { Box::from_raw(ptr) }
88     }
89 }
90