[stream] use buffer pool (#3033)

Co-authored-by: Patrick O'Grady <me@patrickogrady.xyz>

Author: andresilva

cryptography/src/handshake.rs

@@ -49,7 +49,7 @@ mod key_exchange;
 use key_exchange::{EphemeralPublicKey, SecretKey};
 
 mod cipher;
-pub use cipher::{RecvCipher, SendCipher, CIPHERTEXT_OVERHEAD};
+pub use cipher::{RecvCipher, SendCipher, TAG_SIZE};
 
 #[cfg(all(test, feature = "arbitrary"))]
 mod conformance;

cryptography/src/handshake/cipher.rs

@@ -4,9 +4,11 @@ use rand_core::CryptoRngCore;
 use std::vec::Vec;
 use zeroize::Zeroizing;
 
-/// The amount of overhead in a ciphertext, compared to the plain message.
-/// ChaCha20-Poly1305 uses a 128-bit (16 byte) authentication tag.
-pub const CIPHERTEXT_OVERHEAD: usize = 16;
+/// Size of the ChaCha20-Poly1305 authentication tag.
+///
+/// This tag is the overhead added to each ciphertext and must be transmitted
+/// alongside it for the receiver to verify integrity and authenticity.
+pub const TAG_SIZE: usize = 16;
 
 /// How many bytes are in a nonce.
 /// ChaCha20-Poly1305 uses a 96-bit (12 byte) nonce.
@@ -55,36 +57,33 @@ cfg_if::cfg_if! {
                 Self(LessSafeKey::new(unbound_key))
             }
 
-            fn encrypt(
+            fn encrypt_in_place(
                 &self,
                 nonce: &[u8; NONCE_SIZE_BYTES],
-                data: &[u8],
-            ) -> Result<Vec<u8>, Error> {
+                data: &mut [u8],
+            ) -> Result<[u8; TAG_SIZE], Error> {
                 let nonce = aead::Nonce::assume_unique_for_key(*nonce);
-                let mut scratch = Vec::with_capacity(data.len() + CIPHERTEXT_OVERHEAD);
-                scratch.extend_from_slice(data);
-                self.0
-                    .seal_in_place_append_tag(nonce, aead::Aad::empty(), &mut scratch)
+                let tag = self
+                    .0
+                    .seal_in_place_separate_tag(nonce, aead::Aad::empty(), data)
                     .map_err(|_| Error::EncryptionFailed)?;
-                Ok(scratch)
+                Ok(tag.as_ref().try_into().expect("tag size mismatch"))
             }
 
-            fn decrypt(
+            fn decrypt_in_place(
                 &self,
                 nonce: &[u8; NONCE_SIZE_BYTES],
-                data: &[u8],
-            ) -> Result<Vec<u8>, Error> {
+                data: &mut [u8],
+            ) -> Result<usize, Error> {
                 let nonce = aead::Nonce::assume_unique_for_key(*nonce);
-                let mut scratch = data.to_vec();
                 self.0
-                    .open_in_place(nonce, aead::Aad::empty(), &mut scratch)
+                    .open_in_place(nonce, aead::Aad::empty(), data)
                     .map_err(|_| Error::DecryptionFailed)?;
-                scratch.truncate(data.len() - CIPHERTEXT_OVERHEAD);
-                Ok(scratch)
+                Ok(data.len() - TAG_SIZE)
             }
         }
     } else {
-        use chacha20poly1305::{aead::Aead, ChaCha20Poly1305, KeyInit as _};
+        use chacha20poly1305::{aead::AeadInPlace, ChaCha20Poly1305, KeyInit as _};
 
         struct Cipher(ChaCha20Poly1305);
 
@@ -93,24 +92,36 @@ cfg_if::cfg_if! {
                 Self(ChaCha20Poly1305::new(key.into()))
             }
 
-            fn encrypt(
+            fn encrypt_in_place(
                 &self,
                 nonce: &[u8; NONCE_SIZE_BYTES],
-                data: &[u8],
-            ) -> Result<Vec<u8>, Error> {
-                self.0
-                    .encrypt(nonce.into(), data)
-                    .map_err(|_| Error::EncryptionFailed)
+                data: &mut [u8],
+            ) -> Result<[u8; TAG_SIZE], Error> {
+                let tag = self
+                    .0
+                    .encrypt_in_place_detached(nonce.into(), &[], data)
+                    .map_err(|_| Error::EncryptionFailed)?;
+                Ok(tag.into())
             }
 
-            fn decrypt(
+            fn decrypt_in_place(
                 &self,
                 nonce: &[u8; NONCE_SIZE_BYTES],
-                data: &[u8],
-            ) -> Result<Vec<u8>, Error> {
+                data: &mut [u8],
+            ) -> Result<usize, Error> {
+                let plaintext_len = data.len() - TAG_SIZE;
+                let tag: [u8; TAG_SIZE] = data[plaintext_len..]
+                    .try_into()
+                    .map_err(|_| Error::DecryptionFailed)?;
                 self.0
-                    .decrypt(nonce.into(), data)
-                    .map_err(|_| Error::DecryptionFailed)
+                    .decrypt_in_place_detached(
+                        nonce.into(),
+                        &[],
+                        &mut data[..plaintext_len],
+                        &tag.into(),
+                    )
+                    .map_err(|_| Error::DecryptionFailed)?;
+                Ok(plaintext_len)
             }
         }
     }
@@ -133,10 +144,23 @@ impl SendCipher {
         }
     }
 
+    /// Encrypts `data` in-place and returns the authentication tag.
+    ///
+    /// The caller is responsible for appending the returned tag to the buffer.
+    #[inline]
+    pub fn send_in_place(&mut self, data: &mut [u8]) -> Result<[u8; TAG_SIZE], Error> {
+        let nonce = self.nonce.inc()?;
+        self.inner
+            .expose(|cipher| cipher.encrypt_in_place(&nonce, data))
+    }
+
     /// Encrypts data and returns the ciphertext.
     pub fn send(&mut self, data: &[u8]) -> Result<Vec<u8>, Error> {
-        let nonce = self.nonce.inc()?;
-        self.inner.expose(|cipher| cipher.encrypt(&nonce, data))
+        let mut buf = vec![0u8; data.len() + TAG_SIZE];
+        buf[..data.len()].copy_from_slice(data);
+        let tag = self.send_in_place(&mut buf[..data.len()])?;
+        buf[data.len()..].copy_from_slice(&tag);
+        Ok(buf)
     }
 }
 
@@ -157,6 +181,33 @@ impl RecvCipher {
         }
     }
 
+    /// Decrypts `encrypted_data` in-place and returns the plaintext length.
+    ///
+    /// The buffer must contain ciphertext with the authentication tag appended
+    /// (last `TAG_SIZE` bytes). After decryption, the plaintext is in
+    /// `encrypted_data[..returned_len]`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - `encrypted_data.len() < TAG_SIZE`
+    /// - Too many messages have been received with this cipher
+    /// - The ciphertext was corrupted or tampered with
+    ///
+    /// In the last two cases, the `RecvCipher` will no longer be able to return
+    /// valid ciphertexts, and will always return an error on subsequent calls
+    /// to [`Self::recv`]. Terminating (and optionally reestablishing) the connection
+    /// is a simple (and safe) way to handle this scenario.
+    #[inline]
+    pub fn recv_in_place(&mut self, encrypted_data: &mut [u8]) -> Result<usize, Error> {
+        if encrypted_data.len() < TAG_SIZE {
+            return Err(Error::DecryptionFailed);
+        }
+        let nonce = self.nonce.inc()?;
+        self.inner
+            .expose(|cipher| cipher.decrypt_in_place(&nonce, encrypted_data))
+    }
+
     /// Decrypts ciphertext and returns the original data.
     ///
     /// # Errors
@@ -171,9 +222,10 @@ impl RecvCipher {
     /// to [`Self::recv`]. Terminating (and optionally reestablishing) the connection
     /// is a simple (and safe) way to handle this scenario.
     pub fn recv(&mut self, encrypted_data: &[u8]) -> Result<Vec<u8>, Error> {
-        let nonce = self.nonce.inc()?;
-        self.inner
-            .expose(|cipher| cipher.decrypt(&nonce, encrypted_data))
+        let mut buf = encrypted_data.to_vec();
+        let plaintext_len = self.recv_in_place(&mut buf)?;
+        buf.truncate(plaintext_len);
+        Ok(buf)
     }
 }
 
@@ -189,7 +241,7 @@ mod tests {
 
         let plaintext = b"hello world";
         let ciphertext = send.send(plaintext).unwrap();
-        assert_eq!(ciphertext.len(), plaintext.len() + CIPHERTEXT_OVERHEAD);
+        assert_eq!(ciphertext.len(), plaintext.len() + TAG_SIZE);
 
         let decrypted = recv.recv(&ciphertext).unwrap();
         assert_eq!(decrypted, plaintext);
@@ -211,7 +263,7 @@ mod tests {
     fn test_recv_ciphertext_too_short() {
         let mut rng = test_rng();
         let mut recv = RecvCipher::new(&mut rng);
-        let short_data = vec![0u8; CIPHERTEXT_OVERHEAD - 1];
+        let short_data = vec![0u8; TAG_SIZE - 1];
         assert!(matches!(
             recv.recv(&short_data),
             Err(Error::DecryptionFailed)
@@ -222,7 +274,70 @@ mod tests {
     fn test_recv_ciphertext_exactly_overhead() {
         let mut rng = test_rng();
         let mut recv = RecvCipher::new(&mut rng);
-        let tag_only = vec![0u8; CIPHERTEXT_OVERHEAD];
+        let tag_only = vec![0u8; TAG_SIZE];
         assert!(matches!(recv.recv(&tag_only), Err(Error::DecryptionFailed)));
     }
+
+    #[test]
+    fn test_send_recv_in_place_roundtrip() {
+        let mut send = SendCipher::new(&mut test_rng());
+        let mut recv = RecvCipher::new(&mut test_rng());
+
+        let plaintext = b"hello world";
+        let mut buf = vec![0u8; plaintext.len() + TAG_SIZE];
+        buf[..plaintext.len()].copy_from_slice(plaintext);
+
+        // Encrypt plaintext in place, get tag back
+        let tag = send.send_in_place(&mut buf[..plaintext.len()]).unwrap();
+        // Append tag to buffer
+        buf[plaintext.len()..].copy_from_slice(&tag);
+
+        // Decrypt ciphertext+tag in place, get plaintext length back
+        let plaintext_len = recv.recv_in_place(&mut buf).unwrap();
+
+        assert_eq!(plaintext_len, plaintext.len());
+        assert_eq!(&buf[..plaintext_len], plaintext);
+    }
+
+    #[test]
+    fn test_recv_in_place_ciphertext_too_short() {
+        let mut recv = RecvCipher::new(&mut test_rng());
+
+        // Buffer smaller than tag size
+        let mut buf = vec![0u8; TAG_SIZE - 1];
+        assert!(matches!(
+            recv.recv_in_place(&mut buf),
+            Err(Error::DecryptionFailed)
+        ));
+    }
+
+    #[test]
+    fn test_send_in_place_recv_compatibility() {
+        let mut send = SendCipher::new(&mut test_rng());
+        let mut recv = RecvCipher::new(&mut test_rng());
+
+        let plaintext = b"cross-api test";
+        let mut buf = vec![0u8; plaintext.len() + TAG_SIZE];
+        buf[..plaintext.len()].copy_from_slice(plaintext);
+
+        let tag = send.send_in_place(&mut buf[..plaintext.len()]).unwrap();
+        buf[plaintext.len()..].copy_from_slice(&tag);
+
+        // Use allocating recv on in-place encrypted data
+        let decrypted = recv.recv(&buf).unwrap();
+        assert_eq!(decrypted, plaintext);
+    }
+
+    #[test]
+    fn test_send_recv_in_place_compatibility() {
+        let mut send = SendCipher::new(&mut test_rng());
+        let mut recv = RecvCipher::new(&mut test_rng());
+
+        let plaintext = b"cross-api test";
+        let mut ciphertext = send.send(plaintext).unwrap();
+
+        // Use in-place recv on allocating send data
+        let plaintext_len = recv.recv_in_place(&mut ciphertext).unwrap();
+        assert_eq!(&ciphertext[..plaintext_len], plaintext);
+    }
 }

p2p/src/authenticated/discovery/actors/dialer.rs

@@ -14,7 +14,8 @@ use crate::authenticated::{
 use commonware_cryptography::Signer;
 use commonware_macros::select_loop;
 use commonware_runtime::{
-    spawn_cell, Clock, ContextCell, Handle, Metrics, Network, Resolver, SinkOf, Spawner, StreamOf,
+    spawn_cell, BufferPooler, Clock, ContextCell, Handle, Metrics, Network, Resolver, SinkOf,
+    Spawner, StreamOf,
 };
 use commonware_stream::encrypted::{dial, Config as StreamConfig};
 use commonware_utils::SystemTimeExt;
@@ -64,7 +65,11 @@ pub struct Actor<E: Spawner + Clock + Network + Resolver + Metrics, C: Signer> {
     attempts: Family<metrics::Peer, Counter>,
 }
 
-impl<E: Spawner + Clock + Network + Resolver + CryptoRngCore + Metrics, C: Signer> Actor<E, C> {
+impl<
+        E: Spawner + BufferPooler + Clock + Network + Resolver + CryptoRngCore + Metrics,
+        C: Signer,
+    > Actor<E, C>
+{
     pub fn new(context: E, cfg: Config<C>) -> Self {
         let attempts = Family::<metrics::Peer, Counter>::default();
         context.register(

p2p/src/authenticated/discovery/actors/listener.rs

@@ -8,8 +8,8 @@ use crate::authenticated::{
 use commonware_cryptography::Signer;
 use commonware_macros::select_loop;
 use commonware_runtime::{
-    spawn_cell, Clock, ContextCell, Handle, KeyedRateLimiter, Listener, Metrics, Network, Quota,
-    SinkOf, Spawner, StreamOf,
+    spawn_cell, BufferPooler, Clock, ContextCell, Handle, KeyedRateLimiter, Listener, Metrics,
+    Network, Quota, SinkOf, Spawner, StreamOf,
 };
 use commonware_stream::encrypted::{listen, Config as StreamConfig};
 use commonware_utils::{concurrency::Limiter, net::SubnetMask, IpAddrExt};
@@ -34,7 +34,7 @@ pub struct Config<C: Signer> {
     pub allowed_handshake_rate_per_subnet: Quota,
 }
 
-pub struct Actor<E: Spawner + Clock + Network + CryptoRngCore + Metrics, C: Signer> {
+pub struct Actor<E: Spawner + BufferPooler + Clock + Network + CryptoRngCore + Metrics, C: Signer> {
     context: ContextCell<E>,
 
     address: SocketAddr,
@@ -49,7 +49,7 @@ pub struct Actor<E: Spawner + Clock + Network + CryptoRngCore + Metrics, C: Sign
     handshakes_subnet_rate_limited: Counter,
 }
 
-impl<E: Spawner + Clock + Network + CryptoRngCore + Metrics, C: Signer> Actor<E, C> {
+impl<E: Spawner + BufferPooler + Clock + Network + CryptoRngCore + Metrics, C: Signer> Actor<E, C> {
     pub fn new(context: E, cfg: Config<C>) -> Self {
         // Create metrics
         let handshakes_blocked = Counter::default();

p2p/src/authenticated/discovery/mod.rs

@@ -242,7 +242,7 @@ mod tests {
     use commonware_cryptography::{ed25519, Signer as _};
     use commonware_macros::{select, select_loop, test_group, test_traced};
     use commonware_runtime::{
-        count_running_tasks, deterministic, tokio, Clock, Handle, IoBuf, Metrics,
+        count_running_tasks, deterministic, tokio, BufferPooler, Clock, Handle, IoBuf, Metrics,
         Network as RNetwork, Quota, Resolver, Runner, Spawner,
     };
     use commonware_utils::{channel::mpsc, hostname, ordered::Set, TryCollect, NZU32};
@@ -284,7 +284,7 @@ mod tests {
     /// We set a unique `base_port` for each test to avoid "address already in use"
     /// errors when tests are run immediately after each other.
     async fn run_network(
-        context: impl Spawner + Clock + CryptoRngCore + RNetwork + Resolver + Metrics,
+        context: impl Spawner + BufferPooler + Clock + CryptoRngCore + RNetwork + Resolver + Metrics,
         max_message_size: u32,
         base_port: u16,
         n: usize,

p2p/src/authenticated/discovery/network.rs

@@ -13,7 +13,8 @@ use crate::{
 use commonware_cryptography::Signer;
 use commonware_macros::select;
 use commonware_runtime::{
-    spawn_cell, Clock, ContextCell, Handle, Metrics, Network as RNetwork, Quota, Resolver, Spawner,
+    spawn_cell, BufferPooler, Clock, ContextCell, Handle, Metrics, Network as RNetwork, Quota,
+    Resolver, Spawner,
 };
 use commonware_stream::encrypted::Config as StreamConfig;
 use commonware_utils::union;
@@ -27,7 +28,10 @@ const TRACKER_SUFFIX: &[u8] = b"_TRACKER";
 const STREAM_SUFFIX: &[u8] = b"_STREAM";
 
 /// Implementation of an `authenticated` network.
-pub struct Network<E: Spawner + Clock + CryptoRngCore + RNetwork + Resolver + Metrics, C: Signer> {
+pub struct Network<
+    E: Spawner + BufferPooler + Clock + CryptoRngCore + RNetwork + Resolver + Metrics,
+    C: Signer,
+> {
     context: ContextCell<E>,
     cfg: Config<C>,
 
@@ -39,7 +43,11 @@ pub struct Network<E: Spawner + Clock + CryptoRngCore + RNetwork + Resolver + Me
     info_verifier: InfoVerifier<C::PublicKey>,
 }
 
-impl<E: Spawner + Clock + CryptoRngCore + RNetwork + Resolver + Metrics, C: Signer> Network<E, C> {
+impl<
+        E: Spawner + BufferPooler + Clock + CryptoRngCore + RNetwork + Resolver + Metrics,
+        C: Signer,
+    > Network<E, C>
+{
     /// Create a new instance of an `authenticated` network.
     ///
     /// # Parameters