From c3934183c8023d7607ae3cb9f6c0f534b7a447a7 Mon Sep 17 00:00:00 2001
From: Xynnn007 <xynnn@linux.alibaba.com>
Date: Wed, 20 May 2026 15:45:47 +0800
Subject: [PATCH] aes-kw: add optional zeroize feature for wrappers

Close #79

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>
---
 Cargo.lock               |  8 ++++++++
 aes-kw/Cargo.toml        |  2 ++
 aes-kw/src/kw.rs         |  3 +++
 aes-kw/src/kwp.rs        |  3 +++
 aes-kw/src/lib.rs        |  3 +++
 aes-kw/tests/kw_tests.rs | 13 +++++++++++++
 6 files changed, 32 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 03aea07..2822f10 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -12,6 +12,7 @@ dependencies = [
  "cipher",
  "cpubits",
  "cpufeatures",
+ "zeroize",
 ]
 
 [[package]]
@@ -21,6 +22,7 @@ dependencies = [
  "aes",
  "const-oid",
  "hex-literal",
+ "zeroize",
 ]
 
 [[package]]
@@ -115,3 +117,9 @@ name = "typenum"
 version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb"
+
+[[package]]
+name = "zeroize"
+version = "1.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0"
diff --git a/aes-kw/Cargo.toml b/aes-kw/Cargo.toml
index 2016254..74226f9 100644
--- a/aes-kw/Cargo.toml
+++ b/aes-kw/Cargo.toml
@@ -15,6 +15,7 @@ rust-version = "1.85"
 [dependencies]
 aes = "0.9"
 const-oid = { version = "0.10", optional = true }
+zeroize = { version = "1.5.6", optional = true, default-features = false }
 
 [dev-dependencies]
 hex-literal = "1"
@@ -22,6 +23,7 @@ hex-literal = "1"
 [features]
 default = ["oid"]
 oid = ["dep:const-oid"]
+zeroize = ["dep:zeroize", "aes/zeroize"]
 
 [package.metadata.docs.rs]
 all-features = true
diff --git a/aes-kw/src/kw.rs b/aes-kw/src/kw.rs
index 172d596..d527ee2 100644
--- a/aes-kw/src/kw.rs
+++ b/aes-kw/src/kw.rs
@@ -224,3 +224,6 @@ impl<C: BlockCipherDecrypt<BlockSize = U16>> AesKw<C> {
         Ok(buf)
     }
 }
+
+#[cfg(feature = "zeroize")]
+impl<C: zeroize::ZeroizeOnDrop> zeroize::ZeroizeOnDrop for AesKw<C> {}
diff --git a/aes-kw/src/kwp.rs b/aes-kw/src/kwp.rs
index 38bda21..79e7eb8 100644
--- a/aes-kw/src/kwp.rs
+++ b/aes-kw/src/kwp.rs
@@ -274,3 +274,6 @@ impl<C: BlockCipherDecrypt<BlockSize = U16>> AesKwp<C> {
             .map(|res| res.try_into().unwrap())
     }
 }
+
+#[cfg(feature = "zeroize")]
+impl<C: zeroize::ZeroizeOnDrop> zeroize::ZeroizeOnDrop for AesKwp<C> {}
diff --git a/aes-kw/src/lib.rs b/aes-kw/src/lib.rs
index 2545891..40b311e 100644
--- a/aes-kw/src/lib.rs
+++ b/aes-kw/src/lib.rs
@@ -26,6 +26,9 @@ pub use aes;
 pub use aes::cipher;
 pub use aes::cipher::{KeyInit, common::InnerInit};
 
+#[cfg(feature = "zeroize")]
+pub use zeroize;
+
 /// AES-128 key wrapping
 pub type KwAes128 = AesKw<aes::Aes128>;
 /// AES-192 key wrapping
diff --git a/aes-kw/tests/kw_tests.rs b/aes-kw/tests/kw_tests.rs
index d215996..6506e42 100644
--- a/aes-kw/tests/kw_tests.rs
+++ b/aes-kw/tests/kw_tests.rs
@@ -121,3 +121,16 @@ fn error_integrity_check_failed() {
 
     assert_eq!(res, Err(Error::IntegrityCheckFailed));
 }
+
+#[cfg(feature = "zeroize")]
+#[test]
+fn zeroize_on_drop() {
+    use zeroize::ZeroizeOnDrop;
+
+    fn assert_zeroize_on_drop<T: ZeroizeOnDrop>(_: T) {}
+
+    let key256 = hex!("000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F");
+    let key128 = hex!("000102030405060708090A0B0C0D0E0F");
+    assert_zeroize_on_drop(KwAes256::new(&key256.into()));
+    assert_zeroize_on_drop(KwAes128::new(&key128.into()));
+}