Change functions to take mutex (#15)

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "ironcore-documents"
 description = "A library for working with IronCore Labs documents and header formats."
-version = "0.1.1-pre"
+version = "0.2.0-pre"
 edition = "2021"
 license = "AGPL-3.0-only"
 documentation = "https://docs.rs/ironcore-documents"

rust-toolchain.toml

@@ -1,4 +1,4 @@
 [toolchain]
 profile = "default"
-channel = "1.75.0"
+channel = "1.77.2"
 components = ["rust-src", "rust-analyzer"]

src/aes.rs

@@ -1,8 +1,13 @@
 // This module is dedicated with things to do with aes encryption/decryption.
 use super::Error;
+use crate::take_lock;
 use aes_gcm::{aead::Aead, aead::Payload, AeadCore, Aes256Gcm, KeyInit, Nonce};
 use bytes::Bytes;
 use rand::{CryptoRng, RngCore};
+use std::{
+    ops::DerefMut,
+    sync::{Arc, Mutex},
+};
 
 type Result<T> = core::result::Result<T, super::Error>;
 pub(crate) const IV_LEN: usize = 12;
@@ -76,21 +81,26 @@ pub fn decrypt_document_with_attached_iv(
 
 /// Encrypt a document and put the iv on the front of it.
 pub fn encrypt_document_and_attach_iv<R: RngCore + CryptoRng>(
-    rng: &mut R,
+    rng: Arc<Mutex<R>>,
     key: EncryptionKey,
     document: PlaintextDocument,
 ) -> Result<IvAndCiphertext> {
-    let (iv, enc_data) = aes_encrypt(key, &document.0, &[], rng)?;
-    Ok(IvAndCiphertext([&iv[..], &enc_data.0[..]].concat().into()))
+    let (iv, mut enc_data) = aes_encrypt(key, &document.0, &[], rng)?;
+    let mut iv_vec = iv.to_vec();
+    iv_vec.append(&mut enc_data.0);
+    Ok(IvAndCiphertext(iv_vec.into()))
 }
 
 pub(crate) fn aes_encrypt<R: RngCore + CryptoRng>(
     key: EncryptionKey,
     plaintext: &[u8],
     associated_data: &[u8],
-    rng: &mut R,
+    rng: Arc<Mutex<R>>,
 ) -> Result<([u8; 12], EncryptedDocument)> {
-    let iv = Aes256Gcm::generate_nonce(rng);
+    let iv = {
+        let mut guard = take_lock(&rng);
+        Aes256Gcm::generate_nonce(&mut *guard.deref_mut())
+    };
     aes_encrypt_with_iv(key, plaintext, iv.into(), associated_data)
 }
 
@@ -149,22 +159,30 @@ mod test {
             "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
         ));
         let plaintext = hex!("112233445566778899aabbccddee");
-        let (iv, encrypt_result) =
-            aes_encrypt(key, &plaintext, &[], &mut rand::thread_rng()).unwrap();
+        let (iv, encrypt_result) = aes_encrypt(
+            key,
+            &plaintext,
+            &[],
+            Arc::new(Mutex::new(rand::thread_rng())),
+        )
+        .unwrap();
         let decrypt_result = aes_decrypt_core(&key, iv, &encrypt_result.0, &[]).unwrap();
         assert_eq!(decrypt_result, plaintext);
     }
 
     #[test]
     fn encrypt_decrypt_attached_roundtrip() {
-        let mut rng = ChaCha20Rng::seed_from_u64(13u64);
+        let rng = ChaCha20Rng::seed_from_u64(13u64);
         let key = EncryptionKey(hex!(
             "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
         ));
         let document = vec![1u8];
-        let encrypted =
-            encrypt_document_and_attach_iv(&mut rng, key, PlaintextDocument(document.clone()))
-                .unwrap();
+        let encrypted = encrypt_document_and_attach_iv(
+            Arc::new(Mutex::new(rng)),
+            key,
+            PlaintextDocument(document.clone()),
+        )
+        .unwrap();
         let result = decrypt_document_with_attached_iv(&key, &encrypted).unwrap();
         assert_eq!(result.0, document);
     }

src/lib.rs

@@ -4,7 +4,10 @@ pub mod v3;
 pub mod v4;
 pub mod v5;
 
-use std::fmt::{Display, Formatter, Result as DisplayResult};
+use std::{
+    fmt::{Display, Formatter, Result as DisplayResult},
+    sync::{Mutex, MutexGuard},
+};
 use thiserror::Error;
 use v5::key_id_header::KEY_ID_HEADER_LEN;
 
@@ -65,3 +68,28 @@ impl Display for Error {
         }
     }
 }
+
+/// Acquire mutex in a blocking fashion. If the Mutex is or becomes poisoned, write out an error
+/// message and panic.
+///
+/// The lock is released when the returned MutexGuard falls out of scope.
+///
+/// # Usage:
+/// single statement (mut)
+/// `let result = take_lock(&t).deref_mut().call_method_on_t();`
+///
+/// multi-statement (mut)
+/// ```ignore
+/// let t = T {};
+/// let result = {
+///     let g = &mut *take_lock(&t);
+///     g.call_method_on_t()
+/// }; // lock released here
+/// ```
+///
+pub fn take_lock<T>(m: &Mutex<T>) -> MutexGuard<T> {
+    m.lock().unwrap_or_else(|e| {
+        let error = format!("Error when acquiring lock: {e}");
+        panic!("{error}");
+    })
+}

src/v4/aes.rs

@@ -9,19 +9,23 @@ use crate::{
         },
         V4DocumentHeader,
     },
-    signing, Error,
+    signing, take_lock, Error,
 };
 use bytes::Bytes;
 use protobuf::Message;
 use rand::{CryptoRng, RngCore};
+use std::{
+    ops::DerefMut,
+    sync::{Arc, Mutex},
+};
 
 type Result<T> = core::result::Result<T, crate::Error>;
 
 /// If `maybe_dek` is None, generate a dek, otherwise use the one provided.
 /// Encrypt the dek using the kek to make an aes edek. The provided id will be put into the Aes256GcmEncryptedDek.
 /// Returns the dek and Aes256GcmEncryptedDek.
 fn generate_aes_edek<R: CryptoRng + RngCore>(
-    rng: &mut R,
+    rng: Arc<Mutex<R>>,
     kek: EncryptionKey,
     maybe_dek: Option<EncryptionKey>,
     id: &str,
@@ -31,7 +35,7 @@ fn generate_aes_edek<R: CryptoRng + RngCore>(
 )> {
     let dek = maybe_dek.unwrap_or_else(|| {
         let mut buffer = [0u8; 32];
-        rng.fill_bytes(&mut buffer);
+        take_lock(&rng).deref_mut().fill_bytes(&mut buffer);
         EncryptionKey(buffer)
     });
     let (iv, edek) = aes_encrypt(kek, &dek.0, &[], rng)?;
@@ -82,7 +86,7 @@ pub fn create_signed_proto(
 /// The edek will be placed into a V4DocumentHeader and the signature will be computed.
 /// The aes dek is the key used to compute the signature.
 pub fn generate_aes_edek_and_sign<R: CryptoRng + RngCore>(
-    rng: &mut R,
+    rng: Arc<Mutex<R>>,
     kek: EncryptionKey,
     maybe_dek: Option<EncryptionKey>,
     id: &str,
@@ -160,24 +164,26 @@ mod test {
     use rand_chacha::ChaCha20Rng;
     #[test]
     fn generate_aes_edek_decrypts() {
-        let mut rng = ChaCha20Rng::seed_from_u64(203u64);
+        let rng = ChaCha20Rng::seed_from_u64(203u64);
         let kek = EncryptionKey(hex!(
             "aabbccddeefaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
         ));
         let id = "hello";
-        let (aes_dek, aes_edek) = generate_aes_edek(&mut rng, kek, None, id).unwrap();
+        let (aes_dek, aes_edek) =
+            generate_aes_edek(Arc::new(Mutex::new(rng)), kek, None, id).unwrap();
         let result = decrypt_aes_edek(&kek, &aes_edek).unwrap();
         assert_eq!(result, aes_dek);
     }
 
     #[test]
     fn signed_aes_edek_verifies_and_decrypts() {
-        let mut rng = ChaCha20Rng::seed_from_u64(203u64);
+        let rng = ChaCha20Rng::seed_from_u64(203u64);
         let kek = EncryptionKey(hex!(
             "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
         ));
         let id = "hello";
-        let (aes_dek, v4_document) = generate_aes_edek_and_sign(&mut rng, kek, None, id).unwrap();
+        let (aes_dek, v4_document) =
+            generate_aes_edek_and_sign(Arc::new(Mutex::new(rng)), kek, None, id).unwrap();
         let aes_edek =
             v4_document.signed_payload.0.clone().unwrap().edeks[0].take_aes_256_gcm_edek();
         let decrypted_aes_dek = decrypt_aes_edek(&kek, &aes_edek).unwrap();
@@ -188,12 +194,13 @@ mod test {
 
     #[test]
     fn signed_aes_edek_decrypts() {
-        let mut rng = ChaCha20Rng::seed_from_u64(203u64);
+        let rng = ChaCha20Rng::seed_from_u64(203u64);
         let kek = EncryptionKey(hex!(
             "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
         ));
         let id = "hello";
-        let (aes_dek, v4_document) = generate_aes_edek_and_sign(&mut rng, kek, None, id).unwrap();
+        let (aes_dek, v4_document) =
+            generate_aes_edek_and_sign(Arc::new(Mutex::new(rng)), kek, None, id).unwrap();
         let aes_edek = v4_document.signed_payload.0.unwrap().edeks[0].take_aes_256_gcm_edek();
         let result = decrypt_aes_edek(&kek, &aes_edek).unwrap();
         assert_eq!(result, aes_dek);

src/v5/mod.rs

@@ -10,6 +10,7 @@ use crate::{
 use bytes::{Buf, Bytes};
 use key_id_header::KeyIdHeader;
 use rand::{CryptoRng, RngCore};
+use std::sync::{Arc, Mutex};
 
 // The V5 data format is defined by 2 data formats. One for the edek and one for encrypted data encrypted with that edek.
 // The edek format is a 6 byte key id (see the key_id_header module) followed by a V4DocumentHeader proto.
@@ -92,7 +93,7 @@ impl EncryptedPayload {
 /// Encrypt a document to be used as a detached document. This means it will have a header of `0IRON` as the first
 /// 5 bytes.
 pub fn encrypt_detached_document<R: RngCore + CryptoRng>(
-    rng: &mut R,
+    rng: Arc<Mutex<R>>,
     key: EncryptionKey,
     document: PlaintextDocument,
 ) -> Result<EncryptedPayload> {
@@ -122,12 +123,13 @@ mod test {
     use rand_chacha::ChaCha20Rng;
     #[test]
     fn encrypt_decrypt_detached_document_roundtrips() {
-        let mut rng = ChaCha20Rng::seed_from_u64(172u64);
+        let rng = ChaCha20Rng::seed_from_u64(172u64);
         let key = EncryptionKey(hex!(
             "fffefdfcfbfaf9f8f7f6f5f4f3f2f1f0f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
         ));
         let plaintext = PlaintextDocument(vec![100u8, 200u8]);
-        let encrypted = encrypt_detached_document(&mut rng, key, plaintext.clone()).unwrap();
+        let encrypted =
+            encrypt_detached_document(Arc::new(Mutex::new(rng)), key, plaintext.clone()).unwrap();
         let result = encrypted.decrypt(&key).unwrap();
         assert_eq!(result, plaintext);
     }