From 25e4935286a823a1f6bf4ebd964860f3e5cf05c5 Mon Sep 17 00:00:00 2001
From: Naomi Plasterer <naomi@xmtp.com>
Date: Fri, 10 May 2024 13:54:32 -0700
Subject: [PATCH] more updates on bytewrapper

---
 memory_storage/src/lib.rs          | 14 +++++++-------
 memory_storage/src/test_store.rs   |  1 -
 openmls/src/group/mls_group/mod.rs |  6 +++---
 openmls/src/storage.rs             | 13 +++----------
 traits/src/storage.rs              |  3 +--
 5 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/memory_storage/src/lib.rs b/memory_storage/src/lib.rs
index 5ed0c3852..e2968c0ba 100644
--- a/memory_storage/src/lib.rs
+++ b/memory_storage/src/lib.rs
@@ -801,7 +801,6 @@ impl StorageProvider<CURRENT_VERSION> for MemoryStorage {
     >(
         &self,
         group_id: &GroupId,
-        proposal_ref: &ProposalRef,
     ) -> Result<(), Self::Error> {
         let mut values = self.values.write().unwrap();
 
@@ -876,12 +875,13 @@ impl StorageProvider<CURRENT_VERSION> for MemoryStorage {
         &self,
         group_id: &GroupId,
     ) -> Result<ByteWrapper, Self::Error> {
-        let key = serde_json::to_vec(group_id)?;
-        match self.read(AAD_LABEL, &key) {
-            Ok(Some(data)) => Ok(data),
-            Ok(None) => Ok(ByteWrapper::from(Vec::new())),
-            Err(e) => Err(e),
-        }
+        let key: Vec<u8> = serde_json::to_vec(group_id)?;
+        let values = self.values.read().unwrap();
+
+        let value = values.get(&key).unwrap();
+        let value = serde_json::from_slice(value).unwrap();
+        
+        Ok(value)
     }
 
     fn write_aad<
diff --git a/memory_storage/src/test_store.rs b/memory_storage/src/test_store.rs
index 97f873c30..569ebdfa7 100644
--- a/memory_storage/src/test_store.rs
+++ b/memory_storage/src/test_store.rs
@@ -444,7 +444,6 @@ impl StorageProvider<V_TEST> for MemoryStorage {
     >(
         &self,
         _group_id: &GroupId,
-        proposal_ref: &ProposalRef,
     ) -> Result<(), Self::Error> {
         todo!()
     }
diff --git a/openmls/src/group/mls_group/mod.rs b/openmls/src/group/mls_group/mod.rs
index 153f3d5ab..15031a873 100644
--- a/openmls/src/group/mls_group/mod.rs
+++ b/openmls/src/group/mls_group/mod.rs
@@ -16,7 +16,7 @@ use crate::{
     storage::{OpenMlsProvider, StorageProvider},
     treesync::{node::leaf_node::LeafNode, RatchetTree},
 };
-use openmls_traits::types::Ciphersuite;
+use openmls_traits::{storage::traits::ByteWrapper, types::Ciphersuite};
 
 // Private
 mod application;
@@ -312,7 +312,7 @@ impl MlsGroup {
             self.proposal_store.empty();
 
             // Clear proposals in storage
-            storage.clear_proposal_queue(self.group_id(), &self.proposal_store)?;
+            storage.clear_proposal_queue::<GroupId, ProposalRef>(self.group_id())?;
         }
 
         Ok(())
@@ -371,7 +371,7 @@ impl MlsGroup {
     ) -> Result<(), StorageProvider::Error> {
         self.group.delete(storage)?;
         storage.delete_group_config(self.group_id())?;
-        storage.clear_proposal_queue(self.group_id(), &self.proposal_store)?;
+        storage.clear_proposal_queue::<GroupId, ProposalRef>(self.group_id())?;
         storage.delete_own_leaf_nodes(self.group_id())?;
         storage.delete_aad(self.group_id())?;
         storage.delete_group_state(self.group_id())?;
diff --git a/openmls/src/storage.rs b/openmls/src/storage.rs
index d25e8c71f..86b07e713 100644
--- a/openmls/src/storage.rs
+++ b/openmls/src/storage.rs
@@ -56,25 +56,18 @@ impl<
     type StorageError = Error;
 }
 
-// Implementations for the Entity and Key traits
-
 #[derive(Debug, Serialize, Deserialize)]
-struct ByteWrapper {
+pub struct ByteWrapper {
     data: Vec<u8>,
 }
 
 impl ByteWrapper {
-    fn from(data: Vec<u8>) -> Self {
-        ByteWrapper { data }
-    }
-}
-
-impl From<Vec<u8>> for ByteWrapper {
-    fn from(data: Vec<u8>) -> Self {
+    pub fn from(data: Vec<u8>) -> Self {
         ByteWrapper { data }
     }
 }
 
+// Implementations for the Entity and Key traits
 impl Entity<CURRENT_VERSION> for ByteWrapper {}
 impl traits::ByteWrapper<CURRENT_VERSION> for ByteWrapper {}
 
diff --git a/traits/src/storage.rs b/traits/src/storage.rs
index 4fad76ac3..eef0fa748 100644
--- a/traits/src/storage.rs
+++ b/traits/src/storage.rs
@@ -3,7 +3,7 @@
 //! implements the [`StorageProvider`] trait. The trait mostly defines getters and setters, but
 //! also a few methods that append to lists (which behave similar to setters).
 
-use serde::{de::DeserializeOwned, Serialize};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 /// The storage version used by OpenMLS
 pub const CURRENT_VERSION: u16 = 1;
 
@@ -537,7 +537,6 @@ pub trait StorageProvider<const VERSION: u16> {
     >(
         &self,
         group_id: &GroupId,
-        proposal_ref: &ProposalRef,
     ) -> Result<(), Self::Error>;
 
     //