ml-kem: validate encryption/encapsulation keys

Adds a check that decoded `EncryptionKey`s successfully round-trip
encode back to the bytes they were decoded from, in order to detect
malformed keys.

Closes #172

Author: tarcieri

ml-kem/src/pke.rs

@@ -156,15 +156,20 @@ where
     }
 
     /// Parse an encryption key from a byte array `(t_hat || rho)`
-    // TODO(tarcieri): validate decoded keys
-    #[allow(clippy::unnecessary_wraps)]
     pub fn from_bytes(enc: &EncodedEncryptionKey<P>) -> Result<Self, Error> {
         let (t_hat, rho) = P::split_ek(enc);
         let t_hat = P::decode_u12(t_hat);
-        Ok(Self {
+        let ret = Self {
             t_hat,
             rho: rho.clone(),
-        })
+        };
+
+        // Ensure key round-trips successfully
+        if &ret.as_bytes() == enc {
+            Ok(ret)
+        } else {
+            Err(Error)
+        }
     }
 }
 
@@ -221,4 +226,12 @@ mod test {
         codec_test::<MlKem768Params>();
         codec_test::<MlKem1024Params>();
     }
+
+    #[test]
+    fn reject_invalid_encryption_keys() {
+        // Create an invalid key: all bytes set to 0xFF
+        // When decoded as 12-bit coefficients, this produces values of 0xFFF = 4095 > 3329
+        let invalid_key = EncodedEncryptionKey::<MlKem768Params>::from_fn(|_| 0xFF);
+        assert!(EncryptionKey::<MlKem768Params>::from_bytes(&invalid_key).is_err());
+    }
 }