From 176d993cf104518651400816bbf92eb253c3bbab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=90=D1=80=D1=82=D1=91=D0=BC=20=D0=9F=D0=B0=D0=B2=D0=BB?=
 =?UTF-8?q?=D0=BE=D0=B2=20=5BArtyom=20Pavlov=5D?= <newpavlov@gmail.com>
Date: Thu, 11 Jan 2024 19:14:54 +0300
Subject: [PATCH] Add zeroize support for blake2

---
 blake2/Cargo.toml         |  1 +
 blake2/src/lib.rs         |  3 +++
 blake2/src/macros.rs      | 35 +++++++++++++++++++++++++++++++++++
 blake2/src/simd/simdty.rs | 13 +++++++++++++
 4 files changed, 52 insertions(+)

diff --git a/blake2/Cargo.toml b/blake2/Cargo.toml
index 485936a32..95d265611 100644
--- a/blake2/Cargo.toml
+++ b/blake2/Cargo.toml
@@ -22,6 +22,7 @@ hex-literal = "0.4"
 [features]
 default = ["std"]
 std = ["digest/std"]
+zeroize = ["digest/zeroize"]
 reset = [] # Enable reset functionality
 #simd = []
 #simd_opt = ["simd"]
diff --git a/blake2/src/lib.rs b/blake2/src/lib.rs
index d435f15b7..266c39ec5 100644
--- a/blake2/src/lib.rs
+++ b/blake2/src/lib.rs
@@ -31,6 +31,9 @@ use digest::{
 #[cfg(feature = "reset")]
 use digest::{FixedOutputReset, Reset};
 
+#[cfg(feature = "zeroize")]
+use digest::zeroize::{Zeroize, ZeroizeOnDrop};
+
 mod as_bytes;
 mod consts;
 
diff --git a/blake2/src/macros.rs b/blake2/src/macros.rs
index bf92d5039..3c9459ff6 100644
--- a/blake2/src/macros.rs
+++ b/blake2/src/macros.rs
@@ -246,6 +246,18 @@ macro_rules! blake2_impl {
                 f.write_str(concat!(stringify!($name), " { ... }"))
             }
         }
+
+        impl Drop for $name {
+            fn drop(&mut self) {
+                #[cfg(feature = "zeroize")]
+                {
+                    self.h.zeroize();
+                    self.t.zeroize();
+                }
+            }
+        }
+        #[cfg(feature = "zeroize")]
+        impl ZeroizeOnDrop for $name {}
     };
 }
 
@@ -426,5 +438,28 @@ macro_rules! blake2_mac_impl {
                 write!(f, "{}{} {{ ... }}", stringify!($name), OutSize::USIZE)
             }
         }
+
+        impl<OutSize> Drop for $name<OutSize>
+        where
+            OutSize: ArraySize + IsLessOrEqual<$max_size>,
+            LeEq<OutSize, $max_size>: NonZero,
+        {
+            fn drop(&mut self) {
+                #[cfg(feature = "zeroize")]
+                {
+                    // `self.core` zeroized by its `Drop` impl
+                    self.buffer.zeroize();
+                    #[cfg(feature = "reset")]
+                    self.key_block.zeroize();
+                }
+            }
+        }
+        #[cfg(feature = "zeroize")]
+        impl<OutSize> ZeroizeOnDrop for $name<OutSize>
+        where
+            OutSize: ArraySize + IsLessOrEqual<$max_size>,
+            LeEq<OutSize, $max_size>: NonZero,
+        {
+        }
     };
 }
diff --git a/blake2/src/simd/simdty.rs b/blake2/src/simd/simdty.rs
index 008b8b48c..3b271bd54 100644
--- a/blake2/src/simd/simdty.rs
+++ b/blake2/src/simd/simdty.rs
@@ -9,6 +9,9 @@
 
 use crate::as_bytes::Safe;
 
+#[cfg(feature = "zeroize")]
+use digest::zeroize::Zeroize;
+
 #[cfg(feature = "simd")]
 macro_rules! decl_simd {
     ($($decl:item)*) => {
@@ -50,6 +53,16 @@ decl_simd! {
                          pub T, pub T, pub T, pub T);
 }
 
+#[cfg(feature = "zeroize")]
+impl<T: Zeroize> Zeroize for Simd4<T> {
+    fn zeroize(&mut self) {
+        self.0.zeroize();
+        self.1.zeroize();
+        self.2.zeroize();
+        self.3.zeroize();
+    }
+}
+
 pub type u64x2 = Simd2<u64>;
 
 pub type u32x4 = Simd4<u32>;