Unit tests

* Unit tests for CipherState
* Unit tests for SymmetricState
* Nonce increment bug fix (found with tests)
* Integration test restructuring

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "clatter"
-version = "0.1.3-alpha"
+version = "0.1.4-alpha"
 edition = "2021"
 license = "MIT"
 description = "no_std compatible implementation of Noise protocol framework with Post-Quantum extensions"

src/bytearray.rs

@@ -1,8 +1,10 @@
+use core::fmt::Debug;
+
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
 /// Simple trait used throughout the codebase to provide
 /// portable array operations
-pub trait ByteArray: Sized + Zeroize {
+pub trait ByteArray: Sized + Zeroize + PartialEq + Debug {
     fn new_zero() -> Self;
     fn new_with(_: u8) -> Self;
     fn from_slice(_: &[u8]) -> Self;
@@ -15,7 +17,7 @@ pub trait ByteArray: Sized + Zeroize {
 }
 
 /// Encapsulation for all [`ByteArray`] types that is automatically zeroized on drop.
-#[derive(ZeroizeOnDrop, Zeroize, Clone)]
+#[derive(ZeroizeOnDrop, Zeroize, Clone, PartialEq, Debug)]
 pub struct SensitiveByteArray<A: ByteArray>(A);
 
 impl<A: ByteArray> SensitiveByteArray<A> {

src/cipherstate.rs

@@ -40,8 +40,11 @@ impl<C: Cipher> CipherState<C> {
     fn nonce_inc_check(&mut self) {
         // "If incrementing n results in 2^(64)-1, then any further EncryptWithAd()
         // or DecryptWithAd() calls will signal an error to the caller"
-        if self.n.checked_add(1).is_none() {
-            self.overflowed = true;
+        match self.n.checked_add(1) {
+            None => self.overflowed = true,
+            Some(n) => {
+                self.n = n;
+            }
         }
     }
 
@@ -137,3 +140,130 @@ impl<C: Cipher> CipherState<C> {
         self.k = C::rekey(&self.k)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use core::u64;
+
+    use super::CipherState;
+    use crate::crypto::cipher::{AesGcm, ChaChaPoly};
+    use crate::traits::Cipher;
+
+    const K: &[u8] = b"Back home.... where I belong....";
+
+    fn cipher_suite<C: Cipher>() {
+        let mut c1 = CipherState::<C>::new(K, 0);
+        let mut c2 = CipherState::<C>::new(K, 0);
+
+        let mut c1_buf = [0u8; 4069];
+        let mut c2_buf = [0u8; 4069];
+
+        let msg = b"Decadent scenes from my memory";
+        let cipher_len = msg.len() + C::tag_len();
+
+        // Normal encrypt-decrypt
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        c2.decrypt_with_ad(&[], &c1_buf[..cipher_len], &mut c2_buf[..msg.len()])
+            .unwrap();
+        assert_eq!(*msg, c2_buf[..msg.len()]);
+        assert!(c1_buf[..msg.len()] != c2_buf[..msg.len()]);
+
+        // With AD
+        c1.encrypt_with_ad(b"Close your eyes", msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        c2.decrypt_with_ad(
+            b"Close your eyes",
+            &c1_buf[..cipher_len],
+            &mut c2_buf[..msg.len()],
+        )
+        .unwrap();
+        assert_eq!(*msg, c2_buf[..msg.len()]);
+
+        // Wrong AD
+        c1.encrypt_with_ad(b"Close your eyes", msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        assert!(c2
+            .decrypt_with_ad(
+                b"Close your eyes and relax",
+                &c1_buf[..cipher_len],
+                &mut c2_buf[..msg.len()]
+            )
+            .is_err());
+
+        // Nonce is now desynchronized
+        assert!(c1.get_nonce() != c2.get_nonce());
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        assert!(c2
+            .decrypt_with_ad(
+                b"Close your eyes and relax",
+                &c1_buf[..cipher_len],
+                &mut c2_buf[..msg.len()]
+            )
+            .is_err());
+
+        // Restore nonce
+        c2.set_nonce(c1.get_nonce());
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        c2.decrypt_with_ad(&[], &c1_buf[..cipher_len], &mut c2_buf[..msg.len()])
+            .unwrap();
+        assert_eq!(*msg, c2_buf[..msg.len()]);
+
+        // Rekey responder
+        c2.rekey();
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        assert!(c2
+            .decrypt_with_ad(
+                b"Close your eyes and relax",
+                &c1_buf[..cipher_len],
+                &mut c2_buf[..msg.len()]
+            )
+            .is_err());
+
+        // Rekey sender (and restore nonce...)
+        c1.rekey();
+        c2.set_nonce(c1.get_nonce());
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        c2.decrypt_with_ad(&[], &c1_buf[..cipher_len], &mut c2_buf[..msg.len()])
+            .unwrap();
+        assert_eq!(*msg, c2_buf[..msg.len()]);
+
+        // Rekey a lot
+        for _ in 0..10000 {
+            c1.rekey();
+            c2.rekey();
+        }
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        c2.decrypt_with_ad(&[], &c1_buf[..cipher_len], &mut c2_buf[..msg.len()])
+            .unwrap();
+        assert_eq!(*msg, c2_buf[..msg.len()]);
+
+        // Nonce overflow
+        c1.set_nonce(u64::MAX);
+        // This should be ok
+        c1.encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .unwrap();
+        // This and all following calls should result in an error
+        assert!(c1
+            .encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .is_err());
+        assert!(c1
+            .encrypt_with_ad(&[], msg, &mut c1_buf[..cipher_len])
+            .is_err());
+    }
+
+    #[test]
+    fn cipher_suite_chacha() {
+        cipher_suite::<ChaChaPoly>();
+    }
+
+    #[test]
+    fn cipher_suite_aes_gcm() {
+        cipher_suite::<AesGcm>();
+    }
+}

src/lib.rs

@@ -131,8 +131,6 @@ pub mod handshakepattern;
 mod handshakestate;
 /// Symmetricstate implementation
 mod symmetricstate;
-#[cfg(test)]
-mod test;
 /// Common traits for supporting crypto algorithms
 pub mod traits;
 /// Transportstate implementation

src/symmetricstate.rs

@@ -111,13 +111,25 @@ where
     }
 
     /// Return [`CipherStates`] for encrypting transport messages
+    ///
+    /// # Panics:
+    /// * If no key material has been stablished
     pub(crate) fn split(&self) -> CipherStates<C> {
-        let (temp_k1, temp_k2) = H::hkdf(self.ck.as_slice(), &[]);
+        // This means that we are still in the initial state
+        if self.ck == self.h {
+            panic!("No key material")
+        }
 
-        CipherStates {
+        let (mut temp_k1, mut temp_k2) = H::hkdf(self.ck.as_slice(), &[]);
+
+        let ct = CipherStates {
             initiator_to_responder: CipherState::new(&temp_k1.as_slice()[..C::key_len()], 0),
             responder_to_initiator: CipherState::new(&temp_k2.as_slice()[..C::key_len()], 0),
-        }
+        };
+
+        temp_k1.zeroize();
+        temp_k2.zeroize();
+        ct
     }
 
     /// Get handshake hash
@@ -130,3 +142,114 @@ where
         self.cipherstate.is_some()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::SymmetricState;
+    use crate::crypto::cipher::{AesGcm, ChaChaPoly};
+    use crate::crypto::hash::{Blake2b, Blake2s, Sha256, Sha512};
+    use crate::traits::{Cipher, Hash};
+
+    impl<C: Cipher, H: Hash> PartialEq for SymmetricState<C, H> {
+        fn eq(&self, other: &Self) -> bool {
+            self.h == other.h && self.ck == other.ck
+        }
+    }
+
+    fn symmetric_suite<C: Cipher, H: Hash>() {
+        let mut s1 = SymmetricState::<C, H>::new("complex delirium");
+        let mut s2 = SymmetricState::<C, H>::new("complex delirium");
+
+        // Don't have keys yet
+        assert!(!s1.has_key());
+        assert!(!s2.has_key());
+
+        // Identical at start
+        assert!(s1 == s2);
+
+        // Mix hash
+        s1.mix_hash(b"all wound up");
+        s2.mix_hash(b"all wound up");
+        assert!(s1 == s2);
+        assert!(!s1.has_key() && !s2.has_key());
+
+        // Mix key
+        s1.mix_key(b"sleep disturbed");
+        s2.mix_key(b"sleep disturbed");
+        assert!(s1 == s2);
+        assert!(s1.has_key() && s2.has_key());
+
+        // Mix key and hash
+        s1.mix_key_and_hash(b"sleep disturbed");
+        s2.mix_key_and_hash(b"sleep disturbed");
+        assert!(s1 == s2);
+
+        // Mix key and hash empty
+        s1.mix_key_and_hash(&[]);
+        s2.mix_key_and_hash(&[]);
+        assert!(s1 == s2);
+
+        // Encrypt and hash
+        let mut buf1 = [0; 4096];
+        let mut buf2 = [0; 4096];
+        let msg = b"caught off guard";
+        s1.encrypt_and_hash(msg, &mut buf1[..msg.len() + C::tag_len()])
+            .unwrap();
+        assert!(s1 != s2);
+        assert!(msg != &buf1[..msg.len()]);
+
+        // Decrypt and hash
+        s2.decrypt_and_hash(&buf1[..msg.len() + C::tag_len()], &mut buf2[..msg.len()])
+            .unwrap();
+        assert_eq!(*msg, buf2[..msg.len()]);
+        assert!(s1 == s2);
+
+        // Split
+        let s1_c = s1.split();
+        let s2_c = s2.split();
+
+        assert!(s1_c.initiator_to_responder.take() == s2_c.initiator_to_responder.take());
+        assert!(s1_c.responder_to_initiator.take() == s2_c.responder_to_initiator.take());
+
+        // Mix in different material
+        s1.mix_key_and_hash(b"run");
+        s2.mix_key_and_hash(b"try to hide");
+        assert!(s1 != s2);
+
+        // Encrypt and hash with wrong hashes
+        s1.encrypt_and_hash(msg, &mut buf1[..msg.len() + C::tag_len()])
+            .unwrap();
+        assert!(s1 != s2);
+        assert!(msg != &buf1[..msg.len()]);
+
+        // Decrypt should fail
+        assert!(s2
+            .decrypt_and_hash(&buf1[..msg.len() + C::tag_len()], &mut buf2[..msg.len()])
+            .is_err());
+
+        // Verify that we panic if no key material is available
+        cant_split_without_key::<C, H>();
+    }
+
+    #[should_panic]
+    fn cant_split_without_key<C: Cipher, H: Hash>() {
+        let mut s1 = SymmetricState::<C, H>::new("complex delirium");
+        s1.mix_hash(b"all wound up");
+        s1.split();
+    }
+
+    #[test]
+    fn symmetric_suites() {
+        // ChaChaPoly
+        symmetric_suite::<ChaChaPoly, Sha256>();
+        symmetric_suite::<ChaChaPoly, Sha512>();
+        symmetric_suite::<ChaChaPoly, Blake2b>();
+        symmetric_suite::<ChaChaPoly, Blake2s>();
+
+        // AES-GCM
+        symmetric_suite::<AesGcm, Sha256>();
+        symmetric_suite::<AesGcm, Sha512>();
+        symmetric_suite::<AesGcm, Blake2b>();
+        symmetric_suite::<AesGcm, Blake2s>();
+    }
+}

src/test/smoke.rs → tests/smoke.rs

@@ -1,15 +1,15 @@
 //! Basic smoke tests - not full coverage on all crypto primitive combinations but good enough
 
-use crate::bytearray::ByteArray;
-use crate::crypto::cipher::{AesGcm, ChaChaPoly};
-use crate::crypto::dh::X25519;
-use crate::crypto::hash::{Blake2b, Sha512};
+use clatter::bytearray::ByteArray;
+use clatter::crypto::cipher::{AesGcm, ChaChaPoly};
+use clatter::crypto::dh::X25519;
+use clatter::crypto::hash::{Blake2b, Sha512};
 #[cfg(feature = "use-argyle-kyber768")]
-use crate::crypto_impl::argyle_software_kyber::Kyber768 as ArgyleKyber;
-use crate::crypto_impl::{pqclean_kyber, rust_crypto_kyber};
-use crate::handshakepattern::*;
-use crate::traits::{Cipher, Dh, Hash, Kem};
-use crate::{Handshaker, NqHandshake, PqHandshake};
+use clatter::crypto::kem::argyle_software_kyber::Kyber768 as ArgyleKyber;
+use clatter::crypto::kem::{pqclean_kyber, rust_crypto_kyber};
+use clatter::handshakepattern::*;
+use clatter::traits::{Cipher, Dh, Hash, Kem};
+use clatter::{Handshaker, NqHandshake, PqHandshake};
 
 const PSKS: &[[u8; 32]] = &[[0; 32], [1; 32], [2; 32], [3; 32]];