Merge pull request #384 from logos-blockchain/Pravdyvy/shared-secret-receiver-usage-fix

Shared secret receiver usage fix

integration_tests/tests/keys_restoration.rs

@@ -1,9 +1,9 @@
 use std::{str::FromStr, time::Duration};
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, format_private_account_id,
-    format_public_account_id, verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx,
+    format_private_account_id, format_public_account_id, verify_commitment_is_in_state,
 };
 use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use log::info;
@@ -15,6 +15,93 @@ use wallet::cli::{
     programs::native_token_transfer::AuthTransferSubcommand,
 };
 
+#[test]
+async fn sync_private_account_with_non_zero_chain_index() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ctx.existing_private_accounts()[0];
+
+    // Create a new private account
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+        cci: None,
+        label: None,
+    }));
+
+    for _ in 0..3 {
+        // Key Tree shift
+        // This way we have account with child index > 0.
+        let result = wallet::cli::execute_subcommand(
+            ctx.wallet_mut(),
+            Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+                cci: None,
+                label: None,
+            })),
+        )
+        .await?;
+        let SubcommandReturnValue::RegisterAccount { account_id: _ } = result else {
+            anyhow::bail!("Expected RegisterAccount return value");
+        };
+    }
+
+    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id,
+    } = sub_ret
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Get the keys for the newly created account
+    let (to_keys, _) = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .get_private_account(to_account_id)
+        .cloned()
+        .context("Failed to get private account")?;
+
+    // Send to this account using claiming path (using npk and vpk instead of account ID)
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(from),
+        to: None,
+        to_npk: Some(hex::encode(to_keys.nullifer_public_key.0)),
+        to_vpk: Some(hex::encode(to_keys.viewing_public_key.0)),
+        amount: 100,
+    });
+
+    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash } = sub_ret else {
+        anyhow::bail!("Expected PrivacyPreservingTransfer return value");
+    };
+
+    let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await;
+
+    // Sync the wallet to claim the new account
+    let command = Command::Account(AccountSubcommand::SyncPrivate {});
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let new_commitment1 = ctx
+        .wallet()
+        .get_private_account_commitment(from)
+        .context("Failed to get private account commitment for sender")?;
+    assert_eq!(tx.message.new_commitments[0], new_commitment1);
+
+    assert_eq!(tx.message.new_commitments.len(), 2);
+    for commitment in tx.message.new_commitments.into_iter() {
+        assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await);
+    }
+
+    let to_res_acc = ctx
+        .wallet()
+        .get_account_private(to_account_id)
+        .context("Failed to get recipient's private account")?;
+    assert_eq!(to_res_acc.balance, 100);
+
+    info!("Successfully transferred using claiming path");
+
+    Ok(())
+}
+
 #[test]
 async fn restore_keys_from_seed() -> Result<()> {
     let mut ctx = TestContext::new().await?;

key_protocol/src/key_management/key_tree/chain_index.rs

@@ -68,6 +68,10 @@ impl ChainIndex {
         &self.0
     }
 
+    pub fn index(&self) -> Option<u32> {
+        self.chain().last().copied()
+    }
+
     pub fn next_in_line(&self) -> ChainIndex {
         let mut chain = self.0.clone();
         // ToDo: Add overflow check

key_protocol/src/key_management/mod.rs

@@ -62,9 +62,10 @@ impl KeyChain {
     pub fn calculate_shared_secret_receiver(
         &self,
         ephemeral_public_key_sender: EphemeralPublicKey,
+        index: Option<u32>,
     ) -> SharedSecretKey {
         SharedSecretKey::new(
-            &self.secret_spending_key.generate_viewing_secret_key(None),
+            &self.secret_spending_key.generate_viewing_secret_key(index),
             &ephemeral_public_key_sender,
         )
     }
@@ -78,6 +79,9 @@ mod tests {
     use rand::RngCore;
 
     use super::*;
+    use crate::key_management::{
+        ephemeral_key_holder::EphemeralKeyHolder, key_tree::KeyTreePrivate,
+    };
 
     #[test]
     fn test_new_os_random() {
@@ -101,8 +105,8 @@ mod tests {
         let ephemeral_public_key_sender = EphemeralPublicKey::from_scalar(scalar);
 
         // Calculate shared secret
-        let _shared_secret =
-            account_id_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
+        let _shared_secret = account_id_key_holder
+            .calculate_shared_secret_receiver(ephemeral_public_key_sender, None);
     }
 
     #[test]
@@ -150,4 +154,40 @@ mod tests {
             hex::encode(viewing_public_key.to_bytes())
         );
     }
+
+    fn account_with_chain_index_2_for_tests() -> KeyChain {
+        let seed = SeedHolder::new_os_random();
+        let mut key_tree_private = KeyTreePrivate::new(&seed);
+
+        // /0
+        key_tree_private.generate_new_node_layered().unwrap();
+        // /1
+        key_tree_private.generate_new_node_layered().unwrap();
+        // /0/0
+        key_tree_private.generate_new_node_layered().unwrap();
+        // /2
+        let (second_child_id, _) = key_tree_private.generate_new_node_layered().unwrap();
+
+        key_tree_private
+            .get_node(second_child_id)
+            .unwrap()
+            .value
+            .0
+            .clone()
+    }
+
+    #[test]
+    fn test_non_trivial_chain_index() {
+        let keys = account_with_chain_index_2_for_tests();
+
+        let eph_key_holder = EphemeralKeyHolder::new(&keys.nullifer_public_key);
+
+        let key_sender = eph_key_holder.calculate_shared_secret_sender(&keys.viewing_public_key);
+        let key_receiver = keys.calculate_shared_secret_receiver(
+            eph_key_holder.generate_ephemeral_public_key(),
+            Some(2),
+        );
+
+        assert_eq!(key_sender.0, key_receiver.0);
+    }
 }

nssa/core/src/encryption/mod.rs

@@ -75,6 +75,17 @@ impl EncryptionScheme {
         Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
 
         let mut cursor = Cursor::new(buffer.as_slice());
-        Account::from_cursor(&mut cursor).ok()
+        Account::from_cursor(&mut cursor)
+            .inspect_err(|err| {
+                println!(
+                    "Failed to decode {ciphertext:?} \n
+                      with secret {:?} ,\n 
+                      commitment {commitment:?} ,\n
+                      and output_index {output_index} ,\n
+                      with error {err:?}",
+                    shared_secret.0
+                )
+            })
+            .ok()
     }
 }

wallet/src/lib.rs

@@ -363,7 +363,7 @@ impl WalletCore {
             );
         let tx = PrivacyPreservingTransaction::new(message, witness_set);
 
-        let shared_secrets = private_account_keys
+        let shared_secrets: Vec<_> = private_account_keys
             .into_iter()
             .map(|keys| keys.ssk)
             .collect();
@@ -419,18 +419,19 @@ impl WalletCore {
             .user_data
             .default_user_private_accounts
             .iter()
-            .map(|(acc_account_id, (key_chain, _))| (*acc_account_id, key_chain))
-            .chain(
-                self.storage
-                    .user_data
-                    .private_key_tree
-                    .key_map
-                    .values()
-                    .map(|keys_node| (keys_node.account_id(), &keys_node.value.0)),
-            );
+            .map(|(acc_account_id, (key_chain, _))| (*acc_account_id, key_chain, None))
+            .chain(self.storage.user_data.private_key_tree.key_map.iter().map(
+                |(chain_index, keys_node)| {
+                    (
+                        keys_node.account_id(),
+                        &keys_node.value.0,
+                        chain_index.index(),
+                    )
+                },
+            ));
 
         let affected_accounts = private_account_key_chains
-            .flat_map(|(acc_account_id, key_chain)| {
+            .flat_map(|(acc_account_id, key_chain, index)| {
                 let view_tag = EncryptedAccountData::compute_view_tag(
                     key_chain.nullifer_public_key.clone(),
                     key_chain.viewing_public_key.clone(),
@@ -444,8 +445,8 @@ impl WalletCore {
                     .filter_map(|(ciph_id, encrypted_data)| {
                         let ciphertext = &encrypted_data.ciphertext;
                         let commitment = &tx.message.new_commitments[ciph_id];
-                        let shared_secret =
-                            key_chain.calculate_shared_secret_receiver(encrypted_data.epk.clone());
+                        let shared_secret = key_chain
+                            .calculate_shared_secret_receiver(encrypted_data.epk.clone(), index);
 
                         nssa_core::EncryptionScheme::decrypt(
                             ciphertext,
@@ -455,6 +456,7 @@ impl WalletCore {
                         )
                     })
                     .map(move |res_acc| (acc_account_id, res_acc))
+                    .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();