age: Use rayon for processing STREAM chunks in parallel

str4d · Aug 7, 2021 · 6d8b2e2 · 6d8b2e2
1 parent d992e13
commit 6d8b2e2
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 18 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/age/Cargo.toml b/age/Cargo.toml
@@ -77,6 +77,7 @@ rust-embed = "5"
 
 # Performance
 num_cpus = "1.0"
+rayon = "1.5"
 
 # Common CLI dependencies
 console = { version = "0.14", optional = true }

diff --git a/age/src/primitives/stream.rs b/age/src/primitives/stream.rs
@@ -6,6 +6,7 @@ use chacha20poly1305::{
 };
 use lazy_static::lazy_static;
 use pin_project::pin_project;
+use rayon::prelude::*;
 use secrecy::{ExposeSecret, SecretVec};
 use std::cmp;
 use std::convert::TryInto;
@@ -53,9 +54,9 @@ impl Nonce {
         self.0 = u128::from(val) << 8;
     }
 
-    fn increment_counter(&mut self) {
+    fn increment_counter(&mut self, by: usize) {
         // Increment the 11-byte counter
-        self.0 += 1 << 8;
+        self.0 += (by as u128) << 8;
         if self.0 >> (8 * 12) != 0 {
             panic!("We overflowed the nonce!");
         }
@@ -196,26 +197,29 @@ impl Stream {
         let num_chunks = chunks.len();
         let mut encrypted = vec![0; chunks_len + TAG_SIZE * num_chunks];
 
-        for (i, (encrypted, chunk)) in encrypted
+        encrypted
             .chunks_mut(ENCRYPTED_CHUNK_SIZE)
             .zip(chunks)
             .enumerate()
-        {
-            if i + 1 == num_chunks {
-                self.nonce.set_last(last).unwrap();
-            }
+            .par_bridge()
+            .for_each_with(self.nonce, |nonce, (i, (encrypted, chunk))| {
+                nonce.increment_counter(i);
+                if i + 1 == num_chunks {
+                    nonce.set_last(last).unwrap();
+                }
 
-            let (buffer, tag) = encrypted.split_at_mut(chunk.len());
-            buffer.copy_from_slice(chunk);
-            tag.copy_from_slice(
-                self.aead
-                    .encrypt_in_place_detached(&self.nonce.to_bytes().into(), &[], buffer)
-                    .expect("we will never hit chacha20::MAX_BLOCKS because of the chunk size")
-                    .as_slice(),
-            );
+                let (buffer, tag) = encrypted.split_at_mut(chunk.len());
+                buffer.copy_from_slice(chunk);
+                tag.copy_from_slice(
+                    self.aead
+                        .encrypt_in_place_detached(&nonce.to_bytes().into(), &[], buffer)
+                        .expect("we will never hit chacha20::MAX_BLOCKS because of the chunk size")
+                        .as_slice(),
+                );
+            });
 
-            self.nonce.increment_counter();
-        }
+        self.nonce.increment_counter(num_chunks);
+        self.nonce.set_last(last).unwrap();
 
         Ok(encrypted)
     }
@@ -250,7 +254,7 @@ impl Stream {
                 )
                 .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "decryption error"))?;
 
-            self.nonce.increment_counter();
+            self.nonce.increment_counter(1);
         }
 
         Ok(SecretVec::new(decrypted))