From 3f8cd84a88b0830ab0ae278dd3d507e6046040cf Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Mon, 6 Jul 2026 18:36:19 +0000 Subject: [PATCH] feat(sequencer): serve proofs and root in one rpc call --- .../tests/auth_transfer/private.rs | 19 +--- lee/state_machine/src/state/mod.rs | 5 + lez/sequencer/service/protocol/src/lib.rs | 2 +- lez/sequencer/service/rpc/src/lib.rs | 12 +-- lez/sequencer/service/src/service.rs | 17 ++-- lez/wallet/src/account_manager.rs | 97 +++++++++---------- lez/wallet/src/lib.rs | 30 ++---- test_fixtures/src/lib.rs | 4 +- 8 files changed, 86 insertions(+), 100 deletions(-) diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index b543a2f6..61b4ec84 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -12,10 +12,9 @@ use lee::{ privacy_preserving_transaction::circuit::ProgramWithDependencies, program::Program, }; use lee_core::{ - DUMMY_COMMITMENT, DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, - ML_KEM_768_CIPHERTEXT_LEN, Nullifier, NullifierPublicKey, + DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, ML_KEM_768_CIPHERTEXT_LEN, + Nullifier, NullifierPublicKey, account::{Account, AccountWithMetadata}, - compute_digest_for_path, encryption::{EphemeralPublicKey, ViewingPublicKey}, }; use log::info; @@ -662,12 +661,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); @@ -691,12 +685,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 dae8fd31..26653fcb 100644 --- a/lez/wallet/src/account_manager.rs +++ b/lez/wallet/src/account_manager.rs @@ -5,7 +5,7 @@ use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder; 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}, encryption::{EncryptedAccountData, EphemeralPublicKey, ViewingPublicKey}, @@ -227,17 +227,17 @@ impl AccountManager { } prepare_public_keycard_state(wallet, account_id, key_path).await? } - AccountIdentity::PrivateOwned(account_id) => State::Private( - private_key_tree_acc_preparation(wallet, account_id, false).await?, - ), + AccountIdentity::PrivateOwned(account_id) => { + State::Private(private_key_tree_acc_preparation(wallet, account_id, false)?) + } AccountIdentity::PrivateForeign { npk, vpk, identifier, } => State::Private(private_foreign_acc_preparation(npk, vpk, identifier)), - AccountIdentity::PrivatePdaOwned(account_id) => State::Private( - private_key_tree_acc_preparation(wallet, account_id, true).await?, - ), + AccountIdentity::PrivatePdaOwned(account_id) => { + State::Private(private_key_tree_acc_preparation(wallet, account_id, true)?) + } AccountIdentity::PrivatePdaForeign { account_id, npk, @@ -253,12 +253,9 @@ impl AccountManager { identifier, } => { let account_id = lee::AccountId::from((&npk, identifier)); - State::Private( - private_shared_acc_preparation( - wallet, account_id, nsk, npk, vpk, identifier, false, - ) - .await?, - ) + State::Private(private_shared_acc_preparation( + wallet, account_id, nsk, npk, vpk, identifier, false, + )) } AccountIdentity::PrivatePdaShared { account_id, @@ -266,29 +263,15 @@ impl AccountManager { npk, vpk, identifier, - } => State::Private( - private_shared_acc_preparation( - wallet, account_id, nsk, npk, vpk, identifier, true, - ) - .await?, - ), + } => State::Private(private_shared_acc_preparation( + wallet, account_id, nsk, npk, vpk, identifier, true, + )), }; 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, @@ -505,7 +488,34 @@ async fn prepare_public_keycard_state( Ok(State::PublicKeycard { account, key_path }) } -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, @@ -520,12 +530,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); @@ -540,14 +544,14 @@ async fn private_key_tree_acc_preparation( identifier: from_identifier, vpk: from_vpk, pre_state: sender_pre, - proof, + proof: None, ssk, epk, is_pda, }) } -async fn private_shared_acc_preparation( +fn private_shared_acc_preparation( wallet: &WalletCore, account_id: AccountId, nsk: NullifierSecretKey, @@ -555,7 +559,7 @@ async fn private_shared_acc_preparation( vpk: ViewingPublicKey, identifier: Identifier, is_pda: bool, -) -> Result { +) -> AccountPreparedData { let acc = wallet .storage() .key_chain() @@ -565,26 +569,21 @@ 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 eph_holder = EphemeralKeyHolder::new(&vpk); let ssk = eph_holder.calculate_shared_secret_sender(); let epk = eph_holder.ephemeral_public_key().clone(); - Ok(AccountPreparedData { + AccountPreparedData { nsk: Some(nsk), npk, identifier, vpk, pre_state, - proof, + proof: None, ssk, epk, is_pda, - }) + } } fn private_foreign_acc_preparation( diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs index b1a71995..e9295510 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() }