From 837cc31a9305ad15e4245a05c32473ef6a0846e3 Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Wed, 7 May 2025 16:34:52 -0300
Subject: [PATCH] use x coordinate as key for aes

---
 .../key_management/ephemeral_key_holder.rs    | 11 ++----
 accounts/src/key_management/mod.rs            | 34 ++++++-------------
 2 files changed, 13 insertions(+), 32 deletions(-)

diff --git a/accounts/src/key_management/ephemeral_key_holder.rs b/accounts/src/key_management/ephemeral_key_holder.rs
index 9e3a9b1..266cfa3 100644
--- a/accounts/src/key_management/ephemeral_key_holder.rs
+++ b/accounts/src/key_management/ephemeral_key_holder.rs
@@ -3,6 +3,7 @@ use elliptic_curve::PrimeField;
 use k256::{AffinePoint, FieldBytes, Scalar};
 use log::info;
 use rand::{rngs::OsRng, RngCore};
+use elliptic_curve::point::AffineCoordinates;
 
 use super::constants_types::{CipherText, Nonce};
 
@@ -39,14 +40,8 @@ impl EphemeralKeyHolder {
         viewing_public_key_receiver: AffinePoint,
         data: &[u8],
     ) -> (CipherText, Nonce) {
-        let key_point = self.calculate_shared_secret_sender(viewing_public_key_receiver);
-        let binding = serde_json::to_vec(&key_point).unwrap();
-        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();
-
-        let cipher = Aes256Gcm::new(&key);
+        let shared_secret = self.calculate_shared_secret_sender(viewing_public_key_receiver);
+        let cipher = Aes256Gcm::new(&shared_secret.x());
         let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
 
         (cipher.encrypt(&nonce, data).unwrap(), nonce)
diff --git a/accounts/src/key_management/mod.rs b/accounts/src/key_management/mod.rs
index 25673f6..79418ec 100644
--- a/accounts/src/key_management/mod.rs
+++ b/accounts/src/key_management/mod.rs
@@ -5,6 +5,7 @@ use ephemeral_key_holder::EphemeralKeyHolder;
 use k256::AffinePoint;
 use log::info;
 use secret_holders::{SeedHolder, TopSecretKeyHolder, UTXOSecretKeyHolder};
+use elliptic_curve::point::AffineCoordinates;
 
 use crate::account_core::PublicKey;
 
@@ -63,14 +64,8 @@ impl AddressKeyHolder {
         ciphertext: CipherText,
         nonce: Nonce,
     ) -> Result<Vec<u8>, aes_gcm::Error> {
-        let key_point = self.calculate_shared_secret_receiver(ephemeral_public_key_sender);
-        let binding = serde_json::to_vec(&key_point).unwrap();
-        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();
-
-        let cipher = Aes256Gcm::new(&key);
+        let shared_secret = self.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+        let cipher = Aes256Gcm::new(&shared_secret.x());
 
         cipher.decrypt(&nonce, ciphertext.as_slice())
     }
@@ -107,18 +102,22 @@ impl AddressKeyHolder {
 
 #[cfg(test)]
 mod tests {
+    use std::io::Read;
+
     use aes_gcm::{
         aead::{Aead, KeyInit, OsRng},
         Aes256Gcm,
     };
     use constants_types::{CipherText, Nonce};
     use constants_types::{NULLIFIER_SECRET_CONST, VIEWING_SECRET_CONST};
-    use elliptic_curve::ff::Field;
+    use elliptic_curve::{ff::Field, PrimeField};
     use elliptic_curve::group::prime::PrimeCurveAffine;
     use k256::{AffinePoint, ProjectivePoint, Scalar};
+    use elliptic_curve::point::AffineCoordinates;
 
     use super::*;
 
+
     #[test]
     fn test_new_os_random() {
         // Ensure that a new AddressKeyHolder instance can be created without errors.
@@ -154,22 +153,14 @@ mod tests {
         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 = serde_json::to_vec(&shared_secret).unwrap();
-        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 cipher = Aes256Gcm::new(&shared_secret.x());
         let nonce = Nonce::from_slice(b"unique nonce");
         let plaintext = b"Sensitive data";
         let ciphertext = cipher
@@ -315,12 +306,7 @@ mod tests {
 
         let shared_secret =
             address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
-        // Prepare the encryption key from shared secret
-        let key_raw = serde_json::to_vec(&shared_secret).unwrap();
-        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 cipher = Aes256Gcm::new(&shared_secret.x());
 
         let ciphertext = cipher
             .encrypt(nonce, plaintext.as_ref())