cocoon_tpm_utils_common/
alloc.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2023-2025 SUSE LLC
3// Author: Nicolai Stange <nstange@suse.de>
4
5//! Helpers related to `alloc`.
6
7extern crate alloc;
8use alloc::{boxed::Box, vec::Vec};
9use core::{mem, ptr};
10
11use crate::zeroize;
12
13/// Memory allocation error.
14#[derive(Clone, Copy, Debug)]
15pub enum TryNewError {
16    /// Memory allocation failure.
17    MemoryAllocationFailure,
18}
19
20/// Try to allocate a `Box`, handling allocation failure gracefully.
21///
22/// Currently `Box::try_new()` is still unstable, so this implements an
23/// alternative Box instantiation primitive enabling graceful memory allocation
24/// failure handling.
25///
26/// # Arguments:
27///
28/// * `v` - The value to wrap in a `Box`.
29///
30/// # Errors:
31///
32/// * [`TryNewError::MemoryAllocationFailure`] - The memory allocation has
33///   failed.
34pub fn box_try_new<T>(v: T) -> Result<Box<T>, TryNewError> {
35    // Box::try_new() is unstable, so do it by ourselves for now.
36    // Refer to https://doc.rust-lang.org/std/boxed/index.html#memory-layout.
37    let p: *mut T = if mem::size_of::<T>() == 0 {
38        // Dangling pointers are valid for ZSTs and the write below is Ok.
39        ptr::NonNull::dangling().as_ptr()
40    } else {
41        let layout = alloc::alloc::Layout::new::<T>();
42        let p: *mut T = unsafe { alloc::alloc::alloc(layout) } as *mut T;
43        if p.is_null() {
44            return Err(TryNewError::MemoryAllocationFailure);
45        }
46        p
47    };
48
49    unsafe { p.write(v) };
50
51    Ok(unsafe { Box::from_raw(p) })
52}
53
54/// Error returned by [`box_try_new_with()`](box_try_new_with).
55#[derive(Clone, Copy, Debug)]
56pub enum TryNewWithError<E> {
57    /// Memory allocation failure.
58    TryNew(TryNewError),
59    /// The object factory callback passed to
60    /// [`box_try_new_with()`](box_try_new_with) returned an error, wrapped
61    /// in the variant.
62    With(E),
63}
64
65/// Try to initialize a `Box` from a provided factory callback.
66///
67/// Invoke `f()` to obtain the `Box`' wrapped value only after the heap memory
68/// allocation has succeeded and store the returned object in the `Box`.
69///
70/// This enables the compiler to elide some stack copies unders certain
71/// conditions, because might be possible for `f()` to construct the object
72/// directly in place on the heap. Note that if `E` is
73/// not [`Infallible`](core::convert::Infallible), a stack copy might still be
74/// needed to unpeel a returned `Ok()` and extract the wrapped value. However,
75/// for huge objects it's been empirically observed that one out of up to two
76/// stack copies can usually get eliminated as compared to the "common" `Box`
77/// creation primitives.
78///
79/// # Arguments:
80///
81/// * `f` - Callback to invoke after memory allocation for obtaining the `Box`'
82///   wrapped value. May return an error, which would get propagated via
83///   [`TryNewWithError::With`] back to the caller.
84///
85/// # Errors:
86///
87/// * [`TryNewWithError::TryNew`] - The memory allocation has failed.
88/// * [`TryNewWithError::With`] - The `f` object factory callback returned an
89///   error.
90pub fn box_try_new_with<T, E, F: FnOnce() -> Result<T, E>>(f: F) -> Result<Box<T>, TryNewWithError<E>> {
91    let mut p = box_try_new::<mem::MaybeUninit<T>>(mem::MaybeUninit::uninit()).map_err(TryNewWithError::TryNew)?;
92    p.write(match f() {
93        Ok(v) => v,
94        Err(e) => {
95            return Err(TryNewWithError::With(e));
96        }
97    });
98    Ok(unsafe { p.assume_init() })
99}
100
101/// Convenience helper to allocate a default-initialized `Vec` of a given
102/// length, handling memory allocation failure gracefully.
103///
104/// # Arguments:
105///
106/// * `len` - The length to resize the `Vec` to.
107///
108/// # Errors:
109///
110/// * [`TryNewError::MemoryAllocationFailure`] - The memory allocation has
111///   failed.
112pub fn try_alloc_vec<T: Default + Clone>(len: usize) -> Result<Vec<T>, TryNewError> {
113    let mut v = Vec::new();
114    v.try_reserve_exact(len)
115        .map_err(|_| TryNewError::MemoryAllocationFailure)?;
116    v.resize(len, T::default());
117    Ok(v)
118}
119
120/// Convenience helper to allocate a default-initialized and
121/// [`Zeroizing`](zeroize::Zeroizing) wrapped `Vec` of a given length
122/// handling memory allocation failure gracefully.
123///
124/// # Arguments:
125///
126/// * `len` - The length to resize the `Vec` to.
127///
128/// # Errors:
129///
130/// * [`TryNewError::MemoryAllocationFailure`] - The memory allocation has
131///   failed.
132pub fn try_alloc_zeroizing_vec<T: zeroize::Zeroize + Default + Clone>(
133    len: usize,
134) -> Result<zeroize::Zeroizing<Vec<T>>, TryNewError> {
135    Ok(try_alloc_vec(len)?.into())
136}
cocoon_tpm_utils_common/alloc.rs

cocoon_tpm_utils_common/
alloc.rs
