Tweak buffer methods

aead/src/lib.rs

@@ -45,7 +45,7 @@ use core::fmt;
 use crypto_common::array::{typenum::Unsigned, Array, ArraySize};
 
 #[cfg(feature = "alloc")]
-use alloc::{vec, vec::Vec};
+use alloc::vec::Vec;
 #[cfg(feature = "getrandom")]
 use crypto_common::getrandom;
 #[cfg(feature = "rand_core")]
@@ -319,7 +319,9 @@ pub trait Aead {
         nonce: &Nonce<Self>,
         pt_payload: impl Into<Payload<'msg, 'aad>>,
     ) -> Result<Vec<u8>> {
-        self.encrypt_to_buffer(nonce, pt_payload)
+        let mut buf = Vec::new();
+        self.encrypt_to_buffer(nonce, pt_payload, &mut buf)?;
+        Ok(buf)
     }
 
     /// Decrypt the given ciphertext slice, and return the resulting plaintext
@@ -346,46 +348,50 @@ pub trait Aead {
         nonce: &Nonce<Self>,
         ct_payload: impl Into<Payload<'msg, 'aad>>,
     ) -> Result<Vec<u8>> {
-        self.decrypt_to_buffer(nonce, ct_payload)
+        let mut buf = Vec::new();
+        self.decrypt_to_buffer(nonce, ct_payload, &mut buf)?;
+        Ok(buf)
     }
 
     #[inline]
     fn encrypt_to_buffer<'msg, 'aad, B: Buffer>(
         &self,
         nonce: &Nonce<Self>,
         pt_payload: impl Into<Payload<'msg, 'aad>>,
-    ) -> Result<B> {
+        buffer: &mut B,
+    ) -> Result<()> {
         let Payload { msg: pt, aad } = pt_payload.into();
         let tag_len = Self::TagSize::USIZE;
-        let mut res = B::zeroed(pt.len() + tag_len)?;
+        buffer.resize(pt.len() + tag_len)?;
         let (ct_dst, tag_dst) = if Self::IS_POSTFIX {
-            res.as_mut().split_at_mut(pt.len())
+            buffer.as_mut().split_at_mut(pt.len())
         } else {
-            res.as_mut().split_at_mut(tag_len)
+            buffer.as_mut().split_at_mut(tag_len)
         };
         let tag = self.detached_encrypt_to_buf(nonce, aad, pt, ct_dst)?;
         tag_dst.copy_from_slice(&tag);
-        Ok(res)
+        Ok(())
     }
 
     #[inline]
     fn decrypt_to_buffer<'msg, 'aad, B: Buffer>(
         &self,
         nonce: &Nonce<Self>,
         ct_payload: impl Into<Payload<'msg, 'aad>>,
-    ) -> Result<B> {
+        buffer: &mut B,
+    ) -> Result<()> {
         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 = B::zeroed(pt_len)?;
+        buffer.resize(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_to_buf(nonce, aad, ct, pt_dst.as_mut(), tag)?;
-        Ok(pt_dst)
+        self.detached_decrypt_to_buf(nonce, aad, ct, buffer.as_mut(), tag)?;
+        Ok(())
     }
 
     /// Generate a random nonce for this AEAD algorithm.
@@ -503,8 +509,10 @@ impl<'msg> From<&'msg [u8]> for Payload<'msg, '_> {
 /// 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>;
+    /// Resizes buffer to the requested length.
+    ///
+    /// If buffer is smaller than `len`, fills it with zeros. Otherwise, truncates it to `len`.
+    fn resize(&mut self, len: usize) -> Result<()>;
 
     /// Extend this buffer from the given slice
     fn extend_from_slice(&mut self, other: &[u8]) -> Result<()>;
@@ -515,8 +523,9 @@ pub trait Buffer: AsMut<[u8]> + Sized {
 
 #[cfg(feature = "alloc")]
 impl Buffer for Vec<u8> {
-    fn zeroed(len: usize) -> Result<Self> {
-        Ok(vec![0; len])
+    fn resize(&mut self, len: usize) -> Result<()> {
+        Vec::resize(self, len, 0);
+        Ok(())
     }
 
     fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
@@ -531,8 +540,9 @@ impl Buffer for Vec<u8> {
 
 #[cfg(feature = "bytes")]
 impl Buffer for BytesMut {
-    fn zeroed(len: usize) -> Result<Self> {
-        Ok(BytesMut::zeroed(len))
+    fn resize(&mut self, len: usize) -> Result<()> {
+        BytesMut::resize(self, len, 0);
+        Ok(())
     }
 
     fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
@@ -547,13 +557,14 @@ impl Buffer for BytesMut {
 
 #[cfg(feature = "arrayvec")]
 impl<const N: usize> Buffer for ArrayVec<u8, N> {
-    fn zeroed(len: usize) -> Result<Self> {
-        if len > N {
-            return Err(Error);
+    fn resize(&mut self, len: usize) -> Result<()> {
+        if let Some(ext_len) = len.checked_sub(self.len()) {
+            let buf = &[0u8; N][..ext_len];
+            self.try_extend_from_slice(buf).map_err(|_| Error)
+        } else {
+            self.truncate(len);
+            Ok(())
         }
-        let mut buf = ArrayVec::from([0u8; N]);
-        buf.truncate(len);
-        Ok(buf)
     }
 
     fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {
@@ -567,10 +578,8 @@ impl<const N: usize> Buffer for ArrayVec<u8, N> {
 
 #[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 resize(&mut self, len: usize) -> Result<()> {
+        heapless::Vec::resize(self, len, 0).map_err(|_| Error)
     }
 
     fn extend_from_slice(&mut self, other: &[u8]) -> Result<()> {