key packing tests

supinie · Apr 17, 2024 · a667a83 · a667a83
1 parent e33d1a0
commit a667a83
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 15 deletions.
diff --git a/src/indcpa.rs b/src/indcpa.rs
@@ -25,11 +25,13 @@ impl PrivateKey {
         self.secret.sec_level()
     }
 
-    fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
+    // buf should be of length k * POLYBYTES
+    pub(crate) fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
         self.secret.pack(buf)
     }
 
-    fn unpack(buf: &[u8]) -> Result<Self, PackingError> {
+    // buf should be of length k * POLYBYTES
+    pub(crate) fn unpack(buf: &[u8]) -> Result<Self, PackingError> {
         let secret = PolyVec::unpack(buf)?.normalise();
         Ok(Self { secret })
     }
@@ -47,7 +49,8 @@ impl PublicKey {
         }
     }
 
-    fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
+    // buf should be of length k * POLYBYTES + SYMBYTES
+    pub fn pack(&self, buf: &mut [u8]) -> Result<(), PackingError> {
         let k: usize = self.sec_level()?.k().into();
 
         let break_point: usize = POLYBYTES * k;
@@ -60,7 +63,8 @@ impl PublicKey {
         }
     }
 
-    fn unpack(buf: &[u8]) -> Result<Self, PackingError> {
+    // buf should be of length k * POLYBYTES + SYMBYTES
+    pub fn unpack(buf: &[u8]) -> Result<Self, PackingError> {
         let k = K::try_from((buf.len() - SYMBYTES) / POLYBYTES)?;
         let k_value: usize = k.into();
         let break_point: usize = POLYBYTES * k_value;

diff --git a/src/tests/indcpa.rs b/src/tests/indcpa.rs
@@ -4,9 +4,11 @@ mod indcpa_tests {
     use crate::{
         indcpa::*,
         params::*,
+        tests::params::params_tests::sec_level_strategy,
     };
     use rand::rngs::StdRng;
     use rand::{Rng, SeedableRng};
+    use proptest::prelude::*;
 
     pub(in crate::tests) fn generate_random_seed() -> [u8; 32] {
         let mut rng = StdRng::from_entropy();
@@ -15,20 +17,43 @@ mod indcpa_tests {
         seed
     }
 
-    #[test]
-    fn key_gen_enc_dec() {
-        let key_seed = generate_random_seed();
-        let cipher_seed = generate_random_seed();
-        // let plaintext = generate_random_seed();
-        let plaintext: [u8; 32] = core::array::from_fn(|i| (i + 1) as u8);
+    proptest! {
+        #[test]
+        fn key_gen_enc_dec(
+                key_seed in prop::array::uniform32(u8::MIN..u8::MAX),
+                cipher_seed in prop::array::uniform32(u8::MIN..u8::MAX)
+            ) {
+            // let plaintext = generate_random_seed();
+            let plaintext: [u8; 32] = core::array::from_fn(|i| (i + 1) as u8);
 
 
-        let (priv_key, pub_key) = generate_key_pair(&key_seed, SecurityLevel::new(K::Three)).unwrap();
-
-        let ciphertext = encrypt(&pub_key, &plaintext, &cipher_seed).unwrap();
+            let (priv_key, pub_key) = generate_key_pair(&key_seed, SecurityLevel::new(K::Three)).unwrap();
+            
+            let ciphertext = encrypt(&pub_key, &plaintext, &cipher_seed).unwrap();
 
-        let message = decrypt(&priv_key, &ciphertext).unwrap();
+            let message = decrypt(&priv_key, &ciphertext).unwrap();
 
-        // assert_eq!(message, plaintext);
+            // assert_eq!(message, plaintext);
+        }
+
+        #[test]
+        fn key_pack_unpack(
+            key_seed in prop::array::uniform32(u8::MIN..u8::MAX),
+            sec_level in sec_level_strategy()
+        ) {
+            let (priv_key, pub_key) = generate_key_pair(&key_seed, sec_level).unwrap();
+
+            let mut buf = [0u8; 2 * (4 * POLYBYTES) + SYMBYTES];
+
+            let k:usize = sec_level.k().into();
+            let _ = priv_key.pack(&mut buf[..k * POLYBYTES]);
+            let _ = pub_key.pack(&mut buf[k * POLYBYTES..2 * (k * POLYBYTES) + SYMBYTES]);
+
+            let unpacked_priv = PrivateKey::unpack(&buf[..k * POLYBYTES]).unwrap();
+            let unpacked_pub = PublicKey::unpack(&buf[k * POLYBYTES..2 * (k * POLYBYTES) + SYMBYTES]).unwrap();
+
+            assert_eq!(unpacked_pub, pub_key);
+            assert_eq!(unpacked_priv, priv_key);
+        }
     }
 }