argon2: detect allocation failures in hash_password_into

argon2/src/block.rs

@@ -154,3 +154,47 @@ impl Zeroize for Block {
         self.0.zeroize();
     }
 }
+
+/// Custom implementation of `Box<[Block]>` until `Box::try_new_zeroed_slice` is stabilized.
+#[cfg(feature = "alloc")]
+pub(crate) struct Blocks {
+    p: core::ptr::NonNull<Block>,
+    len: usize,
+}
+
+#[cfg(feature = "alloc")]
+impl Blocks {
+    pub fn new(len: usize) -> Option<Self> {
+        use alloc::alloc::{Layout, alloc_zeroed};
+        use core::ptr::NonNull;
+
+        if len == 0 {
+            return None;
+        }
+
+        let layout = Layout::array::<Block>(len).ok()?;
+        // SAFETY: `alloc_zeroed` is used correctly with non-zero layout
+        let p = unsafe { alloc_zeroed(layout) };
+
+        let p = NonNull::new(p.cast())?;
+        Some(Self { p, len })
+    }
+
+    pub fn as_slice(&mut self) -> &mut [Block] {
+        // SAFETY: `self.p` is a valid non-zero pointer that points to memory of the necessary size
+        unsafe { slice::from_raw_parts_mut(self.p.as_ptr(), self.len) }
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl Drop for Blocks {
+    fn drop(&mut self) {
+        use alloc::alloc::{Layout, dealloc};
+        // SAFETY: layout was checked during construction
+        let layout = unsafe { Layout::array::<Block>(self.len).unwrap_unchecked() };
+        // SAFETY: we use `dealloc` correctly with the previously allocated pointer
+        unsafe {
+            dealloc(self.p.as_ptr().cast(), layout);
+        }
+    }
+}

argon2/src/error.rs

@@ -56,8 +56,11 @@ pub enum Error {
     /// Time cost is too small.
     TimeTooSmall,
 
-    /// Invalid version
+    /// Invalid version.
     VersionInvalid,
+
+    /// Out of memory (heap allocation failure).
+    OutOfMemory,
 }
 
 impl fmt::Display for Error {
@@ -79,6 +82,7 @@ impl fmt::Display for Error {
             Error::ThreadsTooMany => "too many threads",
             Error::TimeTooSmall => "time cost is too small",
             Error::VersionInvalid => "invalid version",
+            Error::OutOfMemory => "out of memory",
         })
     }
 }
@@ -116,6 +120,8 @@ impl From<Error> for password_hash::Error {
             Error::ThreadsTooMany => InvalidValue::TooLong.param_error(),
             Error::TimeTooSmall => InvalidValue::TooShort.param_error(),
             Error::VersionInvalid => password_hash::Error::Version,
+            // TODO: fix after `password_hash::Error::OutOfMemory` will be added
+            Error::OutOfMemory => InvalidValue::TooLong.param_error(),
         }
     }
 }

argon2/src/lib.rs

@@ -144,7 +144,6 @@
 compile_error!("this crate builds on 32-bit and 64-bit platforms only");
 
 #[cfg(feature = "alloc")]
-#[macro_use]
 extern crate alloc;
 
 #[cfg(feature = "std")]
@@ -286,8 +285,9 @@ impl<'key> Argon2<'key> {
     #[cfg(feature = "alloc")]
     #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
     pub fn hash_password_into(&self, pwd: &[u8], salt: &[u8], out: &mut [u8]) -> Result<()> {
-        let mut blocks = vec![Block::default(); self.params.block_count()];
-        self.hash_password_into_with_memory(pwd, salt, out, &mut blocks)
+        let blocks_len = self.params.block_count();
+        let mut blocks = block::Blocks::new(blocks_len).ok_or(Error::OutOfMemory)?;
+        self.hash_password_into_with_memory(pwd, salt, out, blocks.as_slice())
     }
 
     /// Hash a password and associated parameters into the provided output buffer.