diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index 205cbfed..43444538 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; @@ -656,12 +655,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); @@ -686,12 +680,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 fcf71a21..4654b6a2 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 914232c0..fe26460a 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 3a48e7cc..2e7914d4 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; @@ -148,12 +148,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 f156a1ef..98d19a88 100644 --- a/lez/wallet/src/account_manager.rs +++ b/lez/wallet/src/account_manager.rs @@ -4,8 +4,8 @@ use anyhow::Result; use keycard_wallet::{KeycardWallet, python_path}; use lee::{AccountId, PrivateKey, PublicKey, Signature}; use lee_core::{ - CommitmentSetDigest, DUMMY_COMMITMENT_HASH, DummyInput, Identifier, InputAccountIdentity, - MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivateAccountKind, SharedSecretKey, + Commitment, CommitmentSetDigest, DummyInput, Identifier, InputAccountIdentity, MembershipProof, + NullifierPublicKey, NullifierSecretKey, PrivateAccountKind, SharedSecretKey, account::{Account, AccountWithMetadata, Nonce}, encryption::{ Ciphertext, EncryptedAccountData, MlKem768EncapsulationKey, ViewTag, ViewingPublicKey, @@ -254,7 +254,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) } @@ -280,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 +313,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 +326,7 @@ impl AccountManager { } => { let pre = private_shared_acc_preparation( wallet, account_id, nsk, npk, vpk, identifier, true, - ) - .await?; + ); State::Private(pre) } @@ -337,18 +335,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, @@ -555,7 +542,34 @@ struct AccountPreparedData { is_pda: bool, } -async fn private_key_tree_acc_preparation( +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) + .await + .map_err(ExecutionFailureKind::SequencerError)?; + + for (pre, proof) in private.iter_mut().zip(proofs) { + pre.proof = proof; + } + + Ok(root) +} + +fn private_key_tree_acc_preparation( wallet: &WalletCore, account_id: AccountId, is_pda: bool, @@ -570,12 +584,6 @@ async fn private_key_tree_acc_preparation( let from_npk = from_keys.nullifier_public_key; let from_vpk = from_keys.viewing_public_key.clone(); - // TODO: Remove this unwrap, error types must be compatible - let proof = wallet - .check_private_account_initialized(account_id) - .await - .unwrap(); - // TODO: Technically we could allow unauthorized owned accounts, but currently we don't have // support from that in the wallet. let sender_pre = AccountWithMetadata::new(from_acc.account.clone(), true, account_id); @@ -588,13 +596,13 @@ async fn private_key_tree_acc_preparation( 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, @@ -602,7 +610,7 @@ async fn private_shared_acc_preparation( vpk: ViewingPublicKey, identifier: Identifier, is_pda: bool, -) -> Result { +) -> AccountPreparedData { let acc = wallet .storage() .key_chain() @@ -612,23 +620,18 @@ 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 random_seed = random_bytes(); - Ok(AccountPreparedData { + AccountPreparedData { nsk: Some(nsk), npk, identifier, vpk, pre_state, - proof, + proof: None, random_seed, is_pda, - }) + } } /// Generate random byte using OS randomness. diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index f57e9e76..cfc3fe18 100644 --- a/lez/wallet/src/lib.rs +++ b/lez/wallet/src/lib.rs @@ -23,8 +23,8 @@ use lee::{ }, }; use lee_core::{ - Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, SharedSecretKey, - account::Nonce, compute_digest_for_path, program::InstructionData, + Commitment, CommitmentSetDigest, MembershipProof, SharedSecretKey, account::Nonce, + program::InstructionData, }; use log::info; use sequencer_service_rpc::{RpcClient as _, SequencerClient, SequencerClientBuilder}; @@ -479,26 +479,14 @@ impl WalletCore { self.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) { - self.sequencer_client - .get_proof_for_commitment(acc_comm) - .await - .map_err(Into::into) - } else { - Ok(None) - } - } - - pub async fn get_commitment_root(&self) -> Result> { - let proof = self - .sequencer_client - .get_proof_for_commitment(DUMMY_COMMITMENT) - .await?; - Ok(proof.map(|p| compute_digest_for_path(&DUMMY_COMMITMENT, &p))) + commitments: Vec, + ) -> Result<(Vec>, CommitmentSetDigest)> { + self.sequencer_client + .get_proofs_and_root(commitments) + .await + .map_err(Into::into) } pub fn decode_insert_privacy_preserving_transaction_results( diff --git a/test_fixtures/src/lib.rs b/test_fixtures/src/lib.rs index 43e2a9db..4323ba66 100644 --- a/test_fixtures/src/lib.rs +++ b/test_fixtures/src/lib.rs @@ -474,10 +474,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() }