Some additions needed to migrate the code

bitwarden · Oct 7, 2024 · 6ea6267 · 6ea6267
1 parent fdb0263
commit 6ea6267
Show file tree

Hide file tree

Showing 5 changed files with 173 additions and 55 deletions.
diff --git a/crates/bitwarden-crypto/benches/new_encryptable.rs b/crates/bitwarden-crypto/benches/new_encryptable.rs
@@ -1,9 +1,8 @@
-use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
-
 use bitwarden_crypto::{
     key_refs, service::*, CryptoError, EncString, Encryptable, KeyDecryptable, KeyEncryptable,
     SymmetricCryptoKey, UsesKey,
 };
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
 
 pub fn criterion_benchmark(c: &mut Criterion) {
     let user_key = SymmetricCryptoKey::generate(rand::thread_rng());
@@ -21,8 +20,14 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     let service: CryptoService<MySymmKeyRef, MyAsymmKeyRef> = CryptoService::new();
     #[allow(deprecated)]
     {
-        service.insert_symmetric_key(MySymmKeyRef::User, user_key.clone());
-        service.insert_symmetric_key(MySymmKeyRef::Organization(org_id), org_key.clone());
+        service
+            .context_mut()
+            .set_symmetric_key(MySymmKeyRef::User, user_key.clone())
+            .expect("");
+        service
+            .context_mut()
+            .set_symmetric_key(MySymmKeyRef::Organization(org_id), org_key.clone())
+            .expect("");
     }
 
     let cipher_views = vec![cipher_view.clone(); 10_000];

diff --git a/crates/bitwarden-crypto/src/keys/encryptable.rs b/crates/bitwarden-crypto/src/keys/encryptable.rs
@@ -86,6 +86,8 @@ pub trait Decryptable<
     ) -> Result<Output, crate::CryptoError>;
 }
 
+// Basic Encryptable/Decryptable implementations to and from bytes
+
 impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
     Decryptable<SymmKeyRef, AsymmKeyRef, SymmKeyRef, Vec<u8>> for EncString
 {
@@ -134,6 +136,8 @@ impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
     }
 }
 
+// Encryptable/Decryptable implementations to and from strings
+
 impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
     Decryptable<SymmKeyRef, AsymmKeyRef, SymmKeyRef, String> for EncString
 {
@@ -183,3 +187,103 @@ impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
         self.as_bytes().encrypt(ctx, key)
     }
 }
+
+impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
+    Encryptable<SymmKeyRef, AsymmKeyRef, SymmKeyRef, EncString> for String
+{
+    fn encrypt(
+        &self,
+        ctx: &mut CryptoServiceContext<SymmKeyRef, AsymmKeyRef>,
+        key: SymmKeyRef,
+    ) -> Result<EncString, crate::CryptoError> {
+        self.as_bytes().encrypt(ctx, key)
+    }
+}
+
+impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
+    Encryptable<SymmKeyRef, AsymmKeyRef, AsymmKeyRef, AsymmetricEncString> for String
+{
+    fn encrypt(
+        &self,
+        ctx: &mut CryptoServiceContext<SymmKeyRef, AsymmKeyRef>,
+        key: AsymmKeyRef,
+    ) -> Result<AsymmetricEncString, crate::CryptoError> {
+        self.as_bytes().encrypt(ctx, key)
+    }
+}
+
+// Generic implementations for Optional values
+
+impl<
+        SymmKeyRef: SymmetricKeyRef,
+        AsymmKeyRef: AsymmetricKeyRef,
+        Key: KeyRef,
+        T: Encryptable<SymmKeyRef, AsymmKeyRef, Key, Output>,
+        Output,
+    > Encryptable<SymmKeyRef, AsymmKeyRef, Key, Option<Output>> for Option<T>
+{
+    fn encrypt(
+        &self,
+        ctx: &mut CryptoServiceContext<SymmKeyRef, AsymmKeyRef>,
+        key: Key,
+    ) -> Result<Option<Output>, crate::CryptoError> {
+        self.as_ref()
+            .map(|value| value.encrypt(ctx, key))
+            .transpose()
+    }
+}
+
+impl<
+        SymmKeyRef: SymmetricKeyRef,
+        AsymmKeyRef: AsymmetricKeyRef,
+        Key: KeyRef,
+        T: Decryptable<SymmKeyRef, AsymmKeyRef, Key, Output>,
+        Output,
+    > Decryptable<SymmKeyRef, AsymmKeyRef, Key, Option<Output>> for Option<T>
+{
+    fn decrypt(
+        &self,
+        ctx: &mut CryptoServiceContext<SymmKeyRef, AsymmKeyRef>,
+        key: Key,
+    ) -> Result<Option<Output>, crate::CryptoError> {
+        self.as_ref()
+            .map(|value| value.decrypt(ctx, key))
+            .transpose()
+    }
+}
+
+// Generic implementations for Vec values
+
+impl<
+        SymmKeyRef: SymmetricKeyRef,
+        AsymmKeyRef: AsymmetricKeyRef,
+        Key: KeyRef,
+        T: Encryptable<SymmKeyRef, AsymmKeyRef, Key, Output>,
+        Output,
+    > Encryptable<SymmKeyRef, AsymmKeyRef, Key, Vec<Output>> for Vec<T>
+{
+    fn encrypt(
+        &self,
+        ctx: &mut CryptoServiceContext<SymmKeyRef, AsymmKeyRef>,
+        key: Key,
+    ) -> Result<Vec<Output>, crate::CryptoError> {
+        self.iter().map(|value| value.encrypt(ctx, key)).collect()
+    }
+}
+
+impl<
+        SymmKeyRef: SymmetricKeyRef,
+        AsymmKeyRef: AsymmetricKeyRef,
+        Key: KeyRef,
+        T: Decryptable<SymmKeyRef, AsymmKeyRef, Key, Output>,
+        Output,
+    > Decryptable<SymmKeyRef, AsymmKeyRef, Key, Vec<Output>> for Vec<T>
+{
+    fn decrypt(
+        &self,
+        ctx: &mut CryptoServiceContext<SymmKeyRef, AsymmKeyRef>,
+        key: Key,
+    ) -> Result<Vec<Output>, crate::CryptoError> {
+        self.iter().map(|value| value.decrypt(ctx, key)).collect()
+    }
+}
diff --git a/crates/bitwarden-crypto/src/service/context.rs b/crates/bitwarden-crypto/src/service/context.rs
@@ -1,9 +1,11 @@
 use std::sync::{RwLockReadGuard, RwLockWriteGuard};
 
 use rsa::Oaep;
+use zeroize::Zeroizing;
 
 use super::Keys;
 use crate::{
+    derive_shareable_key,
     service::{key_store::KeyStore, AsymmetricKeyRef, SymmetricKeyRef},
     AsymmetricCryptoKey, AsymmetricEncString, CryptoError, EncString, Result, SymmetricCryptoKey,
 };
@@ -84,6 +86,20 @@ impl<
         self.local_asymmetric_keys.clear();
     }
 
+    pub fn retain_symmetric_keys(&mut self, f: fn(SymmKeyRef) -> bool) {
+        if let Ok(keys) = self.global.get_mut() {
+            keys.symmetric_keys.retain(f);
+        }
+        self.local_symmetric_keys.retain(f);
+    }
+
+    pub fn retain_asymmetric_keys(&mut self, f: fn(AsymmKeyRef) -> bool) {
+        if let Ok(keys) = self.global.get_mut() {
+            keys.asymmetric_keys.retain(f);
+        }
+        self.local_asymmetric_keys.retain(f);
+    }
+
     /// TODO: All these encrypt x key with x key look like they need to be made generic,
     /// but I haven't found the best way to do that yet.
 
@@ -96,6 +112,7 @@ impl<
         let mut new_key_material =
             self.decrypt_data_with_symmetric_key(encryption_key, encrypted_key)?;
 
+        #[allow(deprecated)]
         self.set_symmetric_key(
             new_key_ref,
             SymmetricCryptoKey::try_from(new_key_material.as_mut_slice())?,
@@ -123,6 +140,7 @@ impl<
         let mut new_key_material =
             self.decrypt_data_with_asymmetric_key(encryption_key, encrypted_key)?;
 
+        #[allow(deprecated)]
         self.set_symmetric_key(
             new_key_ref,
             SymmetricCryptoKey::try_from(new_key_material.as_mut_slice())?,
@@ -150,6 +168,7 @@ impl<
         let new_key_material =
             self.decrypt_data_with_asymmetric_key(encryption_key, encrypted_key)?;
 
+        #[allow(deprecated)]
         self.set_asymmetric_key(
             new_key_ref,
             AsymmetricCryptoKey::from_der(&new_key_material)?,
@@ -181,12 +200,31 @@ impl<
         self.get_asymmetric_key(key_ref).is_ok()
     }
 
-    #[deprecated(note = "This function should never be used outside this crate")]
+    pub fn generate_symmetric_key(&mut self, key_ref: SymmKeyRef) -> Result<SymmKeyRef> {
+        let key = SymmetricCryptoKey::generate(rand::thread_rng());
+        #[allow(deprecated)]
+        self.set_symmetric_key(key_ref, key)?;
+        Ok(key_ref)
+    }
+
+    pub fn derive_shareable_key(
+        &mut self,
+        key_ref: SymmKeyRef,
+        secret: Zeroizing<[u8; 16]>,
+        name: &str,
+        info: Option<&str>,
+    ) -> Result<SymmKeyRef> {
+        #[allow(deprecated)]
+        self.set_symmetric_key(key_ref, derive_shareable_key(secret, name, info))?;
+        Ok(key_ref)
+    }
+
+    #[deprecated(note = "This function should ideally never be used outside this crate")]
     pub fn dangerous_get_symmetric_key(&self, key_ref: SymmKeyRef) -> Result<&SymmetricCryptoKey> {
         self.get_symmetric_key(key_ref)
     }
 
-    #[deprecated(note = "This function should never be used outside this crate")]
+    #[deprecated(note = "This function should ideally never be used outside this crate")]
     pub fn dangerous_get_asymmetric_key(
         &self,
         key_ref: AsymmKeyRef,
@@ -212,7 +250,12 @@ impl<
         .ok_or_else(|| crate::CryptoError::MissingKey2(format!("{key_ref:?}")))
     }
 
-    fn set_symmetric_key(&mut self, key_ref: SymmKeyRef, key: SymmetricCryptoKey) -> Result<()> {
+    #[deprecated(note = "This function should ideally never be used outside this crate")]
+    pub fn set_symmetric_key(
+        &mut self,
+        key_ref: SymmKeyRef,
+        key: SymmetricCryptoKey,
+    ) -> Result<()> {
         if key_ref.is_local() {
             self.local_symmetric_keys.insert(key_ref, key);
         } else {
@@ -221,7 +264,12 @@ impl<
         Ok(())
     }
 
-    fn set_asymmetric_key(&mut self, key_ref: AsymmKeyRef, key: AsymmetricCryptoKey) -> Result<()> {
+    #[deprecated(note = "This function should ideally never be used outside this crate")]
+    pub fn set_asymmetric_key(
+        &mut self,
+        key_ref: AsymmKeyRef,
+        key: AsymmetricCryptoKey,
+    ) -> Result<()> {
         if key_ref.is_local() {
             self.local_asymmetric_keys.insert(key_ref, key);
         } else {

diff --git a/crates/bitwarden-crypto/src/service/key_store/slice.rs b/crates/bitwarden-crypto/src/service/key_store/slice.rs
@@ -2,9 +2,8 @@ use std::marker::PhantomData;
 
 use zeroize::ZeroizeOnDrop;
 
-use crate::KeyRef;
-
 use super::KeyStore;
+use crate::KeyRef;
 
 /// This trait represents some data stored sequentially in memory, with a fixed size.
 /// We use this to abstract the implementation over Vec/Box<[u8]/NonNull<[u8]>, which

diff --git a/crates/bitwarden-crypto/src/service/mod.rs b/crates/bitwarden-crypto/src/service/mod.rs
@@ -1,17 +1,13 @@
 use std::sync::{Arc, RwLock};
 
-use crate::{
-    AsymmetricCryptoKey, AsymmetricKeyRef, Decryptable, Encryptable, KeyRef, SymmetricCryptoKey,
-    SymmetricKeyRef, UsesKey,
-};
+use crate::{AsymmetricKeyRef, Decryptable, Encryptable, KeyRef, SymmetricKeyRef, UsesKey};
 
 mod context;
 
 mod key_store;
 
 use context::ReadWriteGlobal;
 pub use context::{CryptoServiceContext, ReadOnlyGlobal};
-
 pub use key_store::create_key_store;
 use key_store::KeyStore;
 
@@ -54,40 +50,6 @@ impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
         keys.asymmetric_keys.clear();
     }
 
-    pub fn retain_symmetric_keys(&self, f: fn(SymmKeyRef) -> bool) {
-        self.key_stores
-            .write()
-            .expect("RwLock is poisoned")
-            .symmetric_keys
-            .retain(f);
-    }
-
-    pub fn retain_asymmetric_keys(&self, f: fn(AsymmKeyRef) -> bool) {
-        self.key_stores
-            .write()
-            .expect("RwLock is poisoned")
-            .asymmetric_keys
-            .retain(f);
-    }
-
-    #[deprecated(note = "We should be generating/decrypting the keys into the service directly")]
-    pub fn insert_symmetric_key(&self, key_ref: SymmKeyRef, key: SymmetricCryptoKey) {
-        self.key_stores
-            .write()
-            .expect("RwLock is poisoned")
-            .symmetric_keys
-            .insert(key_ref, key);
-    }
-
-    #[deprecated(note = "We should be generating/decrypting the keys into the service directly")]
-    pub fn insert_asymmetric_key(&self, key_ref: AsymmKeyRef, key: AsymmetricCryptoKey) {
-        self.key_stores
-            .write()
-            .expect("RwLock is poisoned")
-            .asymmetric_keys
-            .insert(key_ref, key);
-    }
-
     /// Initiate an encryption/decryption context. This context will have read only access to the
     /// global keys, and will have its own local key stores with read/write access. This
     /// context-local store will be cleared up when the context is dropped.
@@ -178,14 +140,14 @@ impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
         let res: Result<Vec<_>, _> = data
             .par_chunks(chunk_size)
             .map(|chunk| {
-                let mut context = self.context();
+                let mut ctx = self.context();
 
                 let mut result = Vec::with_capacity(chunk.len());
 
                 for item in chunk {
                     let key = item.uses_key();
-                    result.push(item.decrypt(&mut context, key));
-                    context.clear();
+                    result.push(item.decrypt(&mut ctx, key));
+                    ctx.clear();
                 }
 
                 result
@@ -215,14 +177,14 @@ impl<SymmKeyRef: SymmetricKeyRef, AsymmKeyRef: AsymmetricKeyRef>
         let res: Result<Vec<_>, _> = data
             .par_chunks(chunk_size)
             .map(|chunk| {
-                let mut context = self.context();
+                let mut ctx = self.context();
 
                 let mut result = Vec::with_capacity(chunk.len());
 
                 for item in chunk {
                     let key = item.uses_key();
-                    result.push(item.encrypt(&mut context, key));
-                    context.clear();
+                    result.push(item.encrypt(&mut ctx, key));
+                    ctx.clear();
                 }
 
                 result