1 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 // Copyright by contributors to this project. 3 // SPDX-License-Identifier: (Apache-2.0 OR MIT) 4 5 #[cfg(target_has_atomic = "ptr")] 6 use alloc::sync::Arc; 7 8 #[cfg(not(target_has_atomic = "ptr"))] 9 use portable_atomic_util::Arc; 10 11 use core::{ 12 convert::Infallible, 13 fmt::{self, Debug}, 14 }; 15 16 use alloc::vec::Vec; 17 use mls_rs_core::key_package::{KeyPackageData, KeyPackageStorage}; 18 19 #[cfg(feature = "std")] 20 use std::sync::{Mutex, MutexGuard}; 21 22 #[cfg(mls_build_async)] 23 use alloc::boxed::Box; 24 #[cfg(not(feature = "std"))] 25 use spin::{Mutex, MutexGuard}; 26 27 use crate::map::LargeMap; 28 29 #[derive(Clone, Default)] 30 /// In memory key package storage backed by a HashMap. 31 /// 32 /// All clones of an instance of this type share the same underlying HashMap. 33 pub struct InMemoryKeyPackageStorage { 34 inner: Arc<Mutex<LargeMap<Vec<u8>, KeyPackageData>>>, 35 } 36 37 impl Debug for InMemoryKeyPackageStorage { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result38 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 39 f.debug_struct("InMemoryKeyPackageStorage") 40 .field( 41 "inner", 42 &mls_rs_core::debug::pretty_with(|f| { 43 f.debug_map() 44 .entries( 45 self.lock() 46 .iter() 47 .map(|(k, v)| (mls_rs_core::debug::pretty_bytes(k), v)), 48 ) 49 .finish() 50 }), 51 ) 52 .finish() 53 } 54 } 55 56 impl InMemoryKeyPackageStorage { 57 /// Create an empty key package storage. new() -> Self58 pub fn new() -> Self { 59 Default::default() 60 } 61 62 /// Insert key package data. insert(&self, id: Vec<u8>, pkg: KeyPackageData)63 pub fn insert(&self, id: Vec<u8>, pkg: KeyPackageData) { 64 self.lock().insert(id, pkg); 65 } 66 67 /// Get a key package data by `id`. get(&self, id: &[u8]) -> Option<KeyPackageData>68 pub fn get(&self, id: &[u8]) -> Option<KeyPackageData> { 69 self.lock().get(id).cloned() 70 } 71 72 /// Delete key package data by `id`. delete(&self, id: &[u8])73 pub fn delete(&self, id: &[u8]) { 74 self.lock().remove(id); 75 } 76 77 /// Get all key packages that are currently stored. key_packages(&self) -> Vec<(Vec<u8>, KeyPackageData)>78 pub fn key_packages(&self) -> Vec<(Vec<u8>, KeyPackageData)> { 79 self.lock() 80 .iter() 81 .map(|(k, v)| (k.clone(), v.clone())) 82 .collect() 83 } 84 lock(&self) -> MutexGuard<'_, LargeMap<Vec<u8>, KeyPackageData>>85 fn lock(&self) -> MutexGuard<'_, LargeMap<Vec<u8>, KeyPackageData>> { 86 #[cfg(feature = "std")] 87 return self.inner.lock().unwrap(); 88 89 #[cfg(not(feature = "std"))] 90 return self.inner.lock(); 91 } 92 } 93 94 #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)] 95 #[cfg_attr(mls_build_async, maybe_async::must_be_async)] 96 impl KeyPackageStorage for InMemoryKeyPackageStorage { 97 type Error = Infallible; 98 delete(&mut self, id: &[u8]) -> Result<(), Self::Error>99 async fn delete(&mut self, id: &[u8]) -> Result<(), Self::Error> { 100 (*self).delete(id); 101 Ok(()) 102 } 103 insert(&mut self, id: Vec<u8>, pkg: KeyPackageData) -> Result<(), Self::Error>104 async fn insert(&mut self, id: Vec<u8>, pkg: KeyPackageData) -> Result<(), Self::Error> { 105 (*self).insert(id, pkg); 106 Ok(()) 107 } 108 get(&self, id: &[u8]) -> Result<Option<KeyPackageData>, Self::Error>109 async fn get(&self, id: &[u8]) -> Result<Option<KeyPackageData>, Self::Error> { 110 Ok(self.get(id)) 111 } 112 } 113