Return Buffer

aead/src/lib.rs

@@ -16,19 +16,23 @@
 #[cfg(feature = "alloc")]
 extern crate alloc;
 
+#[cfg(feature = "bytes")]
+use bytes::BytesMut;
+
+#[cfg(feature = "arrayvec")]
+use arrayvec::ArrayVec;
+
 #[cfg(feature = "dev")]
 pub mod dev;
 
 // pub mod stream;
 
 pub use crypto_common::{
+    self,
     array::{self, typenum::consts},
     Key, KeyInit, KeySizeUser,
 };
 
-#[cfg(feature = "getrandom")]
-pub use crypto_common::rand_core::OsRng;
-
 #[cfg(feature = "rand_core")]
 pub use crypto_common::rand_core;
 use inout::{InOutBuf, InOutBufReserved};
@@ -100,7 +104,7 @@ pub trait Aead {
 
     /// Encrypt the data in-place, returning the authentication tag.
     #[inline]
-    fn detached_encrypt(
+    fn detached_encrypt_inplace(
         &self,
         nonce: &Nonce<Self>,
         associated_data: &[u8],
@@ -112,7 +116,7 @@ pub trait Aead {
     /// Encrypt the data in-place, returning an error in the event the provided
     /// authentication tag does not match the given ciphertext.
     #[inline]
-    fn detached_decrypt(
+    fn detached_decrypt_inplace(
         &self,
         nonce: &Nonce<Self>,
         associated_data: &[u8],
@@ -124,7 +128,7 @@ pub trait Aead {
 
     /// Encrypt the data buffer-to-buffer, returning the authentication tag.
     #[inline]
-    fn detached_encrypt_b2b(
+    fn detached_encrypt_buf2buf(
         &self,
         nonce: &Nonce<Self>,
         associated_data: &[u8],
@@ -257,33 +261,25 @@ pub trait Aead {
         Ok(pt_dst)
     }
 
-    #[cfg(feature = "alloc")]
     #[inline]
-    fn postfix_encrypt_vec(
+    fn postfix_encrypt_buffer(
         &self,
         nonce: &Nonce<Self>,
         associated_data: &[u8],
-        plaintext: &[u8],
-    ) -> Result<Vec<u8>> {
-        let tag_len = Self::TagSize::USIZE;
-        let pt_len = plaintext.len();
-        let mut buf = vec![0u8; pt_len + tag_len];
-        let (ct_dst, tag_dst) = buf.split_at_mut(pt_len);
-        let inout_buf = InOutBuf::new(plaintext, ct_dst).expect("ct_dst has correct length");
-        let tag = self.detached_encrypt_inout(nonce, associated_data, inout_buf)?;
-        tag_dst.copy_from_slice(&tag);
-        Ok(buf)
+        buffer: &mut impl Buffer,
+    ) -> Result<()> {
+        let tag = self.detached_encrypt_inplace(nonce, associated_data, buffer.as_mut())?;
+        buffer.extend_from_slice(&tag)
     }
 
-    #[cfg(feature = "alloc")]
     #[inline]
-    fn postfix_decrypt_vec(
+    fn postfix_decrypt_buffer(
         &self,
         nonce: &Nonce<Self>,
         associated_data: &[u8],
-        buffer: &mut Vec<u8>,
+        buffer: &mut impl Buffer,
     ) -> Result<()> {
-        let pt = self.postfix_decrypt(nonce, associated_data, buffer)?;
+        let pt = self.postfix_decrypt(nonce, associated_data, buffer.as_mut())?;
         let pt_len = pt.len();
         buffer.truncate(pt_len);
         Ok(())
@@ -312,22 +308,14 @@ pub trait Aead {
     /// AES-GCM-SIV, ChaCha20Poly1305). [`Aead`] implementations which do not
     /// use a postfix tag will need to override this to correctly assemble the
     /// ciphertext message.
-    fn encrypt<'msg, 'aad>(
+    #[cfg(feature = "alloc")]
+    #[inline]
+    fn encrypt_to_vec<'msg, 'aad, B: Buffer>(
         &self,
         nonce: &Nonce<Self>,
         pt_payload: impl Into<Payload<'msg, 'aad>>,
     ) -> Result<Vec<u8>> {
-        let Payload { msg: pt, aad } = pt_payload.into();
-        let tag_len = Self::TagSize::USIZE;
-        let mut res = vec![0u8; pt.len() + tag_len];
-        let (ct_dst, tag_dst) = if Self::IS_POSTFIX {
-            res.split_at_mut(pt.len())
-        } else {
-            res.split_at_mut(tag_len)
-        };
-        let tag = self.detached_encrypt_b2b(nonce, aad, pt, ct_dst)?;
-        tag_dst.copy_from_slice(&tag);
-        Ok(res)
+        self.encrypt_to_buffer(nonce, pt_payload)
     }
 
     /// Decrypt the given ciphertext slice, and return the resulting plaintext
@@ -347,22 +335,51 @@ pub trait Aead {
     /// AES-GCM-SIV, ChaCha20Poly1305). [`Aead`] implementations which do not
     /// use a postfix tag will need to override this to correctly parse the
     /// ciphertext message.
-    fn decrypt<'msg, 'aad>(
+    #[inline]
+    fn decrypt_to_vec<'msg, 'aad>(
         &self,
         nonce: &Nonce<Self>,
         ct_payload: impl Into<Payload<'msg, 'aad>>,
     ) -> Result<Vec<u8>> {
+        self.decrypt_to_buffer(nonce, ct_payload)
+    }
+
+    #[inline]
+    fn encrypt_to_buffer<'msg, 'aad, B: Buffer>(
+        &self,
+        nonce: &Nonce<Self>,
+        pt_payload: impl Into<Payload<'msg, 'aad>>,
+    ) -> Result<B> {
+        let Payload { msg: pt, aad } = pt_payload.into();
+        let tag_len = Self::TagSize::USIZE;
+        let mut res = B::zeroed(pt.len() + tag_len)?;
+        let (ct_dst, tag_dst) = if Self::IS_POSTFIX {
+            res.as_mut().split_at_mut(pt.len())
+        } else {
+            res.as_mut().split_at_mut(tag_len)
+        };
+        let tag = self.detached_encrypt_buf2buf(nonce, aad, pt, ct_dst)?;
+        tag_dst.copy_from_slice(&tag);
+        Ok(res)
+    }
+
+    #[inline]
+    fn decrypt_to_buffer<'msg, 'aad, B: Buffer>(
+        &self,
+        nonce: &Nonce<Self>,
+        ct_payload: impl Into<Payload<'msg, 'aad>>,
+    ) -> Result<B> {
         let Payload { msg: ct_tag, aad } = ct_payload.into();
         let tag_len = Self::TagSize::USIZE;
         let pt_len = ct_tag.len().checked_sub(tag_len).ok_or(Error)?;
-        let mut pt_dst = vec![0u8; pt_len];
+        let mut pt_dst = B::zeroed(pt_len)?;
         let (ct, tag) = if Self::IS_POSTFIX {
             ct_tag.split_at(pt_len)
         } else {
             ct_tag.split_at(tag_len)
         };
         let tag = tag.try_into().expect("tag has correct length");
-        self.detached_decrypt_b2b(nonce, aad, ct, &mut pt_dst, tag)?;
+        self.detached_decrypt_b2b(nonce, aad, ct, pt_dst.as_mut(), tag)?;
         Ok(pt_dst)
     }
 
@@ -410,7 +427,7 @@ pub trait Aead {
 
     /// Generate a random nonce for this AEAD algorithm using the specified [`CryptoRngCore`].
     ///
-    /// See [`AeadCore::generate_nonce`] documentation for requirements for random nonces.
+    /// See [`Aead::generate_nonce`] documentation for requirements for random nonces.
     #[cfg(feature = "rand_core")]
     fn generate_nonce_with_rng(
         rng: &mut impl CryptoRngCore,
@@ -476,3 +493,87 @@ impl<'msg> From<&'msg [u8]> for Payload<'msg, '_> {
         Self { msg, aad: b"" }
     }
 }
+
+/// In-place encryption/decryption byte buffers.
+///
+/// This trait defines the set of methods needed to support in-place operations
+/// on a `Vec`-like data type.
+pub trait Buffer: AsMut<[u8]> + Sized {
+    /// Creates new buffer with the requested length.
+    fn zeroed(len: usize) -> Result<Self>;
+
+    /// Extend this buffer from the given slice
+    fn extend_from_slice(&mut self, other: &[u8]) -> Result<()>;
+
+    /// Truncate this buffer to the given size
+    fn truncate(&mut self, len: usize);
+}
+
+#[cfg(feature = "alloc")]
+impl Buffer for Vec<u8> {
+    fn zeroed(len: usize) -> Result<Self> {
+        Ok(vec![0; len])
+    }
+
+    fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
+        Vec::extend_from_slice(self, other);
+        Ok(())
+    }
+
+    fn truncate(&mut self, len: usize) {
+        Vec::truncate(self, len);
+    }
+}
+
+#[cfg(feature = "bytes")]
+impl Buffer for BytesMut {
+    fn zeroed(len: usize) -> Result<Self> {
+        Ok(BytesMut::zeroed(len))
+    }
+
+    fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
+        BytesMut::extend_from_slice(self, other);
+        Ok(())
+    }
+
+    fn truncate(&mut self, len: usize) {
+        BytesMut::truncate(self, len);
+    }
+}
+
+#[cfg(feature = "arrayvec")]
+impl<const N: usize> Buffer for ArrayVec<u8, N> {
+    fn zeroed(len: usize) -> Result<Self> {
+        if len > N {
+            return Err(Error);
+        }
+        let mut buf = ArrayVec::from([0u8; N]);
+        buf.truncate(len);
+        Ok(buf)
+    }
+
+    fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
+        ArrayVec::try_extend_from_slice(self, other).map_err(|_| Error)
+    }
+
+    fn truncate(&mut self, len: usize) {
+        ArrayVec::truncate(self, len);
+    }
+}
+
+#[cfg(feature = "heapless")]
+impl<const N: usize> Buffer for heapless::Vec<u8, N> {
+    fn zeroed(len: usize) -> Result<Self> {
+        let mut buf = heapless::Vec::<u8, N>::new();
+        buf.resize_default(len).map_err(|()| Error)?;
+        Ok(buf)
+    }
+
+    fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
+        heapless::Vec::extend_from_slice(self, other).map_err(|_| Error)
+    }
+
+    fn truncate(&mut self, len: usize) {
+        heapless::Vec::truncate(self, len);
+    }
+}