Merge pull request #20 from vacp2p/key_management_tests

Key management tests

accounts/src/key_management/mod.rs

@@ -62,8 +62,9 @@ impl AddressKeyHolder {
         nonce: Nonce,
     ) -> Vec<u8> {
         let key_point = self.calculate_shared_secret_receiver(ephemeral_public_key_sender);
-        let key_raw = key_point.to_bytes();
-        let key_raw_adjust: [u8; 32] = key_raw.as_slice().try_into().unwrap();
+        let binding = key_point.to_bytes();
+        let key_raw = &binding.as_slice()[..32];
+        let key_raw_adjust: [u8; 32] = key_raw.try_into().unwrap();
 
         let key: Key<Aes256Gcm> = key_raw_adjust.into();
 
@@ -75,11 +76,231 @@ impl AddressKeyHolder {
 
 #[cfg(test)]
 mod tests {
+    use aes_gcm::{
+        aead::{Aead, KeyInit, OsRng},
+        Aes256Gcm,
+    };
+    use constants_types::{CipherText, Nonce};
     use constants_types::{NULLIFIER_SECRET_CONST, VIEVING_SECRET_CONST};
+    use elliptic_curve::ff::Field;
+    use elliptic_curve::group::prime::PrimeCurveAffine;
     use elliptic_curve::group::GroupEncoding;
+    use k256::{AffinePoint, ProjectivePoint, Scalar};
 
     use super::*;
 
+    #[test]
+    fn test_new_os_random() {
+        // Ensure that a new AddressKeyHolder instance can be created without errors.
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Check that key holder fields are initialized with expected types
+        assert!(!Into::<bool>::into(
+            address_key_holder.nullifer_public_key.is_identity()
+        ));
+        assert!(!Into::<bool>::into(
+            address_key_holder.viewing_public_key.is_identity()
+        ));
+    }
+
+    #[test]
+    fn test_calculate_shared_secret_receiver() {
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Generate a random ephemeral public key sender
+        let scalar = Scalar::random(&mut OsRng);
+        let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
+
+        // Calculate shared secret
+        let shared_secret =
+            address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+
+        // Ensure the shared secret is not an identity point (suggesting non-zero output)
+        assert!(!Into::<bool>::into(shared_secret.is_identity()));
+    }
+
+    #[test]
+    fn test_decrypt_data() {
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Generate an ephemeral key and shared secret
+        let scalar = Scalar::random(OsRng);
+        let ephemeral_public_key_sender = address_key_holder
+            .produce_ephemeral_key_holder()
+            .generate_ephemeral_public_key();
+        let shared_secret =
+            address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+
+        // Prepare the encryption key from shared secret
+        let key_raw = shared_secret.to_bytes();
+        let key_raw_adjust_pre = &key_raw.as_slice()[..32];
+        let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
+        let key: Key<Aes256Gcm> = key_raw_adjust.into();
+
+        let cipher = Aes256Gcm::new(&key);
+
+        // Encrypt sample data
+        let nonce = Nonce::from_slice(b"unique nonce");
+        let plaintext = b"Sensitive data";
+        let ciphertext = cipher
+            .encrypt(nonce, plaintext.as_ref())
+            .expect("encryption failure");
+
+        // Attempt decryption
+        let decrypted_data: Vec<u8> = address_key_holder.decrypt_data(
+            ephemeral_public_key_sender,
+            CipherText::from(ciphertext),
+            nonce.clone(),
+        );
+
+        // Verify decryption is successful and matches original plaintext
+        assert_eq!(decrypted_data, plaintext);
+    }
+
+    #[test]
+    fn test_new_os_random_initialization() {
+        // Ensure that AddressKeyHolder is initialized correctly
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Check that key holder fields are initialized with expected types and values
+        assert!(!Into::<bool>::into(
+            address_key_holder.nullifer_public_key.is_identity()
+        ));
+        assert!(!Into::<bool>::into(
+            address_key_holder.viewing_public_key.is_identity()
+        ));
+        assert!(address_key_holder.address.as_slice().len() > 0); // Assume TreeHashType has non-zero length for a valid address
+    }
+
+    #[test]
+    fn test_calculate_shared_secret_with_identity_point() {
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Use identity point as ephemeral public key
+        let identity_point = AffinePoint::identity();
+
+        // Calculate shared secret
+        let shared_secret = address_key_holder.calculate_shared_secret_receiver(identity_point);
+
+        // The shared secret with the identity point should also result in the identity point
+        assert!(Into::<bool>::into(shared_secret.is_identity()));
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_decrypt_data_with_incorrect_nonce() {
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Generate ephemeral public key and shared secret
+        let scalar = Scalar::random(OsRng);
+        let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
+        let shared_secret =
+            address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+
+        // Prepare the encryption key from shared secret
+        let key_raw = shared_secret.to_bytes();
+        let key_raw_adjust_pre = &key_raw.as_slice()[..32];
+        let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
+        let key: Key<Aes256Gcm> = key_raw_adjust.into();
+
+        let cipher = Aes256Gcm::new(&key);
+
+        // Encrypt sample data with a specific nonce
+        let nonce = Nonce::from_slice(b"unique nonce");
+        let plaintext = b"Sensitive data";
+        let ciphertext = cipher
+            .encrypt(nonce, plaintext.as_ref())
+            .expect("encryption failure");
+
+        // Attempt decryption with an incorrect nonce
+        let incorrect_nonce = Nonce::from_slice(b"wrong nonce");
+        let decrypted_data = address_key_holder.decrypt_data(
+            ephemeral_public_key_sender,
+            CipherText::from(ciphertext.clone()),
+            incorrect_nonce.clone(),
+        );
+
+        // The decryption should fail or produce incorrect output due to nonce mismatch
+        assert_ne!(decrypted_data, plaintext);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_decrypt_data_with_incorrect_ciphertext() {
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Generate ephemeral public key and shared secret
+        let scalar = Scalar::random(OsRng);
+        let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
+        let shared_secret =
+            address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+
+        // Prepare the encryption key from shared secret
+        let key_raw = shared_secret.to_bytes();
+        let key_raw_adjust_pre = &key_raw.as_slice()[..32];
+        let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
+        let key: Key<Aes256Gcm> = key_raw_adjust.into();
+
+        let cipher = Aes256Gcm::new(&key);
+
+        // Encrypt sample data
+        let nonce = Nonce::from_slice(b"unique nonce");
+        let plaintext = b"Sensitive data";
+        let ciphertext = cipher
+            .encrypt(nonce, plaintext.as_ref())
+            .expect("encryption failure");
+
+        // Tamper with the ciphertext to simulate corruption
+        let mut corrupted_ciphertext = ciphertext.clone();
+        corrupted_ciphertext[0] ^= 1; // Flip a bit in the ciphertext
+
+        // Attempt decryption
+        let result = address_key_holder.decrypt_data(
+            ephemeral_public_key_sender,
+            CipherText::from(corrupted_ciphertext),
+            nonce.clone(),
+        );
+
+        // The decryption should fail or produce incorrect output due to tampered ciphertext
+        assert_ne!(result, plaintext);
+    }
+
+    #[test]
+    fn test_encryption_decryption_round_trip() {
+        let address_key_holder = AddressKeyHolder::new_os_random();
+
+        // Generate ephemeral key and shared secret
+        let scalar = Scalar::random(OsRng);
+        let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
+
+        // Encrypt sample data
+        let plaintext = b"Round-trip test data";
+        let nonce = Nonce::from_slice(b"unique nonce");
+
+        let shared_secret =
+            address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+        // Prepare the encryption key from shared secret
+        let key_raw = shared_secret.to_bytes();
+        let key_raw_adjust_pre = &key_raw.as_slice()[..32];
+        let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
+        let key: Key<Aes256Gcm> = key_raw_adjust.into();
+        let cipher = Aes256Gcm::new(&key);
+
+        let ciphertext = cipher
+            .encrypt(nonce, plaintext.as_ref())
+            .expect("encryption failure");
+
+        // Decrypt the data using the `AddressKeyHolder` instance
+        let decrypted_data = address_key_holder.decrypt_data(
+            ephemeral_public_key_sender,
+            CipherText::from(ciphertext),
+            nonce.clone(),
+        );
+
+        // Verify the decrypted data matches the original plaintext
+        assert_eq!(decrypted_data, plaintext);
+    }
+
     #[test]
     fn key_generation_test() {
         let seed_holder = SeedHolder::new_os_random();