diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index af016dc6..ddea3ab8 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 a77db776..f21530e0 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-ffi/src/lib.rs b/lez/wallet-ffi/src/lib.rs index 361ae672..c91185e6 100644 --- a/lez/wallet-ffi/src/lib.rs +++ b/lez/wallet-ffi/src/lib.rs @@ -47,6 +47,7 @@ pub mod error; pub mod generic_transaction; pub mod keys; pub mod label; +pub mod pda; pub mod pinata; pub mod program_deployment; pub mod sync; diff --git a/lez/wallet-ffi/src/pda.rs b/lez/wallet-ffi/src/pda.rs new file mode 100644 index 00000000..35a9482b --- /dev/null +++ b/lez/wallet-ffi/src/pda.rs @@ -0,0 +1,142 @@ +use lee::AccountId; + +use crate::{ + error::WalletFfiError, FfiBytes32, FfiNullifierPublicKey, FfiPdaSeed, FfiPrivateAccountKeys, + FfiProgramId, FfiU128, +}; + +/// Produce account id for public PDA. +/// +/// # Parameters +/// - `program_id`: Id of the owner program +/// - `pda_seed`: 32 byte seed +/// +/// # Returns +/// - `FfiBytes32` representing account id bytes +#[no_mangle] +pub extern "C" fn wallet_ffi_account_id_for_public_pda( + program_id: FfiProgramId, + pda_seed: FfiPdaSeed, +) -> FfiBytes32 { + AccountId::for_public_pda(&program_id.data, &pda_seed.into()).into() +} + +/// Produce account id for private PDA. +/// +/// # Parameters +/// - `program_id`: Id of the owner program +/// - `pda_seed`: 32 byte seed +/// - `npk`: 32 byte nullifier public key (can be obtained from +/// `wallet_ffi_get_private_account_keys`) +/// - `viewing_public_key`: pointer to u8 (can be obtained from +/// `wallet_ffi_get_private_account_keys`) +/// - `viewing_public_key_len`: length of a `viewing_public_key` (can be obtained from +/// `wallet_ffi_get_private_account_keys`), must be `1184` +/// - `identifier`: little endian encoded `u128` +/// - `account_id`: valid pointer to `FfiBytes32` +/// +/// # Returns +/// - `Success` on successful parsing +/// - Error code on failure +/// +/// # Safety +/// - `viewing_public_key` must be a valid pointer to a `u8` +/// - `account_id` must be a valid pointer to a `FfiBytes32` struct +#[no_mangle] +pub unsafe extern "C" fn wallet_ffi_account_id_for_private_pda( + program_id: FfiProgramId, + pda_seed: FfiPdaSeed, + npk: FfiNullifierPublicKey, + viewing_public_key: *const u8, + viewing_public_key_len: usize, + identifier: FfiU128, + account_id: *mut FfiBytes32, +) -> WalletFfiError { + if viewing_public_key.is_null() { + return WalletFfiError::NullPointer; + } + + let ffi_private_keys = FfiPrivateAccountKeys { + nullifier_public_key: npk, + viewing_public_key, + viewing_public_key_len, + }; + + let vpk = ffi_private_keys.vpk(); + + if vpk.is_err() { + return vpk.err().unwrap(); + } + + unsafe { + *account_id = AccountId::for_private_pda( + &program_id.data, + &pda_seed.into(), + &ffi_private_keys.npk(), + &vpk.unwrap(), + identifier.into(), + ) + .into(); + } + + WalletFfiError::Success +} + +#[cfg(test)] +mod tests { + use lee::AccountId; + use lee_core::{encryption::ViewingPublicKey, NullifierPublicKey}; + use vault_core::PdaSeed; + + use crate::{ + error::WalletFfiError, + pda::{wallet_ffi_account_id_for_private_pda, wallet_ffi_account_id_for_public_pda}, + FfiBytes32, + }; + + #[test] + fn public_pda_consistent_derivation() { + let program_id = [100_u32, 101, 102, 103, 104, 105, 106, 107]; + let pda_seed = PdaSeed::new([42; 32]); + + let pda_id = AccountId::for_public_pda(&program_id, &pda_seed); + let ffi_pda_id = wallet_ffi_account_id_for_public_pda(program_id.into(), pda_seed.into()); + + assert_eq!(pda_id.into_value(), ffi_pda_id.data); + } + + #[test] + fn private_pda_consistent_derivation() { + let program_id = [100_u32, 101, 102, 103, 104, 105, 106, 107]; + let pda_seed = PdaSeed::new([42; 32]); + let vpk = ViewingPublicKey::from_bytes(vec![43; 1184]).unwrap(); + let npk = NullifierPublicKey([44; 32]); + let identifier = 100_000_u128; + + let pda_id = AccountId::for_private_pda(&program_id, &pda_seed, &npk, &vpk, identifier); + + let vpk_ptr = Box::into_raw(vpk.to_bytes().to_vec().into_boxed_slice()) as *const u8; + + let mut ffi_pda_id_base = FfiBytes32 { data: [0; 32] }; + let ffi_pda_id = &raw mut ffi_pda_id_base; + + let err = unsafe { + wallet_ffi_account_id_for_private_pda( + program_id.into(), + pda_seed.into(), + npk.into(), + vpk_ptr, + 1184, + identifier.into(), + ffi_pda_id, + ) + }; + + assert_eq!(err, WalletFfiError::Success); + + assert_eq!(pda_id.into_value(), unsafe { (*ffi_pda_id).data }); + + let vpk_slice = unsafe { std::slice::from_raw_parts_mut(vpk_ptr.cast_mut(), 1184) }; + drop(unsafe { Box::from_raw(std::ptr::from_mut(vpk_slice)) }); + } +} diff --git a/lez/wallet-ffi/src/types.rs b/lez/wallet-ffi/src/types.rs index 4f041c55..3779ba01 100644 --- a/lez/wallet-ffi/src/types.rs +++ b/lez/wallet-ffi/src/types.rs @@ -8,7 +8,7 @@ use std::{ }; use lee::{Data, ProgramId, SharedSecretKey}; -use lee_core::{encryption::MlKem768EncapsulationKey, NullifierPublicKey}; +use lee_core::{encryption::MlKem768EncapsulationKey, program::PdaSeed, NullifierPublicKey}; use wallet::{account::AccountIdWithPrivacy, AccountIdentity}; use crate::error::WalletFfiError; @@ -29,6 +29,36 @@ pub struct FfiBytes32 { pub data: [u8; 32], } +pub type FfiPdaSeed = FfiBytes32; + +impl From for PdaSeed { + fn from(value: FfiPdaSeed) -> Self { + Self::new(value.data) + } +} + +impl From for FfiPdaSeed { + fn from(value: PdaSeed) -> Self { + Self { + data: *value.as_bytes(), + } + } +} + +pub type FfiNullifierPublicKey = FfiBytes32; + +impl From for NullifierPublicKey { + fn from(value: FfiNullifierPublicKey) -> Self { + Self(value.data) + } +} + +impl From for FfiNullifierPublicKey { + fn from(value: NullifierPublicKey) -> Self { + Self { data: value.0 } + } +} + /// Program ID - 8 u32 values (32 bytes total). #[repr(C)] #[derive(Clone, Copy, Default)] diff --git a/lez/wallet-ffi/wallet_ffi.h b/lez/wallet-ffi/wallet_ffi.h index 8f15a888..4104fb61 100644 --- a/lez/wallet-ffi/wallet_ffi.h +++ b/lez/wallet-ffi/wallet_ffi.h @@ -320,6 +320,10 @@ typedef struct LabelList { enum WalletFfiError error; } LabelList; +typedef struct FfiBytes32 FfiPdaSeed; + +typedef struct FfiBytes32 FfiNullifierPublicKey; + typedef struct FfiCreateWalletOutput { struct WalletHandle *wallet; /** @@ -914,6 +918,50 @@ struct LabelList wallet_ffi_get_all_labels_for_account(struct WalletHandle *hand */ enum WalletFfiError wallet_ffi_free_label_list(struct LabelList *label_list); +/** + * Produce account id for public PDA. + * + * # Parameters + * - `program_id`: Id of the owner program + * - `pda_seed`: 32 byte seed + * + * # Returns + * - `FfiBytes32` representing account id bytes + */ +struct FfiBytes32 wallet_ffi_account_id_for_public_pda(struct FfiProgramId program_id, + FfiPdaSeed pda_seed); + +/** + * Produce account id for private PDA. + * + * # Parameters + * - `program_id`: Id of the owner program + * - `pda_seed`: 32 byte seed + * - `npk`: 32 byte nullifier public key (can be obtained from + * `wallet_ffi_get_private_account_keys`) + * - `viewing_public_key`: pointer to u8 (can be obtained from + * `wallet_ffi_get_private_account_keys`) + * - `viewing_public_key_len`: length of a `viewing_public_key` (can be obtained from + * `wallet_ffi_get_private_account_keys`), must be `1184` + * - `identifier`: little endian encoded `u128` + * - `account_id`: valid pointer to `FfiBytes32` + * + * # Returns + * - `Success` on successful parsing + * - Error code on failure + * + * # Safety + * - `viewing_public_key` must be a valid pointer to a `u8` + * - `account_id` must be a valid pointer to a `FfiBytes32` struct + */ +enum WalletFfiError wallet_ffi_account_id_for_private_pda(struct FfiProgramId program_id, + FfiPdaSeed pda_seed, + FfiNullifierPublicKey npk, + const uint8_t *viewing_public_key, + uintptr_t viewing_public_key_len, + struct FfiU128 identifier, + struct FfiBytes32 *account_id); + /** * Claim a pinata reward using a public transaction. * diff --git a/lez/wallet/src/account_manager.rs b/lez/wallet/src/account_manager.rs index a5f900d6..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, @@ -555,12 +543,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); @@ -574,13 +556,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, @@ -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 e0b89077..fdd208af 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}; @@ -474,26 +474,14 @@ impl WalletCore { Some(Commitment::new(&account_id, account)) } - 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/fixtures/prebuilt_sequencer_db.dump b/test_fixtures/fixtures/prebuilt_sequencer_db.dump index 2e0d4b90..38aeb5ad 100644 Binary files a/test_fixtures/fixtures/prebuilt_sequencer_db.dump and b/test_fixtures/fixtures/prebuilt_sequencer_db.dump differ diff --git a/test_fixtures/src/lib.rs b/test_fixtures/src/lib.rs index 118fa33a..f2d06b69 100644 --- a/test_fixtures/src/lib.rs +++ b/test_fixtures/src/lib.rs @@ -527,10 +527,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() }