diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index fb2fde5b..655ca1a9 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -12,9 +12,8 @@ use lee::{ program::Program, }; use lee_core::{ - DUMMY_COMMITMENT, DUMMY_COMMITMENT_HASH, InputAccountIdentity, Nullifier, NullifierPublicKey, + DUMMY_COMMITMENT_HASH, InputAccountIdentity, Nullifier, NullifierPublicKey, account::{Account, AccountWithMetadata}, - compute_digest_for_path, encryption::ViewingPublicKey, }; use log::info; @@ -652,12 +651,7 @@ async fn prove_init_with_commitment_root( async fn init_with_dummy_commitment_root_produces_valid_root() -> Result<()> { let ctx = TestContext::new().await?; - let dummy_proof = ctx - .sequencer_client() - .get_proof_for_commitment(DUMMY_COMMITMENT) - .await? - .expect("DUMMY_COMMITMENT must be in genesis commitment set"); - let expected_digest = compute_digest_for_path(&DUMMY_COMMITMENT, &dummy_proof); + let (_, expected_digest) = ctx.sequencer_client().get_proofs_and_root(vec![]).await?; let nsk: lee_core::NullifierSecretKey = [7; 32]; let npk = NullifierPublicKey::from(&nsk); @@ -682,12 +676,7 @@ async fn init_with_dummy_commitment_root_produces_valid_root() -> Result<()> { async fn init_nullifier_digest_is_bound_to_commitment_root() -> Result<()> { let ctx = TestContext::new().await?; - let dummy_proof = ctx - .sequencer_client() - .get_proof_for_commitment(DUMMY_COMMITMENT) - .await? - .expect("DUMMY_COMMITMENT must be in genesis commitment set"); - let expected_digest = compute_digest_for_path(&DUMMY_COMMITMENT, &dummy_proof); + let (_, expected_digest) = ctx.sequencer_client().get_proofs_and_root(vec![]).await?; let output_with_root = prove_init_with_commitment_root(&ctx, expected_digest).await?; let output_without_root = prove_init_with_commitment_root(&ctx, DUMMY_COMMITMENT_HASH).await?; diff --git a/lee/state_machine/src/state/mod.rs b/lee/state_machine/src/state/mod.rs index b0daf97b..94d919ca 100644 --- a/lee/state_machine/src/state/mod.rs +++ b/lee/state_machine/src/state/mod.rs @@ -138,6 +138,11 @@ impl V03State { Self::default() } + #[must_use] + pub fn commitment_root(&self) -> CommitmentSetDigest { + self.private_state.0.digest() + } + /// Initializes state with given public account balances leaving other account fields at their /// default values. #[must_use] diff --git a/lez/sequencer/service/protocol/src/lib.rs b/lez/sequencer/service/protocol/src/lib.rs index 58e300f6..ce669d31 100644 --- a/lez/sequencer/service/protocol/src/lib.rs +++ b/lez/sequencer/service/protocol/src/lib.rs @@ -4,7 +4,7 @@ use std::{fmt::Display, str::FromStr}; pub use common::{HashType, block::Block, transaction::LeeTransaction}; pub use lee::{Account, AccountId, ProgramId}; -pub use lee_core::{BlockId, Commitment, MembershipProof, account::Nonce}; +pub use lee_core::{BlockId, Commitment, CommitmentSetDigest, MembershipProof, account::Nonce}; use serde_with::{DeserializeFromStr, SerializeDisplay}; #[derive(Debug, Clone, PartialEq, Eq, Hash, SerializeDisplay, DeserializeFromStr)] diff --git a/lez/sequencer/service/rpc/src/lib.rs b/lez/sequencer/service/rpc/src/lib.rs index f7aa8b56..f32f6bd2 100644 --- a/lez/sequencer/service/rpc/src/lib.rs +++ b/lez/sequencer/service/rpc/src/lib.rs @@ -6,8 +6,8 @@ use jsonrpsee::types::ErrorObjectOwned; #[cfg(feature = "client")] pub use jsonrpsee::{core::ClientError, http_client::HttpClientBuilder as SequencerClientBuilder}; use sequencer_service_protocol::{ - Account, AccountId, Block, BlockId, ChannelId, Commitment, HashType, LeeTransaction, - MembershipProof, Nonce, ProgramId, + Account, AccountId, Block, BlockId, ChannelId, Commitment, CommitmentSetDigest, HashType, + LeeTransaction, MembershipProof, Nonce, ProgramId, }; #[cfg(all(not(feature = "server"), not(feature = "client")))] @@ -76,11 +76,11 @@ pub trait Rpc { account_ids: Vec, ) -> Result, ErrorObjectOwned>; - #[method(name = "getProofForCommitment")] - async fn get_proof_for_commitment( + #[method(name = "getProofsAndRoot")] + async fn get_proofs_and_root( &self, - commitment: Commitment, - ) -> Result, ErrorObjectOwned>; + commitments: Vec, + ) -> Result<(Vec>, CommitmentSetDigest), ErrorObjectOwned>; #[method(name = "getAccount")] async fn get_account(&self, account_id: AccountId) -> Result; diff --git a/lez/sequencer/service/src/service.rs b/lez/sequencer/service/src/service.rs index 60be5fdd..d0850837 100644 --- a/lez/sequencer/service/src/service.rs +++ b/lez/sequencer/service/src/service.rs @@ -12,8 +12,8 @@ use sequencer_core::{ DbError, SequencerCore, TransactionOrigin, block_publisher::BlockPublisherTrait, }; use sequencer_service_protocol::{ - Account, AccountId, Block, BlockId, ChannelId, Commitment, HashType, MembershipProof, Nonce, - ProgramId, + Account, AccountId, Block, BlockId, ChannelId, Commitment, CommitmentSetDigest, HashType, + MembershipProof, Nonce, ProgramId, }; use tokio::sync::Mutex; @@ -162,12 +162,17 @@ impl sequencer_service_rpc::RpcServer Ok(nonces) } - async fn get_proof_for_commitment( + async fn get_proofs_and_root( &self, - commitment: Commitment, - ) -> Result, ErrorObjectOwned> { + commitments: Vec, + ) -> Result<(Vec>, CommitmentSetDigest), ErrorObjectOwned> { let sequencer = self.sequencer.lock().await; - Ok(sequencer.state().get_proof_for_commitment(&commitment)) + let state = sequencer.state(); + let proofs = commitments + .iter() + .map(|commitment| state.get_proof_for_commitment(commitment)) + .collect(); + Ok((proofs, state.commitment_root())) } async fn get_account(&self, account_id: AccountId) -> Result { diff --git a/lez/wallet/src/account_manager.rs b/lez/wallet/src/account_manager.rs index 20d5586d..824a994d 100644 --- a/lez/wallet/src/account_manager.rs +++ b/lez/wallet/src/account_manager.rs @@ -4,9 +4,10 @@ use anyhow::Result; use keycard_wallet::{KeycardWallet, python_path}; use lee::{AccountId, PrivateKey, PublicKey, Signature}; use lee_core::{ - CommitmentSetDigest, DUMMY_COMMITMENT_HASH, Identifier, InputAccountIdentity, MembershipProof, + Commitment, CommitmentSetDigest, Identifier, InputAccountIdentity, MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, account::{AccountWithMetadata, Nonce}, + compute_digest_for_path, encryption::ViewingPublicKey, }; use rand::{RngCore as _, rngs::OsRng}; @@ -252,7 +253,7 @@ impl AccountManager { State::PublicKeycard { account, key_path } } AccountIdentity::PrivateOwned(account_id) => { - let pre = private_key_tree_acc_preparation(wallet, account_id, false).await?; + let pre = private_key_tree_acc_preparation(wallet, account_id, false)?; State::Private(pre) } @@ -279,7 +280,7 @@ impl AccountManager { State::Private(pre) } AccountIdentity::PrivatePdaOwned(account_id) => { - let pre = private_key_tree_acc_preparation(wallet, account_id, true).await?; + let pre = private_key_tree_acc_preparation(wallet, account_id, true)?; State::Private(pre) } AccountIdentity::PrivatePdaForeign { @@ -313,8 +314,7 @@ impl AccountManager { let account_id = lee::AccountId::from((&npk, &vpk, identifier)); let pre = private_shared_acc_preparation( wallet, account_id, nsk, npk, vpk, identifier, false, - ) - .await?; + ); State::Private(pre) } @@ -327,8 +327,7 @@ impl AccountManager { } => { let pre = private_shared_acc_preparation( wallet, account_id, nsk, npk, vpk, identifier, true, - ) - .await?; + ); State::Private(pre) } @@ -337,18 +336,7 @@ impl AccountManager { states.push(state); } - let has_init_account = states - .iter() - .any(|s| matches!(s, State::Private(pre) if pre.proof.is_none())); - let dummy_commitment_root = if has_init_account { - wallet - .get_commitment_root() - .await - .map_err(ExecutionFailureKind::SequencerError)? - .unwrap_or(DUMMY_COMMITMENT_HASH) - } else { - DUMMY_COMMITMENT_HASH - }; + let dummy_commitment_root = fetch_private_proofs_and_root(wallet, &mut states).await?; Ok(Self { states, @@ -540,7 +528,7 @@ struct AccountPreparedData { is_pda: bool, } -async fn private_key_tree_acc_preparation( +fn private_key_tree_acc_preparation( wallet: &WalletCore, account_id: AccountId, is_pda: bool, @@ -562,25 +550,19 @@ async fn private_key_tree_acc_preparation( let mut random_seed: [u8; 32] = [0; 32]; OsRng.fill_bytes(&mut random_seed); - // TODO: Remove this unwrap, error types must be compatible - let proof = wallet - .check_private_account_initialized(account_id) - .await - .unwrap(); - Ok(AccountPreparedData { nsk: Some(nsk), npk: from_npk, identifier: from_identifier, vpk: from_vpk, pre_state: sender_pre, - proof, + proof: None, random_seed, is_pda, }) } -async fn private_shared_acc_preparation( +fn private_shared_acc_preparation( wallet: &WalletCore, account_id: AccountId, nsk: NullifierSecretKey, @@ -588,7 +570,7 @@ async fn private_shared_acc_preparation( vpk: ViewingPublicKey, identifier: Identifier, is_pda: bool, -) -> Result { +) -> AccountPreparedData { let acc = wallet .storage() .key_chain() @@ -598,24 +580,74 @@ async fn private_shared_acc_preparation( let pre_state = AccountWithMetadata::new(acc, true, account_id); - let proof = wallet - .check_private_account_initialized(account_id) - .await - .unwrap_or(None); - let mut random_seed: [u8; 32] = [0; 32]; OsRng.fill_bytes(&mut random_seed); - Ok(AccountPreparedData { + AccountPreparedData { nsk: Some(nsk), npk, identifier, vpk, pre_state, - proof, + proof: None, random_seed, is_pda, - }) + } +} + +async fn fetch_private_proofs_and_root( + wallet: &WalletCore, + states: &mut [State], +) -> Result { + let (mut private, commitments): (Vec<&mut AccountPreparedData>, Vec) = states + .iter_mut() + .filter_map(|state| match state { + State::Private(pre) => { + let commitment = wallet.get_private_account_commitment(pre.pre_state.account_id)?; + Some((pre, commitment)) + } + State::Public { .. } | State::PublicKeycard { .. } => None, + }) + .unzip(); + + let (proofs, root) = wallet + .get_proofs_and_root(commitments.clone()) + .await + .map_err(ExecutionFailureKind::SequencerError)?; + + validate_proofs_against_root(&commitments, &proofs, root)?; + + for (pre, proof) in private.iter_mut().zip(proofs) { + pre.proof = proof; + } + + Ok(root) +} + +fn validate_proofs_against_root( + commitments: &[Commitment], + proofs: &[Option], + root: CommitmentSetDigest, +) -> Result<(), ExecutionFailureKind> { + if proofs.len() != commitments.len() { + return Err(ExecutionFailureKind::SequencerError(anyhow::anyhow!( + "Sequencer returned {} proofs for {} commitments.", + proofs.len(), + commitments.len(), + ))); + } + + for (commitment, proof) in commitments.iter().zip(proofs) { + if let Some(proof) = proof + && compute_digest_for_path(commitment, proof) != root + { + return Err(ExecutionFailureKind::SequencerError(anyhow::anyhow!( + "Membership proof for {commitment:?} does not reproduce the appropriate root {root:?}.", + ))); + } + } + + Ok(()) } #[cfg(test)] diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index c4b37f61..d9f1b891 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -27,8 +27,8 @@ use lee::{ }, }; use lee_core::{ - BlockId, Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, SharedSecretKey, - account::Nonce, compute_digest_for_path, program::InstructionData, + BlockId, Commitment, CommitmentSetDigest, MembershipProof, SharedSecretKey, account::Nonce, + program::InstructionData, }; use log::info; use sequencer_service_rpc::{RpcClient as _, SequencerClient}; @@ -643,39 +643,22 @@ impl WalletCore { self.optimal_poller().poll_tx(hash).await } - pub async fn check_private_account_initialized( + pub async fn get_proofs_and_root( &self, - account_id: AccountId, - ) -> Result> { - if let Some(acc_comm) = self.get_private_account_commitment(account_id) { - let call_f = - async |client: &SequencerClient| client.get_proof_for_commitment(acc_comm).await; + commitments: Vec, + ) -> Result<(Vec>, CommitmentSetDigest)> { + let call_f = + async |client: &SequencerClient| client.get_proofs_and_root(commitments.clone()).await; - let (call_res, metrics_update) = self.multi_sequencer_client.metered_call(call_f).await; - - { - let mut metric_updates_guard = self.metric_updates.write().await; - metric_updates_guard.push(metrics_update); - } - - call_res.map_err(Into::into) - } else { - Ok(None) - } - } - - pub async fn get_commitment_root(&self) -> Result> { - let call_f = async |client: &SequencerClient| { - client.get_proof_for_commitment(DUMMY_COMMITMENT).await - }; - - let (proof, metrics_update) = self.multi_sequencer_client.metered_call(call_f).await; + let (proofs_and_root, metrics_update) = + self.multi_sequencer_client.metered_call(call_f).await; { let mut metric_updates_guard = self.metric_updates.write().await; metric_updates_guard.push(metrics_update); } - Ok(proof?.map(|p| compute_digest_for_path(&DUMMY_COMMITMENT, &p))) + + Ok(proofs_and_root?) } pub fn decode_insert_privacy_preserving_transaction_results( diff --git a/test_fixtures/src/lib.rs b/test_fixtures/src/lib.rs index e7045d4e..8aa201a7 100644 --- a/test_fixtures/src/lib.rs +++ b/test_fixtures/src/lib.rs @@ -532,10 +532,10 @@ pub async fn verify_commitment_is_in_state( seq_client: &SequencerClient, ) -> bool { seq_client - .get_proof_for_commitment(commitment) + .get_proofs_and_root(vec![commitment]) .await .ok() - .flatten() + .and_then(|(proofs, _)| proofs.into_iter().next().flatten()) .is_some() }