Use rayon to multithread encryption/decryption

bitwarden · Feb 6, 2024 · e515bc5 · e515bc5
1 parent 714c2d4
commit e515bc5
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/bitwarden-crypto/Cargo.toml b/crates/bitwarden-crypto/Cargo.toml
@@ -32,6 +32,7 @@ num-bigint = ">=0.4, <0.5"
 num-traits = ">=0.2.15, <0.3"
 pbkdf2 = { version = ">=0.12.1, <0.13", default-features = false }
 rand = ">=0.8.5, <0.9"
+rayon = ">=1.8.1, <2.0"
 rsa = ">=0.9.2, <0.10"
 schemars = { version = ">=0.8, <0.9", features = ["uuid1"] }
 serde = { version = ">=1.0, <2.0", features = ["derive"] }

diff --git a/crates/bitwarden-crypto/src/encryptable.rs b/crates/bitwarden-crypto/src/encryptable.rs
@@ -1,10 +1,11 @@
 use std::{collections::HashMap, hash::Hash};
 
+use rayon::prelude::*;
 use uuid::Uuid;
 
 use crate::{CryptoError, KeyDecryptable, KeyEncryptable, Result, SymmetricCryptoKey};
 
-pub trait KeyContainer {
+pub trait KeyContainer: Send + Sync {
     fn get_key(&self, org_id: &Option<Uuid>) -> Option<&SymmetricCryptoKey>;
 }
 
@@ -46,37 +47,48 @@ impl<T: KeyDecryptable<SymmetricCryptoKey, Output> + LocateKey, Output> Decrypta
     }
 }
 
-impl<T: Encryptable<Output>, Output> Encryptable<Vec<Output>> for Vec<T> {
+impl<T: Encryptable<Output> + Send + Sync, Output: Send + Sync> Encryptable<Vec<Output>>
+    for Vec<T>
+{
     fn encrypt(self, enc: &dyn KeyContainer, org_id: &Option<Uuid>) -> Result<Vec<Output>> {
-        self.into_iter().map(|e| e.encrypt(enc, org_id)).collect()
+        self.into_par_iter()
+            .map(|e| e.encrypt(enc, org_id))
+            .collect()
     }
 }
 
-impl<T: Decryptable<Output>, Output> Decryptable<Vec<Output>> for Vec<T> {
+impl<T: Decryptable<Output> + Send + Sync, Output: Send + Sync> Decryptable<Vec<Output>>
+    for Vec<T>
+{
     fn decrypt(&self, enc: &dyn KeyContainer, org_id: &Option<Uuid>) -> Result<Vec<Output>> {
-        self.iter().map(|e| e.decrypt(enc, org_id)).collect()
+        self.into_par_iter()
+            .map(|e| e.decrypt(enc, org_id))
+            .collect()
     }
 }
 
-impl<T: Encryptable<Output>, Output, Id: Hash + Eq> Encryptable<HashMap<Id, Output>>
-    for HashMap<Id, T>
+impl<T: Encryptable<Output> + Send + Sync, Output: Send + Sync, Id: Hash + Eq + Send + Sync>
+    Encryptable<HashMap<Id, Output>> for HashMap<Id, T>
 {
     fn encrypt(self, enc: &dyn KeyContainer, org_id: &Option<Uuid>) -> Result<HashMap<Id, Output>> {
-        self.into_iter()
+        self.into_par_iter()
             .map(|(id, e)| Ok((id, e.encrypt(enc, org_id)?)))
             .collect()
     }
 }
 
-impl<T: Decryptable<Output>, Output, Id: Hash + Eq + Copy> Decryptable<HashMap<Id, Output>>
-    for HashMap<Id, T>
+impl<
+        T: Decryptable<Output> + Send + Sync,
+        Output: Send + Sync,
+        Id: Hash + Eq + Copy + Send + Sync,
+    > Decryptable<HashMap<Id, Output>> for HashMap<Id, T>
 {
     fn decrypt(
         &self,
         enc: &dyn KeyContainer,
         org_id: &Option<Uuid>,
     ) -> Result<HashMap<Id, Output>> {
-        self.iter()
+        self.into_par_iter()
             .map(|(id, e)| Ok((*id, e.decrypt(enc, org_id)?)))
             .collect()
     }

diff --git a/crates/bitwarden-crypto/src/keys/key_encryptable.rs b/crates/bitwarden-crypto/src/keys/key_encryptable.rs
@@ -1,5 +1,7 @@
 use std::{collections::HashMap, hash::Hash};
 
+use rayon::prelude::*;
+
 use crate::error::Result;
 
 pub trait CryptoKey {}
@@ -44,37 +46,55 @@ impl<T: KeyDecryptable<Key, Output>, Key: CryptoKey, Output> KeyDecryptable<Key,
     }
 }
 
-impl<T: KeyEncryptable<Key, Output>, Key: CryptoKey, Output> KeyEncryptable<Key, Vec<Output>>
-    for Vec<T>
+impl<
+        T: KeyEncryptable<Key, Output> + Send + Sync,
+        Key: CryptoKey + Send + Sync,
+        Output: Send + Sync,
+    > KeyEncryptable<Key, Vec<Output>> for Vec<T>
 {
     fn encrypt_with_key(self, key: &Key) -> Result<Vec<Output>> {
-        self.into_iter().map(|e| e.encrypt_with_key(key)).collect()
+        self.into_par_iter()
+            .map(|e| e.encrypt_with_key(key))
+            .collect()
     }
 }
 
-impl<T: KeyDecryptable<Key, Output>, Key: CryptoKey, Output> KeyDecryptable<Key, Vec<Output>>
-    for Vec<T>
+impl<
+        T: KeyDecryptable<Key, Output> + Send + Sync,
+        Key: CryptoKey + Send + Sync,
+        Output: Send + Sync,
+    > KeyDecryptable<Key, Vec<Output>> for Vec<T>
 {
     fn decrypt_with_key(&self, key: &Key) -> Result<Vec<Output>> {
-        self.iter().map(|e| e.decrypt_with_key(key)).collect()
+        self.into_par_iter()
+            .map(|e| e.decrypt_with_key(key))
+            .collect()
     }
 }
 
-impl<T: KeyEncryptable<Key, Output>, Key: CryptoKey, Output, Id: Hash + Eq>
-    KeyEncryptable<Key, HashMap<Id, Output>> for HashMap<Id, T>
+impl<
+        T: KeyEncryptable<Key, Output> + Send + Sync,
+        Key: CryptoKey + Send + Sync,
+        Output: Send + Sync,
+        Id: Hash + Eq + Send + Sync,
+    > KeyEncryptable<Key, HashMap<Id, Output>> for HashMap<Id, T>
 {
     fn encrypt_with_key(self, key: &Key) -> Result<HashMap<Id, Output>> {
-        self.into_iter()
+        self.into_par_iter()
             .map(|(id, e)| Ok((id, e.encrypt_with_key(key)?)))
             .collect()
     }
 }
 
-impl<T: KeyDecryptable<Key, Output>, Key: CryptoKey, Output, Id: Hash + Eq + Copy>
-    KeyDecryptable<Key, HashMap<Id, Output>> for HashMap<Id, T>
+impl<
+        T: KeyDecryptable<Key, Output> + Send + Sync,
+        Key: CryptoKey + Send + Sync,
+        Output: Send + Sync,
+        Id: Hash + Eq + Copy + Send + Sync,
+    > KeyDecryptable<Key, HashMap<Id, Output>> for HashMap<Id, T>
 {
     fn decrypt_with_key(&self, key: &Key) -> Result<HashMap<Id, Output>> {
-        self.iter()
+        self.into_par_iter()
             .map(|(id, e)| Ok((*id, e.decrypt_with_key(key)?)))
             .collect()
     }