flex_alloc_secure/
bytes.rs

1use core::mem::MaybeUninit;
2use core::slice;
3
4use flex_alloc::alloc::Allocator;
5use flex_alloc::boxed::Box;
6use flex_alloc::vec::{config::VecConfig, Vec};
7use rand_core::RngCore;
8
9/// Access a value as a mutable slice of bytes.
10///
11/// # Safety
12/// This trait must only be implemented for types which can accept any
13/// pattern of bits as a legitimate value.
14pub unsafe trait FillBytes {
15    /// Access the value as a mutable slice of bytes.
16    #[inline(always)]
17    fn as_bytes_mut(&mut self) -> &mut [u8] {
18        let len: usize = size_of_val(self);
19        unsafe { slice::from_raw_parts_mut(self as *mut Self as *mut u8, len) }
20    }
21
22    /// Fill the value with repeated bytes.
23    #[inline(always)]
24    fn fill_bytes(&mut self, value: u8) {
25        self.as_bytes_mut().fill(value);
26        self.canonicalize_bytes();
27    }
28
29    /// Fill the value with random bytes.
30    #[inline(always)]
31    fn fill_random(&mut self, mut rng: impl RngCore) {
32        rng.fill_bytes(self.as_bytes_mut());
33        self.canonicalize_bytes();
34    }
35
36    /// Fill the value using a closure.
37    #[inline(always)]
38    fn fill_with<R>(&mut self, f: impl FnOnce(&mut [u8]) -> R) -> R {
39        f(self.as_bytes_mut())
40    }
41
42    /// After filling the bytes, try to ensure consistency across platforms.
43    fn canonicalize_bytes(&mut self) {}
44}
45
46macro_rules! impl_fill_bytes_int {
47    ($name:ty) => {
48        unsafe impl FillBytes for $name {
49            #[inline(always)]
50            fn canonicalize_bytes(&mut self) {
51                *self = self.to_le();
52            }
53        }
54
55        unsafe impl FillBytes for Option<::core::num::NonZero<$name>> {
56            #[inline(always)]
57            fn canonicalize_bytes(&mut self) {
58                *self = ::core::num::NonZero::new(
59                    <$name>::from_ne_bytes(self.as_bytes_mut().try_into().unwrap()).to_le()
60                );
61            }
62        }
63
64        unsafe impl FillBytes for ::core::num::Saturating<$name> {
65            #[inline(always)]
66            fn canonicalize_bytes(&mut self) {
67                *self = ::core::num::Saturating(self.0.to_le());
68            }
69        }
70
71        unsafe impl FillBytes for ::core::num::Wrapping<$name> {
72            #[inline(always)]
73            fn canonicalize_bytes(&mut self) {
74                *self = ::core::num::Wrapping(self.0.to_le());
75            }
76        }
77    };
78    ($name:ty, $($rest:ty),+) => {
79        impl_fill_bytes_int!($name);
80        impl_fill_bytes_int!($($rest),+);
81    };
82}
83
84impl_fill_bytes_int!(u8, u16, u32, u64, u128, usize);
85impl_fill_bytes_int!(i8, i16, i32, i64, i128, isize);
86
87unsafe impl<T: Copy + FillBytes> FillBytes for MaybeUninit<T> {
88    #[inline(always)]
89    fn canonicalize_bytes(&mut self) {
90        unsafe { self.assume_init_mut() }.canonicalize_bytes();
91    }
92}
93
94unsafe impl<T: Copy + FillBytes, const N: usize> FillBytes for [T; N] {
95    #[inline(always)]
96    fn canonicalize_bytes(&mut self) {
97        for item in self.iter_mut() {
98            item.canonicalize_bytes();
99        }
100    }
101}
102
103unsafe impl<T: Copy + FillBytes> FillBytes for [T] {
104    #[inline(always)]
105    fn canonicalize_bytes(&mut self) {
106        for item in self.iter_mut() {
107            item.canonicalize_bytes();
108        }
109    }
110}
111
112unsafe impl<T: FillBytes + ?Sized, A: Allocator> FillBytes for Box<T, A> {
113    #[inline(always)]
114    fn as_bytes_mut(&mut self) -> &mut [u8] {
115        let len: usize = size_of_val(self.as_ref());
116        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr().cast::<u8>(), len) }
117    }
118
119    #[inline(always)]
120    fn canonicalize_bytes(&mut self) {
121        self.as_mut().canonicalize_bytes();
122    }
123}
124
125unsafe impl<T: FillBytes, C: VecConfig> FillBytes for Vec<T, C> {
126    #[inline(always)]
127    fn as_bytes_mut(&mut self) -> &mut [u8] {
128        let len: usize = size_of_val(self.as_ref());
129        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr().cast::<u8>(), len) }
130    }
131
132    #[inline(always)]
133    fn canonicalize_bytes(&mut self) {
134        for item in self.iter_mut() {
135            item.canonicalize_bytes();
136        }
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use crate::{alloc::UNINIT_ALLOC_BYTE, FillBytes};
143    use core::num::NonZero;
144
145    #[test]
146    fn nonzero_check() {
147        let mut a: Option<NonZero<u64>> = None;
148        a.fill_bytes(UNINIT_ALLOC_BYTE);
149        assert_eq!(
150            a.as_ref().copied(),
151            NonZero::new(u64::from_ne_bytes([UNINIT_ALLOC_BYTE; 8]).to_le())
152        );
153
154        let mut b = [Option::<NonZero<u64>>::None; 10];
155        b.fill_bytes(UNINIT_ALLOC_BYTE);
156        assert_eq!(
157            b[0].as_ref().copied(),
158            NonZero::new(u64::from_ne_bytes([UNINIT_ALLOC_BYTE; 8]).to_le())
159        );
160        b[..2].fill_bytes(UNINIT_ALLOC_BYTE);
161    }
162}
flex_alloc_secure/bytes.rs

flex_alloc_secure/
bytes.rs
