From e5085479143564c40b057bfb808d25b2044686d3 Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Tue, 26 Aug 2025 13:50:52 -0300
Subject: [PATCH] add npk to kdf

---
 nssa/core/src/lib.rs                               | 13 +++++++++++--
 .../guest/src/bin/privacy_preserving_circuit.rs    |  2 +-
 nssa/src/privacy_preserving_transaction/circuit.rs | 14 +++++++-------
 nssa/src/privacy_preserving_transaction/message.rs |  5 +++--
 4 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/nssa/core/src/lib.rs b/nssa/core/src/lib.rs
index 2856d18..aec9ab9 100644
--- a/nssa/core/src/lib.rs
+++ b/nssa/core/src/lib.rs
@@ -67,9 +67,15 @@ pub struct Ciphertext(Vec<u8>);
 
 impl Ciphertext {
     #[cfg(feature = "host")]
-    pub fn decrypt(self, shared_secret: &[u8; 32], output_index: u32) -> Option<Account> {
+    pub fn decrypt(
+        self,
+        shared_secret: &[u8; 32],
+        npk: &NullifierPublicKey,
+        output_index: u32,
+    ) -> Option<Account> {
         let key = Self::kdf(
             &shared_secret,
+            npk,
             // &ipk,
             // &commitment.to_byte_array(),
             output_index,
@@ -85,7 +91,7 @@ impl Ciphertext {
     pub fn new(
         account: &Account,
         shared_secret: &[u8; 32],
-        // npk: &NullifierPublicKey,
+        npk: &NullifierPublicKey,
         // ipk: &IncomingViewingPublicKey,
         output_index: u32,
     ) -> Self {
@@ -93,6 +99,7 @@ impl Ciphertext {
 
         let key = Self::kdf(
             shared_secret,
+            npk,
             // ipk,
             // &commitment.to_byte_array(),
             output_index,
@@ -106,6 +113,7 @@ impl Ciphertext {
 
     pub fn kdf(
         ss_bytes: &[u8; 32],
+        npk: &NullifierPublicKey,
         // epk: &EphemeralPublicKey,
         // ipk: &IncomingViewingPublicKey,
         // commitment: &[u8; 32],
@@ -115,6 +123,7 @@ impl Ciphertext {
 
         bytes.extend_from_slice(b"NSSA/v0.1/KDF-SHA256");
         bytes.extend_from_slice(ss_bytes);
+        bytes.extend_from_slice(&npk.to_byte_array());
         // bytes.extend_from_slice(&epk.0[..]);
         // bytes.extend_from_slice(&ipk.0[..]);
         // bytes.extend_from_slice(&commitment[..]);
diff --git a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs
index 44071f4..cda5053 100644
--- a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs
+++ b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs
@@ -121,7 +121,7 @@ fn main() {
                     shared_secret,
                     // &commitment_post,
                     // esk,
-                    // Npk,
+                    Npk,
                     // Ipk,
                     output_index,
                 );
diff --git a/nssa/src/privacy_preserving_transaction/circuit.rs b/nssa/src/privacy_preserving_transaction/circuit.rs
index 01eb144..38407ce 100644
--- a/nssa/src/privacy_preserving_transaction/circuit.rs
+++ b/nssa/src/privacy_preserving_transaction/circuit.rs
@@ -180,7 +180,7 @@ mod tests {
 
         let recipient_post = output.ciphertexts[0]
             .clone()
-            .decrypt(&shared_secret, 0)
+            .decrypt(&shared_secret, &recipient_keys.npk(), 0)
             .unwrap();
         assert_eq!(recipient_post, expected_recipient_post);
     }
@@ -266,16 +266,16 @@ mod tests {
         assert_eq!(output.new_nullifiers, expected_new_nullifiers);
         assert_eq!(output.ciphertexts.len(), 2);
 
-        let recipient_post_1 = output.ciphertexts[0]
+        let sender_post = output.ciphertexts[0]
             .clone()
-            .decrypt(&shared_secret_1, 0)
+            .decrypt(&shared_secret_1, &sender_keys.npk(), 0)
             .unwrap();
-        assert_eq!(recipient_post_1, expected_private_account_1);
+        assert_eq!(sender_post, expected_private_account_1);
 
-        let recipient_post_2 = output.ciphertexts[1]
+        let recipient_post = output.ciphertexts[1]
             .clone()
-            .decrypt(&shared_secret_2, 1)
+            .decrypt(&shared_secret_2, &recipient_keys.npk(), 1)
             .unwrap();
-        assert_eq!(recipient_post_2, expected_private_account_2);
+        assert_eq!(recipient_post, expected_private_account_2);
     }
 }
diff --git a/nssa/src/privacy_preserving_transaction/message.rs b/nssa/src/privacy_preserving_transaction/message.rs
index 8d31bc9..245297f 100644
--- a/nssa/src/privacy_preserving_transaction/message.rs
+++ b/nssa/src/privacy_preserving_transaction/message.rs
@@ -9,7 +9,7 @@ use k256::{
 };
 use nssa_core::{
     Ciphertext, CommitmentSetDigest, PrivacyPreservingCircuitOutput, SharedSecretKey,
-    account::{Account, Commitment, Nonce, Nullifier},
+    account::{Account, Commitment, Nonce, Nullifier, NullifierPublicKey},
 };
 use serde::{Deserialize, Serialize};
 
@@ -51,10 +51,11 @@ impl EncryptedAccountData {
         self,
         isk: &[u8; 32],
         epk: &EphemeralPublicKey,
+        npk: &NullifierPublicKey,
         output_index: u32,
     ) -> Option<Account> {
         let shared_secret = Self::compute_shared_secret(isk, &epk);
-        self.ciphertext.decrypt(&shared_secret, output_index)
+        self.ciphertext.decrypt(&shared_secret, npk, output_index)
     }
 
     pub fn compute_shared_secret(scalar: &[u8; 32], point: &Secp256k1Point) -> SharedSecretKey {