ml_kem: Enable zeroization (#51)

Co-authored-by: Tony Arcieri <[email protected]>
RustCrypto · Aug 17, 2024 · a75d842 · a75d842
1 parent 1cbc77c
commit a75d842
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 0 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/ml-kem/Cargo.toml b/ml-kem/Cargo.toml
@@ -19,12 +19,14 @@ exclude = ["tests/key-gen.rs", "tests/key-gen.json", "tests/encap-decap.rs", "te
 default = ["std"]
 std = ["sha3/std"]
 deterministic = [] # Expose deterministic generation and encapsulation functions
+zeroize = ["dep:zeroize"]
 
 [dependencies]
 kem = "0.3.0-pre.0"
 hybrid-array = { version = "0.2.0-rc.9", features = ["extra-sizes"] }
 rand_core = "0.6.4"
 sha3 = { version = "0.10.8", default-features = false }
+zeroize = { version = "1.8.1", optional = true, default-features = false }
 
 [dev-dependencies]
 criterion = "0.5.1"

diff --git a/ml-kem/src/algebra.rs b/ml-kem/src/algebra.rs
@@ -7,13 +7,23 @@ use crate::encode::Encode;
 use crate::param::{ArraySize, CbdSamplingSize};
 use crate::util::{Truncate, B32};
 
+#[cfg(feature = "zeroize")]
+use zeroize::Zeroize;
+
 pub type Integer = u16;
 
 /// An element of GF(q).  Although `q` is only 16 bits wide, we use a wider uint type to so that we
 /// can defer modular reductions.
 #[derive(Copy, Clone, Debug, Default, PartialEq)]
 pub struct FieldElement(pub Integer);
 
+#[cfg(feature = "zeroize")]
+impl Zeroize for FieldElement {
+    fn zeroize(&mut self) {
+        self.0.zeroize();
+    }
+}
+
 impl FieldElement {
     pub const Q: Integer = 3329;
     pub const Q32: u32 = Self::Q as u32;
@@ -174,6 +184,15 @@ impl<K: ArraySize> PolynomialVector<K> {
 #[derive(Clone, Default, Debug, PartialEq)]
 pub struct NttPolynomial(pub Array<FieldElement, U256>);
 
+#[cfg(feature = "zeroize")]
+impl Zeroize for NttPolynomial {
+    fn zeroize(&mut self) {
+        for fe in self.0.iter_mut() {
+            fe.zeroize()
+        }
+    }
+}
+
 impl Add<&NttPolynomial> for &NttPolynomial {
     type Output = NttPolynomial;
 
@@ -410,6 +429,18 @@ impl<K: ArraySize> NttVector<K> {
     }
 }
 
+#[cfg(feature = "zeroize")]
+impl<K> Zeroize for NttVector<K>
+where
+    K: ArraySize,
+{
+    fn zeroize(&mut self) {
+        for poly in self.0.iter_mut() {
+            poly.zeroize();
+        }
+    }
+}
+
 impl<K: ArraySize> Add<&NttVector<K>> for &NttVector<K> {
     type Output = NttVector<K>;
 

diff --git a/ml-kem/src/kem.rs b/ml-kem/src/kem.rs
@@ -8,6 +8,9 @@ use crate::pke::{DecryptionKey, EncryptionKey};
 use crate::util::B32;
 use crate::{Encoded, EncodedSizeUser};
 
+#[cfg(feature = "zeroize")]
+use zeroize::{Zeroize, ZeroizeOnDrop};
+
 // Re-export traits from the `kem` crate
 pub use ::kem::{Decapsulate, Encapsulate};
 
@@ -26,6 +29,20 @@ where
     z: B32,
 }
 
+#[cfg(feature = "zeroize")]
+impl<P> Drop for DecapsulationKey<P>
+where
+    P: KemParams,
+{
+    fn drop(&mut self) {
+        self.dk_pke.zeroize();
+        self.z.zeroize();
+    }
+}
+
+#[cfg(feature = "zeroize")]
+impl<P> ZeroizeOnDrop for DecapsulationKey<P> where P: KemParams {}
+
 impl<P> EncodedSizeUser for DecapsulationKey<P>
 where
     P: KemParams,

diff --git a/ml-kem/src/pke.rs b/ml-kem/src/pke.rs
@@ -7,6 +7,9 @@ use crate::encode::Encode;
 use crate::param::{EncodedCiphertext, EncodedDecryptionKey, EncodedEncryptionKey, PkeParams};
 use crate::util::B32;
 
+#[cfg(feature = "zeroize")]
+use zeroize::Zeroize;
+
 /// A `DecryptionKey` provides the ability to generate a new key pair, and decrypt an
 /// encrypted value.
 #[derive(Clone, Default, Debug, PartialEq)]
@@ -17,6 +20,16 @@ where
     s_hat: NttVector<P::K>,
 }
 
+#[cfg(feature = "zeroize")]
+impl<P> Zeroize for DecryptionKey<P>
+where
+    P: PkeParams,
+{
+    fn zeroize(&mut self) {
+        self.s_hat.zeroize();
+    }
+}
+
 impl<P> DecryptionKey<P>
 where
     P: PkeParams,