1 // Copyright 2024 Google LLC 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not 4 // use this file except in compliance with the License. You may obtain a copy of 5 // the License at 6 // 7 // https://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, WITHOUT 11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12 // License for the specific language governing permissions and limitations under 13 // the License. 14 15 //! Types and functions to help with memory buffer management. 16 17 use crate::constants::*; 18 use crate::error::{DpeResult, ErrCode}; 19 use heapless::Vec; 20 use zeroize::ZeroizeOnDrop; 21 22 /// Creates a byte array wrapper type for the sake of precise typing. 23 #[macro_export] 24 macro_rules! byte_array_wrapper { 25 ($type_name:ident, $len:ident, $desc:expr) => { 26 #[doc = "A byte array wrapper to represent a "] 27 #[doc = $desc] 28 #[doc = "."] 29 #[derive(Clone, Debug, Eq, PartialEq, Hash, ZeroizeOnDrop)] 30 pub struct $type_name([u8; $len]); 31 impl $type_name { 32 #[doc = "Returns the length of the array."] 33 pub fn len(&self) -> usize { 34 self.0.len() 35 } 36 37 #[doc = "Whether the array is empty."] 38 pub fn is_empty(&self) -> bool { 39 self.0.is_empty() 40 } 41 42 #[doc = "Borrows the array as a slice."] 43 pub fn as_slice(&self) -> &[u8] { 44 self.0.as_slice() 45 } 46 47 #[doc = "Mutably borrows the array as a slice."] 48 pub fn as_mut_slice(&mut self) -> &mut [u8] { 49 &mut self.0 50 } 51 52 #[doc = "Borrows the array."] 53 pub fn as_array(&self) -> &[u8; $len] { 54 &self.0 55 } 56 57 #[doc = "Creates a "] 58 #[doc = stringify!($type_name)] 59 #[doc = " from a slice. Fails if the slice length is not "] 60 #[doc = stringify!($len)] 61 #[doc = "."] 62 pub fn from_slice(s: &[u8]) -> DpeResult<Self> { 63 Self::try_from(s) 64 } 65 66 #[doc = "Creates a "] 67 #[doc = stringify!($type_name)] 68 #[doc = " from a slice infallibly. If the length of the slice is less than "] 69 #[doc = stringify!($len)] 70 #[doc = ", the remainder of the array is the default value. If the length "] 71 #[doc = "of the slice is more than "] 72 #[doc = stringify!($len)] 73 #[doc = ", only the first "] 74 #[doc = stringify!($len)] 75 #[doc = " bytes are used. This method is infallible."] 76 pub fn from_slice_infallible(value: &[u8]) -> Self { 77 #![allow(clippy::indexing_slicing)] 78 let mut tmp: Self = Default::default(); 79 if value.len() < $len { 80 tmp.0[..value.len()].copy_from_slice(value); 81 } else { 82 tmp.0.copy_from_slice(&value[..$len]); 83 } 84 tmp 85 } 86 87 #[doc = "Creates a "] 88 #[doc = stringify!($type_name)] 89 #[doc = " from an array."] 90 pub fn from_array(value: &[u8; $len]) -> Self { 91 Self(*value) 92 } 93 } 94 95 impl Default for $type_name { 96 fn default() -> Self { 97 Self([0; $len]) 98 } 99 } 100 101 impl TryFrom<&[u8]> for $type_name { 102 type Error = $crate::error::ErrCode; 103 104 fn try_from(value: &[u8]) -> Result<Self, Self::Error> { 105 value.try_into().map(Self).map_err(|_| { 106 log::error!("Invalid length for fixed length value: {}", $desc); 107 $crate::error::ErrCode::InvalidArgument 108 }) 109 } 110 } 111 112 impl From<[u8; $len]> for $type_name { 113 fn from(value: [u8; $len]) -> Self { 114 Self(value) 115 } 116 } 117 }; 118 } 119 120 /// Wraps a [heapless::Vec] of bytes and provides various convenience methods 121 /// that are useful when processing DPE messages. The inner `vec` is also 122 /// accessible directly. 123 #[derive(Clone, Debug, Default, Eq, PartialEq, Hash, ZeroizeOnDrop)] 124 pub struct SizedMessage<const S: usize> { 125 /// The wrapped Vec. 126 pub vec: Vec<u8, S>, 127 } 128 129 impl<const S: usize> SizedMessage<S> { 130 /// Creates a new, empty instance. 131 /// 132 /// # Example 133 /// 134 /// ```rust 135 /// type MyMessage = dpe_rs::memory::SizedMessage<200>; 136 /// 137 /// assert_eq!(MyMessage::new().len(), 0); 138 /// ``` new() -> Self139 pub fn new() -> Self { 140 Default::default() 141 } 142 143 /// Creates a new instance from a slice. 144 /// 145 /// # Errors 146 /// 147 /// If `value` exceeds the available capacity, returns an OutOfMemory error. 148 /// 149 /// # Example 150 /// 151 /// ```rust 152 /// type MyMessage = dpe_rs::memory::SizedMessage<200>; 153 /// 154 /// assert_eq!(MyMessage::from_slice(&[0; 12]).unwrap().as_slice(), &[0; 12]); 155 /// ``` from_slice(value: &[u8]) -> DpeResult<Self>156 pub fn from_slice(value: &[u8]) -> DpeResult<Self> { 157 Ok(Self { 158 vec: Vec::from_slice(value).map_err(|_| ErrCode::OutOfMemory)?, 159 }) 160 } 161 162 /// Clones `slice`, replacing any existing content. 163 /// 164 /// # Errors 165 /// 166 /// If `slice` exceeds the available capacity, returns an OutOfMemory error. 167 /// 168 /// # Example 169 /// 170 /// ```rust 171 /// type MyMessage = dpe_rs::memory::SizedMessage<200>; 172 /// 173 /// let mut m = MyMessage::from_slice(&[0; 12]).unwrap(); 174 /// m.clone_from_slice(&[1; 3]).unwrap(); 175 /// assert_eq!(m.as_slice(), &[1; 3]); 176 /// ``` clone_from_slice(&mut self, slice: &[u8]) -> DpeResult<()>177 pub fn clone_from_slice(&mut self, slice: &[u8]) -> DpeResult<()> { 178 self.clear(); 179 self.vec.extend_from_slice(slice).map_err(|_| ErrCode::OutOfMemory) 180 } 181 182 /// Borrows the inner byte array. as_slice(&self) -> &[u8]183 pub fn as_slice(&self) -> &[u8] { 184 self.vec.as_slice() 185 } 186 187 /// Mutably borrows the inner byte array after resizing. This is useful when 188 /// using the type as an output buffer. 189 /// 190 /// # Errors 191 /// 192 /// If `size` exceeds the available capacity, returns an OutOfMemory error. 193 /// 194 /// # Example 195 /// 196 /// ```rust 197 /// use rand_core::{RngCore, SeedableRng}; 198 /// 199 /// type MyMessage = dpe_rs::memory::SizedMessage<200>; 200 /// 201 /// let mut buffer = MyMessage::new(); 202 /// <rand_chacha::ChaCha12Rng as SeedableRng>::seed_from_u64(0) 203 /// .fill_bytes(buffer.as_mut_sized(100).unwrap()); 204 /// assert_eq!(buffer.len(), 100); 205 /// ``` as_mut_sized(&mut self, size: usize) -> DpeResult<&mut [u8]>206 pub fn as_mut_sized(&mut self, size: usize) -> DpeResult<&mut [u8]> { 207 self.vec.resize_default(size).map_err(|_| ErrCode::OutOfMemory)?; 208 Ok(self.vec.as_mut()) 209 } 210 211 /// Returns the length of the inner vec. len(&self) -> usize212 pub fn len(&self) -> usize { 213 self.vec.len() 214 } 215 216 /// Whether the inner vec is empty. is_empty(&self) -> bool217 pub fn is_empty(&self) -> bool { 218 self.vec.is_empty() 219 } 220 221 /// Clears the inner vec. clear(&mut self)222 pub fn clear(&mut self) { 223 self.vec.clear() 224 } 225 226 /// Removes the first `prefix_size` bytes from the message. This carries the 227 /// cost of moving the remaining bytes to the front of the buffer. 228 /// 229 /// # Errors 230 /// 231 /// If `prefix_size` is larger than the current length, returns an 232 /// InternalError error. 233 /// 234 /// # Example 235 /// 236 /// ```rust 237 /// type MyMessage = dpe_rs::memory::SizedMessage<200>; 238 /// 239 /// let mut m = MyMessage::from_slice("prefixdata".as_bytes()).unwrap(); 240 /// m.remove_prefix(6).unwrap(); 241 /// assert_eq!(m.as_slice(), "data".as_bytes()); 242 /// ``` remove_prefix(&mut self, prefix_size: usize) -> DpeResult<()>243 pub fn remove_prefix(&mut self, prefix_size: usize) -> DpeResult<()> { 244 if prefix_size > self.len() { 245 return Err(ErrCode::InternalError); 246 } 247 if prefix_size == self.len() { 248 self.clear(); 249 } else if prefix_size > 0 { 250 let slice: &mut [u8] = self.vec.as_mut(); 251 slice.copy_within(prefix_size.., 0); 252 self.vec.truncate(self.len() - prefix_size); 253 } 254 Ok(()) 255 } 256 257 /// Inserts `prefix` at the start of the message. This carries the cost of 258 /// moving the existing bytes to make room for the prefix. 259 /// 260 /// # Errors 261 /// 262 /// If inserting `prefix` overflows the available capacity, returns an 263 /// OutOfMemory error. 264 /// 265 /// # Example 266 /// 267 /// ```rust 268 /// type MyMessage = dpe_rs::memory::SizedMessage<200>; 269 /// 270 /// let mut m = MyMessage::from_slice("data".as_bytes()).unwrap(); 271 /// m.insert_prefix("prefix".as_bytes()).unwrap(); 272 /// assert_eq!(m.as_slice(), "prefixdata".as_bytes()); 273 /// ``` insert_prefix(&mut self, prefix: &[u8]) -> DpeResult<()>274 pub fn insert_prefix(&mut self, prefix: &[u8]) -> DpeResult<()> { 275 let old_len = self.len(); 276 self.vec 277 .resize_default(self.len() + prefix.len()) 278 .map_err(|_| ErrCode::OutOfMemory)?; 279 let slice: &mut [u8] = self.vec.as_mut(); 280 slice.copy_within(0..old_len, prefix.len()); 281 slice 282 .get_mut(..prefix.len()) 283 .ok_or(ErrCode::InternalError)? 284 .copy_from_slice(prefix); 285 Ok(()) 286 } 287 } 288 289 /// Represents a DPE command/response message. This type is large and should not 290 /// be instantiated unnecessarily. 291 pub type Message = SizedMessage<MAX_MESSAGE_SIZE>; 292