diff --git a/accounts/src/key_management/mod.rs b/accounts/src/key_management/mod.rs index 42988a2..fcfcedd 100644 --- a/accounts/src/key_management/mod.rs +++ b/accounts/src/key_management/mod.rs @@ -62,8 +62,9 @@ impl AddressKeyHolder { nonce: Nonce, ) -> Vec { 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 = 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::::into( + address_key_holder.nullifer_public_key.is_identity() + )); + assert!(!Into::::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::::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 = 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 = 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::::into( + address_key_holder.nullifer_public_key.is_identity() + )); + assert!(!Into::::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::::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 = 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 = 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 = 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();