diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a8c588b0..ed0de65d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,6 +2,7 @@ on: push: branches: - main + - dev paths-ignore: - "**.md" - "!.github/workflows/*.yml" @@ -94,7 +95,7 @@ jobs: uses: Swatinem/rust-cache@v2 with: shared-key: ci-rust-cache - save-if: ${{ github.ref == 'refs/heads/main' }} + save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }} - name: Lint workspace env: @@ -125,7 +126,7 @@ jobs: uses: Swatinem/rust-cache@v2 with: shared-key: ci-rust-cache - save-if: ${{ github.ref == 'refs/heads/main' }} + save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }} - name: Install nextest run: cargo install --locked cargo-nextest @@ -156,7 +157,7 @@ jobs: uses: Swatinem/rust-cache@v2 with: shared-key: ci-rust-cache - save-if: ${{ github.ref == 'refs/heads/main' }} + save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }} - name: Install nextest run: cargo install --locked cargo-nextest @@ -245,7 +246,7 @@ jobs: uses: Swatinem/rust-cache@v2 with: shared-key: ci-rust-cache - save-if: ${{ github.ref == 'refs/heads/main' }} + save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }} - name: Test valid proof env: @@ -268,7 +269,7 @@ jobs: uses: Swatinem/rust-cache@v2 with: shared-key: ci-rust-cache - save-if: ${{ github.ref == 'refs/heads/main' }} + save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }} - name: Install just run: cargo install --locked just diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 284e3798..b15693b8 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -57,12 +57,16 @@ Before merging a PR, consider squashing non-meaningful commits. E.g.: Could be squashed to an empty commit if they belong to the same PR. +## Default branch + +By default all PRs must be directed into the `dev` branch. This helps us to keep releases stable. + ## Branch workflow -When bringing your feature branch up to date, prefer rebasing on top of `main`. +When bringing your feature branch up to date, prefer rebasing on top of `dev`. -- Preferred: `git rebase main` -- Avoid: `git merge main` in feature branches +- Preferred: `git rebase dev` +- Avoid: `git merge dev` in feature branches This keeps commit history cleaner and makes reviews easier. diff --git a/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin b/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin index 8db9385b..2b6bee3f 100644 Binary files a/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin and b/artifacts/lee/privacy_preserving_circuit/privacy_preserving_circuit.bin differ diff --git a/artifacts/lez/programs/amm.bin b/artifacts/lez/programs/amm.bin index 00f5343d..50a86ac3 100644 Binary files a/artifacts/lez/programs/amm.bin and b/artifacts/lez/programs/amm.bin differ diff --git a/artifacts/lez/programs/associated_token_account.bin b/artifacts/lez/programs/associated_token_account.bin index 8f4b95ed..e48d0762 100644 Binary files a/artifacts/lez/programs/associated_token_account.bin and b/artifacts/lez/programs/associated_token_account.bin differ diff --git a/artifacts/lez/programs/authenticated_transfer.bin b/artifacts/lez/programs/authenticated_transfer.bin index 4f56b0f7..3236f568 100644 Binary files a/artifacts/lez/programs/authenticated_transfer.bin and b/artifacts/lez/programs/authenticated_transfer.bin differ diff --git a/artifacts/lez/programs/bridge.bin b/artifacts/lez/programs/bridge.bin index e4e4ec5e..16c22fc0 100644 Binary files a/artifacts/lez/programs/bridge.bin and b/artifacts/lez/programs/bridge.bin differ diff --git a/artifacts/lez/programs/clock.bin b/artifacts/lez/programs/clock.bin index 663cc59b..52ed4169 100644 Binary files a/artifacts/lez/programs/clock.bin and b/artifacts/lez/programs/clock.bin differ diff --git a/artifacts/lez/programs/faucet.bin b/artifacts/lez/programs/faucet.bin index b26cfc6f..89e153b5 100644 Binary files a/artifacts/lez/programs/faucet.bin and b/artifacts/lez/programs/faucet.bin differ diff --git a/artifacts/lez/programs/pinata.bin b/artifacts/lez/programs/pinata.bin index f93e2c37..809feb82 100644 Binary files a/artifacts/lez/programs/pinata.bin and b/artifacts/lez/programs/pinata.bin differ diff --git a/artifacts/lez/programs/pinata_token.bin b/artifacts/lez/programs/pinata_token.bin index 1ffa430d..39aa583f 100644 Binary files a/artifacts/lez/programs/pinata_token.bin and b/artifacts/lez/programs/pinata_token.bin differ diff --git a/artifacts/lez/programs/token.bin b/artifacts/lez/programs/token.bin index 210c7c3a..bfd11c88 100644 Binary files a/artifacts/lez/programs/token.bin and b/artifacts/lez/programs/token.bin differ diff --git a/artifacts/lez/programs/vault.bin b/artifacts/lez/programs/vault.bin index e6766101..a08868e5 100644 Binary files a/artifacts/lez/programs/vault.bin and b/artifacts/lez/programs/vault.bin differ diff --git a/integration_tests/src/lib.rs b/integration_tests/src/lib.rs index 07212251..bf395101 100644 --- a/integration_tests/src/lib.rs +++ b/integration_tests/src/lib.rs @@ -6,11 +6,168 @@ use std::time::Duration; use anyhow::{Context as _, Result}; +use key_protocol::key_management::key_tree::chain_index::ChainIndex; +use lee::AccountId; use log::info; +use sequencer_service_rpc::RpcClient as _; pub use test_fixtures::*; +use wallet::{ + cli::{ + CliAccountMention, Command, SubcommandReturnValue, + account::{AccountSubcommand, NewSubcommand}, + programs::{ + native_token_transfer::AuthTransferSubcommand, token::TokenProgramAgnosticSubcommand, + }, + }, + storage::key_chain::FoundPrivateAccount, +}; /// Maximum time to wait for the indexer to catch up to the sequencer. -pub const L2_TO_L1_TIMEOUT: Duration = Duration::from_mins(7); +pub const L2_TO_L1_TIMEOUT: Duration = Duration::from_mins(6); + +/// Create a private or public account at the given chain index and return its ID. +/// Pass `cci: None` to use the wallet's next available chain index. +pub async fn new_account( + ctx: &mut TestContext, + private: bool, + cci: Option, +) -> Result { + let subcommand = if private { + NewSubcommand::Private { cci, label: None } + } else { + NewSubcommand::Public { cci, label: None } + }; + let result = wallet::cli::execute_subcommand( + ctx.wallet_mut(), + Command::Account(AccountSubcommand::New(subcommand)), + ) + .await?; + let SubcommandReturnValue::RegisterAccount { account_id } = result else { + anyhow::bail!("Expected RegisterAccount return value"); + }; + Ok(account_id) +} + +/// Send `amount` from `from` to `to` via an authenticated transfer (identifier 0). +pub async fn send( + ctx: &mut TestContext, + from: CliAccountMention, + to: CliAccountMention, + amount: u128, +) -> Result<()> { + let command = Command::AuthTransfer(AuthTransferSubcommand::Send { + from, + to: Some(to), + to_npk: None, + to_vpk: None, + to_keys: None, + to_identifier: Some(0), + amount, + }); + wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + Ok(()) +} + +/// Create a token (New) and wait for the block to be included. +pub async fn create_token( + ctx: &mut TestContext, + definition_account_id: CliAccountMention, + supply_account_id: CliAccountMention, + name: impl Into, + total_supply: u128, +) -> Result<()> { + let subcommand = TokenProgramAgnosticSubcommand::New { + definition_account_id, + supply_account_id, + name: name.into(), + total_supply, + }; + wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; + info!("Waiting for next block creation"); + tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + Ok(()) +} + +/// Send tokens and wait for the block to be included. +pub async fn token_send( + ctx: &mut TestContext, + from: CliAccountMention, + to: CliAccountMention, + amount: u128, +) -> Result<()> { + let subcommand = TokenProgramAgnosticSubcommand::Send { + from, + to: Some(to), + to_npk: None, + to_vpk: None, + to_keys: None, + to_identifier: Some(0), + amount, + }; + wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; + info!("Waiting for next block creation"); + tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + Ok(()) +} + +/// Retrieve the native token balance for `account_id`. +pub async fn account_balance(ctx: &TestContext, account_id: AccountId) -> Result { + Ok(ctx + .sequencer_client() + .get_account_balance(account_id) + .await?) +} + +/// Fetch the full account state for `account_id` from the sequencer. +pub async fn get_account(ctx: &TestContext, account_id: AccountId) -> Result { + Ok(ctx.sequencer_client().get_account(account_id).await?) +} + +/// Fetch the current commitment for `account_id` and assert it is present in the sequencer state. +pub async fn assert_private_commitment_in_state( + ctx: &TestContext, + account_id: AccountId, + label: &str, +) -> Result<()> { + let commitment = ctx + .wallet() + .get_private_account_commitment(account_id) + .with_context(|| format!("Failed to get commitment for {label}"))?; + assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await); + Ok(()) +} + +/// Sync the wallet's private accounts. +pub async fn sync_private(ctx: &mut TestContext) -> Result<()> { + wallet::cli::execute_subcommand( + ctx.wallet_mut(), + Command::Account(AccountSubcommand::SyncPrivate {}), + ) + .await?; + Ok(()) +} + +/// Look up a restored private account for `account_id`, panicking with `label` if absent. +pub fn restored_private_account<'ctx>( + ctx: &'ctx TestContext, + account_id: AccountId, + label: &str, +) -> FoundPrivateAccount<'ctx> { + ctx.wallet() + .storage() + .key_chain() + .private_account(account_id) + .unwrap_or_else(|| panic!("{label} should be restored")) +} + +/// Assert that a restored public account's signing key exists, panicking with `label` if absent. +pub fn assert_public_account_restored(ctx: &TestContext, account_id: AccountId, label: &str) { + ctx.wallet() + .storage() + .key_chain() + .pub_account_signing_key(account_id) + .unwrap_or_else(|| panic!("{label} should be restored")); +} /// Poll the indexer until its last finalized block id reaches the sequencer's /// current last block id or until [`L2_TO_L1_TIMEOUT`] elapses. diff --git a/integration_tests/tests/account.rs b/integration_tests/tests/account.rs index 0de8a9e2..dd6ba4ed 100644 --- a/integration_tests/tests/account.rs +++ b/integration_tests/tests/account.rs @@ -4,12 +4,11 @@ )] use anyhow::{Context as _, Result}; -use integration_tests::{TestContext, private_mention}; +use integration_tests::{TestContext, get_account, new_account, private_mention}; use key_protocol::key_management::KeyChain; use lee::Data; use lee_core::account::Nonce; use log::info; -use sequencer_service_rpc::RpcClient as _; use tokio::test; use wallet::{ account::{AccountIdWithPrivacy, HumanReadableAccount, Label}, @@ -24,10 +23,7 @@ use wallet::{ async fn get_existing_account() -> Result<()> { let ctx = TestContext::new().await?; - let account = ctx - .sequencer_client() - .get_account(ctx.existing_public_accounts()[0]) - .await?; + let account = get_account(&ctx, ctx.existing_public_accounts()[0]).await?; assert_eq!( account.program_owner, @@ -95,18 +91,7 @@ async fn add_label_to_existing_account() -> Result<()> { async fn new_public_account_without_label() -> Result<()> { let mut ctx = TestContext::new().await?; - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })); - - let result = execute_subcommand(ctx.wallet_mut(), command).await?; - - // Extract the account_id from the result - - let wallet::cli::SubcommandReturnValue::RegisterAccount { account_id } = result else { - panic!("Expected RegisterAccount return value") - }; + let account_id = new_account(&mut ctx, false, None).await?; // Verify no label was stored for the account id assert!( diff --git a/integration_tests/tests/amm.rs b/integration_tests/tests/amm.rs index 9f953001..6e270413 100644 --- a/integration_tests/tests/amm.rs +++ b/integration_tests/tests/amm.rs @@ -7,16 +7,18 @@ use std::time::Duration; use anyhow::Result; -use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, public_mention}; +use integration_tests::{ + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, create_token, get_account, new_account, + public_mention, token_send, +}; use log::info; -use sequencer_service_rpc::RpcClient as _; use tokio::test; use wallet::{ account::Label, cli::{ Command, SubcommandReturnValue, account::{AccountSubcommand, NewSubcommand}, - programs::{amm::AmmProgramAgnosticSubcommand, token::TokenProgramAgnosticSubcommand}, + programs::amm::AmmProgramAgnosticSubcommand, }, }; @@ -25,148 +27,60 @@ async fn amm_public() -> Result<()> { let mut ctx = TestContext::new().await?; // Create new account for the token definition - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id_1, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id_1 = new_account(&mut ctx, false, None).await?; // Create new account for the token supply holder - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id_1, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id_1 = new_account(&mut ctx, false, None).await?; // Create new account for receiving a token transaction - let SubcommandReturnValue::RegisterAccount { - account_id: recipient_account_id_1, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id_1 = new_account(&mut ctx, false, None).await?; // Create new account for the token definition - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id_2, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id_2 = new_account(&mut ctx, false, None).await?; // Create new account for the token supply holder - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id_2, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id_2 = new_account(&mut ctx, false, None).await?; // Create new account for receiving a token transaction - let SubcommandReturnValue::RegisterAccount { - account_id: recipient_account_id_2, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id_2 = new_account(&mut ctx, false, None).await?; // Create new token - let subcommand = TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id_1), - supply_account_id: public_mention(supply_account_id_1), - name: "A NAM1".to_owned(), - - total_supply: 37, - }; - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + create_token( + &mut ctx, + public_mention(definition_account_id_1), + public_mention(supply_account_id_1), + "A NAM1".to_owned(), + 37, + ) + .await?; // Transfer 7 tokens from `supply_acc` to the account at account_id `recipient_account_id_1` - let subcommand = TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id_1), - to: Some(public_mention(recipient_account_id_1)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 7, - }; - - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + token_send( + &mut ctx, + public_mention(supply_account_id_1), + public_mention(recipient_account_id_1), + 7, + ) + .await?; // Create new token - let subcommand = TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id_2), - supply_account_id: public_mention(supply_account_id_2), - name: "A NAM2".to_owned(), - - total_supply: 37, - }; - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + create_token( + &mut ctx, + public_mention(definition_account_id_2), + public_mention(supply_account_id_2), + "A NAM2".to_owned(), + 37, + ) + .await?; // Transfer 7 tokens from `supply_acc` to the account at account_id `recipient_account_id_2` - let subcommand = TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id_2), - to: Some(public_mention(recipient_account_id_2)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 7, - }; - - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + token_send( + &mut ctx, + public_mention(supply_account_id_2), + public_mention(recipient_account_id_2), + 7, + ) + .await?; info!("=================== SETUP FINISHED ==============="); @@ -174,19 +88,7 @@ async fn amm_public() -> Result<()> { // Setup accounts // Create new account for the user holding lp - let SubcommandReturnValue::RegisterAccount { - account_id: user_holding_lp, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let user_holding_lp = new_account(&mut ctx, false, None).await?; // Send creation tx let subcommand = AmmProgramAgnosticSubcommand::New { @@ -201,17 +103,11 @@ async fn amm_public() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let user_holding_a_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_1) - .await?; + let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?; - let user_holding_b_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_2) - .await?; + let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?; - let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?; + let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?; assert_eq!( u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()), @@ -244,17 +140,11 @@ async fn amm_public() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let user_holding_a_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_1) - .await?; + let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?; - let user_holding_b_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_2) - .await?; + let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?; - let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?; + let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?; assert_eq!( u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()), @@ -287,17 +177,11 @@ async fn amm_public() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let user_holding_a_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_1) - .await?; + let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?; - let user_holding_b_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_2) - .await?; + let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?; - let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?; + let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?; assert_eq!( u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()), @@ -331,17 +215,11 @@ async fn amm_public() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let user_holding_a_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_1) - .await?; + let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?; - let user_holding_b_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_2) - .await?; + let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?; - let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?; + let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?; assert_eq!( u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()), @@ -375,17 +253,11 @@ async fn amm_public() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let user_holding_a_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_1) - .await?; + let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?; - let user_holding_b_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_2) - .await?; + let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?; - let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?; + let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?; assert_eq!( u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()), @@ -412,33 +284,9 @@ async fn amm_new_pool_using_labels() -> Result<()> { let mut ctx = TestContext::new().await?; // Create token 1 accounts - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id_1, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id_1 = new_account(&mut ctx, false, None).await?; - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id_1, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id_1 = new_account(&mut ctx, false, None).await?; // Create holding_a with a label let holding_a_label = Label::new("amm-holding-a-label"); @@ -457,33 +305,9 @@ async fn amm_new_pool_using_labels() -> Result<()> { }; // Create token 2 accounts - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id_2, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id_2 = new_account(&mut ctx, false, None).await?; - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id_2, - } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id_2 = new_account(&mut ctx, false, None).await?; // Create holding_b with a label let holding_b_label = Label::new("amm-holding-b-label"); @@ -518,48 +342,40 @@ async fn amm_new_pool_using_labels() -> Result<()> { }; // Create token 1 and distribute to holding_a - let subcommand = TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id_1), - supply_account_id: public_mention(supply_account_id_1), - name: "TOKEN1".to_owned(), - total_supply: 10, - }; - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + create_token( + &mut ctx, + public_mention(definition_account_id_1), + public_mention(supply_account_id_1), + "TOKEN1".to_owned(), + 10, + ) + .await?; - let subcommand = TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id_1), - to: Some(public_mention(holding_a_id)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 5, - }; - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + token_send( + &mut ctx, + public_mention(supply_account_id_1), + public_mention(holding_a_id), + 5, + ) + .await?; // Create token 2 and distribute to holding_b - let subcommand = TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id_2), - supply_account_id: public_mention(supply_account_id_2), - name: "TOKEN2".to_owned(), - total_supply: 10, - }; - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + create_token( + &mut ctx, + public_mention(definition_account_id_2), + public_mention(supply_account_id_2), + "TOKEN2".to_owned(), + 10, + ) + .await?; - let subcommand = TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id_2), - to: Some(public_mention(holding_b_id)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 5, - }; - wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?; - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; + token_send( + &mut ctx, + public_mention(supply_account_id_2), + public_mention(holding_b_id), + 5, + ) + .await?; // Create AMM pool using account labels instead of IDs let subcommand = AmmProgramAgnosticSubcommand::New { @@ -572,7 +388,7 @@ async fn amm_new_pool_using_labels() -> Result<()> { wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?; tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let holding_lp_acc = ctx.sequencer_client().get_account(holding_lp_id).await?; + let holding_lp_acc = get_account(&ctx, holding_lp_id).await?; // LP balance should be 3 (geometric mean of 3, 3) assert_eq!( diff --git a/integration_tests/tests/ata.rs b/integration_tests/tests/ata.rs index 7faac67e..21905fb9 100644 --- a/integration_tests/tests/ata.rs +++ b/integration_tests/tests/ata.rs @@ -9,75 +9,34 @@ use std::time::Duration; use anyhow::{Context as _, Result}; use associated_token_account_core::{compute_ata_seed, get_associated_token_account_id}; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, - verify_commitment_is_in_state, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, create_token, get_account, new_account, + private_mention, public_mention, token_send, verify_commitment_is_in_state, }; use log::info; use sequencer_service_rpc::RpcClient as _; use token_core::{TokenDefinition, TokenHolding}; use tokio::test; -use wallet::cli::{ - Command, SubcommandReturnValue, - account::{AccountSubcommand, NewSubcommand}, - programs::{ata::AtaSubcommand, token::TokenProgramAgnosticSubcommand}, -}; - -/// Create a public account and return its ID. -async fn new_public_account(ctx: &mut TestContext) -> Result { - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { account_id } = result else { - anyhow::bail!("Expected RegisterAccount return value"); - }; - Ok(account_id) -} - -/// Create a private account and return its ID. -async fn new_private_account(ctx: &mut TestContext) -> Result { - 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"); - }; - Ok(account_id) -} +use wallet::cli::{Command, programs::ata::AtaSubcommand}; #[test] async fn create_ata_initializes_holding_account() -> Result<()> { let mut ctx = TestContext::new().await?; - let definition_account_id = new_public_account(&mut ctx).await?; - let supply_account_id = new_public_account(&mut ctx).await?; - let owner_account_id = new_public_account(&mut ctx).await?; + let definition_account_id = new_account(&mut ctx, false, None).await?; + let supply_account_id = new_account(&mut ctx, false, None).await?; + let owner_account_id = new_account(&mut ctx, false, None).await?; // Create a fungible token let total_supply = 100_u128; - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id), - supply_account_id: public_mention(supply_account_id), - name: "TEST".to_owned(), - total_supply, - }), + create_token( + &mut ctx, + public_mention(definition_account_id), + public_mention(supply_account_id), + "TEST".to_owned(), + total_supply, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Create the ATA for owner + definition wallet::cli::execute_subcommand( ctx.wallet_mut(), @@ -121,25 +80,20 @@ async fn create_ata_initializes_holding_account() -> Result<()> { async fn create_ata_is_idempotent() -> Result<()> { let mut ctx = TestContext::new().await?; - let definition_account_id = new_public_account(&mut ctx).await?; - let supply_account_id = new_public_account(&mut ctx).await?; - let owner_account_id = new_public_account(&mut ctx).await?; + let definition_account_id = new_account(&mut ctx, false, None).await?; + let supply_account_id = new_account(&mut ctx, false, None).await?; + let owner_account_id = new_account(&mut ctx, false, None).await?; // Create a fungible token - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id), - supply_account_id: public_mention(supply_account_id), - name: "TEST".to_owned(), - total_supply: 100, - }), + create_token( + &mut ctx, + public_mention(definition_account_id), + public_mention(supply_account_id), + "TEST".to_owned(), + 100, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Create the ATA once wallet::cli::execute_subcommand( ctx.wallet_mut(), @@ -196,28 +150,23 @@ async fn create_ata_is_idempotent() -> Result<()> { async fn transfer_and_burn_via_ata() -> Result<()> { let mut ctx = TestContext::new().await?; - let definition_account_id = new_public_account(&mut ctx).await?; - let supply_account_id = new_public_account(&mut ctx).await?; - let sender_account_id = new_public_account(&mut ctx).await?; - let recipient_account_id = new_public_account(&mut ctx).await?; + let definition_account_id = new_account(&mut ctx, false, None).await?; + let supply_account_id = new_account(&mut ctx, false, None).await?; + let sender_account_id = new_account(&mut ctx, false, None).await?; + let recipient_account_id = new_account(&mut ctx, false, None).await?; let total_supply = 1000_u128; // Create a fungible token, supply goes to supply_account_id - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id), - supply_account_id: public_mention(supply_account_id), - name: "TEST".to_owned(), - total_supply, - }), + create_token( + &mut ctx, + public_mention(definition_account_id), + public_mention(supply_account_id), + "TEST".to_owned(), + total_supply, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Derive ATA addresses let ata_program_id = programs::ata().id(); let sender_ata_id = get_associated_token_account_id( @@ -252,23 +201,14 @@ async fn transfer_and_burn_via_ata() -> Result<()> { // Fund sender's ATA from the supply account (direct token transfer) let fund_amount = 200_u128; - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id), - to: Some(public_mention(sender_ata_id)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: fund_amount, - }), + token_send( + &mut ctx, + public_mention(supply_account_id), + public_mention(sender_ata_id), + fund_amount, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Transfer from sender's ATA to recipient's ATA via the ATA program let transfer_amount = 50_u128; wallet::cli::execute_subcommand( @@ -286,7 +226,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Verify sender ATA balance decreased - let sender_ata_acc = ctx.sequencer_client().get_account(sender_ata_id).await?; + let sender_ata_acc = get_account(&ctx, sender_ata_id).await?; let sender_holding = TokenHolding::try_from(&sender_ata_acc.data)?; assert_eq!( sender_holding, @@ -297,7 +237,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> { ); // Verify recipient ATA balance increased - let recipient_ata_acc = ctx.sequencer_client().get_account(recipient_ata_id).await?; + let recipient_ata_acc = get_account(&ctx, recipient_ata_id).await?; let recipient_holding = TokenHolding::try_from(&recipient_ata_acc.data)?; assert_eq!( recipient_holding, @@ -323,7 +263,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Verify sender ATA balance after burn - let sender_ata_acc = ctx.sequencer_client().get_account(sender_ata_id).await?; + let sender_ata_acc = get_account(&ctx, sender_ata_id).await?; let sender_holding = TokenHolding::try_from(&sender_ata_acc.data)?; assert_eq!( sender_holding, @@ -334,10 +274,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> { ); // Verify the token definition total_supply decreased by burn_amount - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!( token_definition, @@ -355,25 +292,20 @@ async fn transfer_and_burn_via_ata() -> Result<()> { async fn create_ata_with_private_owner() -> Result<()> { let mut ctx = TestContext::new().await?; - let definition_account_id = new_public_account(&mut ctx).await?; - let supply_account_id = new_public_account(&mut ctx).await?; - let owner_account_id = new_private_account(&mut ctx).await?; + let definition_account_id = new_account(&mut ctx, false, None).await?; + let supply_account_id = new_account(&mut ctx, false, None).await?; + let owner_account_id = new_account(&mut ctx, true, None).await?; // Create a fungible token - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id), - supply_account_id: public_mention(supply_account_id), - name: "TEST".to_owned(), - total_supply: 100, - }), + create_token( + &mut ctx, + public_mention(definition_account_id), + public_mention(supply_account_id), + "TEST".to_owned(), + 100, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Create the ATA for the private owner + definition wallet::cli::execute_subcommand( ctx.wallet_mut(), @@ -424,28 +356,23 @@ async fn create_ata_with_private_owner() -> Result<()> { async fn transfer_via_ata_private_owner() -> Result<()> { let mut ctx = TestContext::new().await?; - let definition_account_id = new_public_account(&mut ctx).await?; - let supply_account_id = new_public_account(&mut ctx).await?; - let sender_account_id = new_private_account(&mut ctx).await?; - let recipient_account_id = new_public_account(&mut ctx).await?; + let definition_account_id = new_account(&mut ctx, false, None).await?; + let supply_account_id = new_account(&mut ctx, false, None).await?; + let sender_account_id = new_account(&mut ctx, true, None).await?; + let recipient_account_id = new_account(&mut ctx, false, None).await?; let total_supply = 1000_u128; // Create a fungible token - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id), - supply_account_id: public_mention(supply_account_id), - name: "TEST".to_owned(), - total_supply, - }), + create_token( + &mut ctx, + public_mention(definition_account_id), + public_mention(supply_account_id), + "TEST".to_owned(), + total_supply, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Derive ATA addresses let ata_program_id = programs::ata().id(); let sender_ata_id = get_associated_token_account_id( @@ -480,23 +407,14 @@ async fn transfer_via_ata_private_owner() -> Result<()> { // Fund sender's ATA from the supply account (direct token transfer) let fund_amount = 200_u128; - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id), - to: Some(public_mention(sender_ata_id)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: fund_amount, - }), + token_send( + &mut ctx, + public_mention(supply_account_id), + public_mention(sender_ata_id), + fund_amount, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Transfer from sender's ATA (private owner) to recipient's ATA let transfer_amount = 50_u128; wallet::cli::execute_subcommand( @@ -514,7 +432,7 @@ async fn transfer_via_ata_private_owner() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Verify sender ATA balance decreased - let sender_ata_acc = ctx.sequencer_client().get_account(sender_ata_id).await?; + let sender_ata_acc = get_account(&ctx, sender_ata_id).await?; let sender_holding = TokenHolding::try_from(&sender_ata_acc.data)?; assert_eq!( sender_holding, @@ -525,7 +443,7 @@ async fn transfer_via_ata_private_owner() -> Result<()> { ); // Verify recipient ATA balance increased - let recipient_ata_acc = ctx.sequencer_client().get_account(recipient_ata_id).await?; + let recipient_ata_acc = get_account(&ctx, recipient_ata_id).await?; let recipient_holding = TokenHolding::try_from(&recipient_ata_acc.data)?; assert_eq!( recipient_holding, @@ -549,27 +467,22 @@ async fn transfer_via_ata_private_owner() -> Result<()> { async fn burn_via_ata_private_owner() -> Result<()> { let mut ctx = TestContext::new().await?; - let definition_account_id = new_public_account(&mut ctx).await?; - let supply_account_id = new_public_account(&mut ctx).await?; - let holder_account_id = new_private_account(&mut ctx).await?; + let definition_account_id = new_account(&mut ctx, false, None).await?; + let supply_account_id = new_account(&mut ctx, false, None).await?; + let holder_account_id = new_account(&mut ctx, true, None).await?; let total_supply = 500_u128; // Create a fungible token - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::New { - definition_account_id: public_mention(definition_account_id), - supply_account_id: public_mention(supply_account_id), - name: "TEST".to_owned(), - total_supply, - }), + create_token( + &mut ctx, + public_mention(definition_account_id), + public_mention(supply_account_id), + "TEST".to_owned(), + total_supply, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Derive holder's ATA address let ata_program_id = programs::ata().id(); let holder_ata_id = get_associated_token_account_id( @@ -592,23 +505,14 @@ async fn burn_via_ata_private_owner() -> Result<()> { // Fund holder's ATA from the supply account let fund_amount = 300_u128; - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Token(TokenProgramAgnosticSubcommand::Send { - from: public_mention(supply_account_id), - to: Some(public_mention(holder_ata_id)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: fund_amount, - }), + token_send( + &mut ctx, + public_mention(supply_account_id), + public_mention(holder_ata_id), + fund_amount, ) .await?; - info!("Waiting for next block creation"); - tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - // Burn from holder's ATA (private owner) let burn_amount = 100_u128; wallet::cli::execute_subcommand( @@ -625,7 +529,7 @@ async fn burn_via_ata_private_owner() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Verify holder ATA balance after burn - let holder_ata_acc = ctx.sequencer_client().get_account(holder_ata_id).await?; + let holder_ata_acc = get_account(&ctx, holder_ata_id).await?; let holder_holding = TokenHolding::try_from(&holder_ata_acc.data)?; assert_eq!( holder_holding, @@ -636,10 +540,7 @@ async fn burn_via_ata_private_owner() -> Result<()> { ); // Verify the token definition total_supply decreased by burn_amount - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!( token_definition, diff --git a/integration_tests/tests/auth_transfer/private.rs b/integration_tests/tests/auth_transfer/private.rs index 30f0cfdd..b543a2f6 100644 --- a/integration_tests/tests/auth_transfer/private.rs +++ b/integration_tests/tests/auth_transfer/private.rs @@ -3,16 +3,19 @@ use std::time::Duration; use anyhow::{Context as _, Result}; use common::transaction::LeeTransaction; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx, private_mention, - public_mention, verify_commitment_is_in_state, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, + assert_private_commitment_in_state, fetch_privacy_preserving_tx, get_account, new_account, + private_mention, public_mention, send, sync_private, verify_commitment_is_in_state, }; use lee::{ AccountId, SharedSecretKey, execute_and_prove, privacy_preserving_transaction::circuit::ProgramWithDependencies, program::Program, }; use lee_core::{ - EncryptedAccountData, InputAccountIdentity, NullifierPublicKey, - account::AccountWithMetadata, + DUMMY_COMMITMENT, 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; @@ -34,32 +37,13 @@ async fn private_transfer_to_owned_account() -> Result<()> { let from: AccountId = ctx.existing_private_accounts()[0]; let to: AccountId = ctx.existing_private_accounts()[1]; - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: private_mention(from), - to: Some(private_mention(to)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send(&mut ctx, private_mention(from), private_mention(to), 100).await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let new_commitment1 = ctx - .wallet() - .get_private_account_commitment(from) - .context("Failed to get private account commitment for sender")?; - assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await); - - let new_commitment2 = ctx - .wallet() - .get_private_account_commitment(to) - .context("Failed to get private account commitment for receiver")?; - assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, from, "sender").await?; + assert_private_commitment_in_state(&ctx, to, "receiver").await?; info!("Successfully transferred privately to owned account"); @@ -125,17 +109,7 @@ async fn deshielded_transfer_to_public_account() -> Result<()> { .context("Failed to get sender's private account")?; assert_eq!(from_acc.balance, 10000); - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: private_mention(from), - to: Some(public_mention(to)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send(&mut ctx, private_mention(from), public_mention(to), 100).await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; @@ -144,13 +118,9 @@ async fn deshielded_transfer_to_public_account() -> Result<()> { .wallet() .get_account_private(from) .context("Failed to get sender's private account")?; - let new_commitment = ctx - .wallet() - .get_private_account_commitment(from) - .context("Failed to get private account commitment")?; - assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, from, "sender").await?; - let acc_2_balance = ctx.sequencer_client().get_account_balance(to).await?; + let acc_2_balance = account_balance(&ctx, to).await?; assert_eq!(from_acc.balance, 9900); assert_eq!(acc_2_balance, 20100); @@ -167,18 +137,7 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> { 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, - })); - - 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"); - }; + let to_account_id = new_account(&mut ctx, true, None).await?; // Get the keys for the newly created account let to = ctx @@ -207,14 +166,13 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> { 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?; + sync_private(&mut ctx).await?; - let new_commitment1 = ctx + let sender_commitment = 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[0], sender_commitment); assert_eq!(tx.message.new_commitments.len(), 2); for commitment in tx.message.new_commitments { @@ -239,17 +197,7 @@ async fn shielded_transfer_to_owned_private_account() -> Result<()> { let from: AccountId = ctx.existing_public_accounts()[0]; let to: AccountId = ctx.existing_private_accounts()[1]; - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(from), - to: Some(private_mention(to)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send(&mut ctx, public_mention(from), private_mention(to), 100).await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; @@ -258,13 +206,9 @@ async fn shielded_transfer_to_owned_private_account() -> Result<()> { .wallet() .get_account_private(to) .context("Failed to get receiver's private account")?; - let new_commitment = ctx - .wallet() - .get_private_account_commitment(to) - .context("Failed to get receiver's commitment")?; - assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, to, "receiver").await?; - let acc_from_balance = ctx.sequencer_client().get_account_balance(from).await?; + let acc_from_balance = account_balance(&ctx, from).await?; assert_eq!(acc_from_balance, 9900); assert_eq!(acc_to.balance, 20100); @@ -303,7 +247,7 @@ async fn shielded_transfer_to_foreign_account() -> Result<()> { let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await; - let acc_1_balance = ctx.sequencer_client().get_account_balance(from).await?; + let acc_1_balance = account_balance(&ctx, from).await?; assert!( verify_commitment_is_in_state( @@ -332,18 +276,7 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> { 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, - })); - 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!("Failed to register account"); - }; + let to_account_id = new_account(&mut ctx, true, None).await?; // Get the newly created account's keys let to = ctx @@ -396,14 +329,7 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> { async fn initialize_private_account() -> Result<()> { let mut ctx = TestContext::new().await?; - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: None, - label: None, - })); - let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - let SubcommandReturnValue::RegisterAccount { account_id } = result else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let account_id = new_account(&mut ctx, true, None).await?; let command = Command::AuthTransfer(AuthTransferSubcommand::Init { account_id: private_mention(account_id), @@ -413,14 +339,9 @@ async fn initialize_private_account() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Syncing private accounts"); - let command = Command::Account(AccountSubcommand::SyncPrivate {}); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + sync_private(&mut ctx).await?; - let new_commitment = ctx - .wallet() - .get_private_account_commitment(account_id) - .context("Failed to get private account commitment")?; - assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, account_id, "account").await?; let account = ctx .wallet() @@ -455,32 +376,19 @@ async fn private_transfer_using_from_label() -> Result<()> { wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; // Send using the label instead of account ID - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: CliAccountMention::Label(label), - to: Some(private_mention(to)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send( + &mut ctx, + CliAccountMention::Label(label), + private_mention(to), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let new_commitment1 = ctx - .wallet() - .get_private_account_commitment(from) - .context("Failed to get private account commitment for sender")?; - assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await); - - let new_commitment2 = ctx - .wallet() - .get_private_account_commitment(to) - .context("Failed to get private account commitment for receiver")?; - assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, from, "sender").await?; + assert_private_commitment_in_state(&ctx, to, "receiver").await?; info!("Successfully transferred privately using from_label"); @@ -510,14 +418,9 @@ async fn initialize_private_account_using_label() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let command = Command::Account(AccountSubcommand::SyncPrivate {}); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + sync_private(&mut ctx).await?; - let new_commitment = ctx - .wallet() - .get_private_account_commitment(account_id) - .context("Failed to get private account commitment")?; - assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, account_id, "account").await?; let account = ctx .wallet() @@ -593,11 +496,7 @@ async fn shielded_transfers_to_two_identifiers_same_npk() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::SyncPrivate {}), - ) - .await?; + sync_private(&mut ctx).await?; // Both accounts must be discovered with the correct balances. let account_id_1 = AccountId::for_regular_private_account(&npk, identifier_1); @@ -664,7 +563,7 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> { let npk = NullifierPublicKey::from(&nsk); let vpk = ViewingPublicKey::from_bytes(vec![4_u8; 1184]).unwrap(); let ssk = SharedSecretKey([55_u8; 32]); - let epk = EphemeralPublicKey(vec![55_u8; 1088]); + let epk = EphemeralPublicKey(vec![55_u8; ML_KEM_768_CIPHERTEXT_LEN]); let attacker_vault_id = { let seed = vault_core::compute_vault_seed(attacker_id); AccountId::for_private_pda(&vault_program_id, &seed, &npk, 1337) @@ -672,16 +571,12 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> { let amount: u128 = 1; let faucet_pre = AccountWithMetadata::new( - ctx.sequencer_client() - .get_account(faucet_account_id) - .await?, + get_account(&ctx, faucet_account_id).await?, false, faucet_account_id, ); let vault_pda_pre = AccountWithMetadata::new( - ctx.sequencer_client() - .get_account(attacker_vault_id) - .await?, + get_account(&ctx, attacker_vault_id).await?, false, attacker_vault_id, ); @@ -710,6 +605,7 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> { npk, ssk, identifier: 1337, + commitment_root: DUMMY_COMMITMENT_HASH, seed: None, }, ], @@ -720,3 +616,100 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> { Ok(()) } + +async fn prove_init_with_commitment_root( + ctx: &TestContext, + commitment_root: lee_core::CommitmentSetDigest, +) -> Result { + let program = programs::authenticated_transfer(); + let sender_id = ctx.existing_public_accounts()[0]; + let sender_pre = AccountWithMetadata::new( + ctx.sequencer_client().get_account(sender_id).await?, + true, + sender_id, + ); + + let nsk: lee_core::NullifierSecretKey = [7; 32]; + let npk = NullifierPublicKey::from(&nsk); + let vpk = ViewingPublicKey::from_bytes(vec![4_u8; 1184]).unwrap(); + let ssk = SharedSecretKey([55_u8; 32]); + let recipient_account_id = AccountId::for_regular_private_account(&npk, 0); + let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id); + + let (output, _) = execute_and_prove( + vec![sender_pre, recipient], + Program::serialize_instruction(authenticated_transfer_core::Instruction::Transfer { + amount: 1, + })?, + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &vpk), + npk, + ssk, + identifier: 0, + commitment_root, + }, + ], + &program.into(), + )?; + + Ok(output) +} + +#[test] +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 nsk: lee_core::NullifierSecretKey = [7; 32]; + let npk = NullifierPublicKey::from(&nsk); + let recipient_account_id = AccountId::for_regular_private_account(&npk, 0); + + let output = prove_init_with_commitment_root(&ctx, expected_digest).await?; + + assert_eq!(output.new_nullifiers.len(), 1); + let (nullifier, digest) = &output.new_nullifiers[0]; + assert_eq!( + *nullifier, + Nullifier::for_account_initialization(&recipient_account_id) + ); + assert_eq!(*digest, expected_digest); + assert_ne!(*digest, DUMMY_COMMITMENT_HASH); + + Ok(()) +} + +#[test] +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 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?; + + assert_eq!(output_with_root.new_nullifiers[0].1, expected_digest); + assert_eq!( + output_without_root.new_nullifiers[0].1, + DUMMY_COMMITMENT_HASH + ); + assert_ne!( + output_with_root.new_nullifiers[0].1, + output_without_root.new_nullifiers[0].1, + ); + + Ok(()) +} diff --git a/integration_tests/tests/auth_transfer/public.rs b/integration_tests/tests/auth_transfer/public.rs index 5bbf0954..c42bc20c 100644 --- a/integration_tests/tests/auth_transfer/public.rs +++ b/integration_tests/tests/auth_transfer/public.rs @@ -2,7 +2,10 @@ use std::time::Duration; use anyhow::Result; use common::transaction::LeeTransaction; -use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, public_mention}; +use integration_tests::{ + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, get_account, new_account, + public_mention, send, +}; use lee::public_transaction; use log::info; use sequencer_service_rpc::RpcClient as _; @@ -10,8 +13,7 @@ use tokio::test; use wallet::{ account::Label, cli::{ - CliAccountMention, Command, SubcommandReturnValue, - account::{AccountSubcommand, NewSubcommand}, + CliAccountMention, Command, account::AccountSubcommand, programs::native_token_transfer::AuthTransferSubcommand, }, }; @@ -20,30 +22,22 @@ use wallet::{ async fn successful_transfer_to_existing_account() -> Result<()> { let mut ctx = TestContext::new().await?; - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(ctx.existing_public_accounts()[0]), - to: Some(public_mention(ctx.existing_public_accounts()[1])), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let sender = ctx.existing_public_accounts()[0]; + let receiver = ctx.existing_public_accounts()[1]; + send( + &mut ctx, + public_mention(sender), + public_mention(receiver), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[1]) - .await?; + let acc_1_balance = account_balance(&ctx, sender).await?; + let acc_2_balance = account_balance(&ctx, receiver).await?; info!("Balance of sender: {acc_1_balance:#?}"); info!("Balance of receiver: {acc_2_balance:#?}"); @@ -58,51 +52,23 @@ async fn successful_transfer_to_existing_account() -> Result<()> { pub async fn successful_transfer_to_new_account() -> Result<()> { let mut ctx = TestContext::new().await?; - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })); + let new_persistent_account_id = new_account(&mut ctx, false, None).await?; - wallet::cli::execute_subcommand(ctx.wallet_mut(), command) - .await - .unwrap(); - - let new_persistent_account_id = ctx - .wallet() - .storage() - .key_chain() - .public_account_ids() - .map(|(account_id, _)| account_id) - .find(|acc_id| { - *acc_id != ctx.existing_public_accounts()[0] - && *acc_id != ctx.existing_public_accounts()[1] - }) - .expect("Failed to find newly created account in the wallet storage"); - - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(ctx.existing_public_accounts()[0]), - to: Some(public_mention(new_persistent_account_id)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let sender = ctx.existing_public_accounts()[0]; + send( + &mut ctx, + public_mention(sender), + public_mention(new_persistent_account_id), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(new_persistent_account_id) - .await?; + let acc_1_balance = account_balance(&ctx, sender).await?; + let acc_2_balance = account_balance(&ctx, new_persistent_account_id).await?; info!("Balance of sender: {acc_1_balance:#?}"); info!("Balance of receiver: {acc_2_balance:#?}"); @@ -134,14 +100,8 @@ async fn failed_transfer_with_insufficient_balance() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking balances unchanged"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[1]) - .await?; + let acc_1_balance = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?; + let acc_2_balance = account_balance(&ctx, ctx.existing_public_accounts()[1]).await?; info!("Balance of sender: {acc_1_balance:#?}"); info!("Balance of receiver: {acc_2_balance:#?}"); @@ -156,31 +116,24 @@ async fn failed_transfer_with_insufficient_balance() -> Result<()> { async fn two_consecutive_successful_transfers() -> Result<()> { let mut ctx = TestContext::new().await?; - // First transfer - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(ctx.existing_public_accounts()[0]), - to: Some(public_mention(ctx.existing_public_accounts()[1])), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); + let sender = ctx.existing_public_accounts()[0]; + let receiver = ctx.existing_public_accounts()[1]; - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + // First transfer + send( + &mut ctx, + public_mention(sender), + public_mention(receiver), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move after first transfer"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[1]) - .await?; + let acc_1_balance = account_balance(&ctx, sender).await?; + let acc_2_balance = account_balance(&ctx, receiver).await?; info!("Balance of sender: {acc_1_balance:#?}"); info!("Balance of receiver: {acc_2_balance:#?}"); @@ -191,30 +144,20 @@ async fn two_consecutive_successful_transfers() -> Result<()> { info!("First TX Success!"); // Second transfer - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(ctx.existing_public_accounts()[0]), - to: Some(public_mention(ctx.existing_public_accounts()[1])), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send( + &mut ctx, + public_mention(sender), + public_mention(receiver), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move after second transfer"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[1]) - .await?; + let acc_1_balance = account_balance(&ctx, sender).await?; + let acc_2_balance = account_balance(&ctx, receiver).await?; info!("Balance of sender: {acc_1_balance:#?}"); info!("Balance of receiver: {acc_2_balance:#?}"); @@ -231,14 +174,7 @@ async fn two_consecutive_successful_transfers() -> Result<()> { async fn initialize_public_account() -> Result<()> { let mut ctx = TestContext::new().await?; - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })); - let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - let SubcommandReturnValue::RegisterAccount { account_id } = result else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let account_id = new_account(&mut ctx, false, None).await?; let command = Command::AuthTransfer(AuthTransferSubcommand::Init { account_id: public_mention(account_id), @@ -246,7 +182,7 @@ async fn initialize_public_account() -> Result<()> { wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; info!("Checking correct execution"); - let account = ctx.sequencer_client().get_account(account_id).await?; + let account = get_account(&ctx, account_id).await?; assert_eq!( account.program_owner, @@ -274,30 +210,22 @@ async fn successful_transfer_using_from_label() -> Result<()> { wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; // Send using the label instead of account ID - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: CliAccountMention::Label(label), - to: Some(public_mention(ctx.existing_public_accounts()[1])), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let sender = ctx.existing_public_accounts()[0]; + let receiver = ctx.existing_public_accounts()[1]; + send( + &mut ctx, + CliAccountMention::Label(label), + public_mention(receiver), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[1]) - .await?; + let acc_1_balance = account_balance(&ctx, sender).await?; + let acc_2_balance = account_balance(&ctx, receiver).await?; assert_eq!(acc_1_balance, 9900); assert_eq!(acc_2_balance, 20100); @@ -320,30 +248,22 @@ async fn successful_transfer_using_to_label() -> Result<()> { wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; // Send using the label for the recipient - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(ctx.existing_public_accounts()[0]), - to: Some(CliAccountMention::Label(label)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let sender = ctx.existing_public_accounts()[0]; + let receiver = ctx.existing_public_accounts()[1]; + send( + &mut ctx, + public_mention(sender), + CliAccountMention::Label(label), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move"); - let acc_1_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; - let acc_2_balance = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[1]) - .await?; + let acc_1_balance = account_balance(&ctx, sender).await?; + let acc_2_balance = account_balance(&ctx, receiver).await?; assert_eq!(acc_1_balance, 9900); assert_eq!(acc_2_balance, 20100); @@ -359,14 +279,8 @@ async fn cannot_transfer_funds_from_system_faucet_account() -> Result<()> { let faucet_account_id = system_accounts::faucet_account_id(); let recipient = ctx.existing_public_accounts()[0]; - let recipient_balance_before = ctx - .sequencer_client() - .get_account_balance(recipient) - .await?; - let faucet_balance_before = ctx - .sequencer_client() - .get_account_balance(faucet_account_id) - .await?; + let recipient_balance_before = account_balance(&ctx, recipient).await?; + let faucet_balance_before = account_balance(&ctx, faucet_account_id).await?; let amount = 1_u128; let message = public_transaction::Message::try_new( @@ -387,14 +301,8 @@ async fn cannot_transfer_funds_from_system_faucet_account() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let recipient_balance_after = ctx - .sequencer_client() - .get_account_balance(recipient) - .await?; - let faucet_balance_after = ctx - .sequencer_client() - .get_account_balance(faucet_account_id) - .await?; + let recipient_balance_after = account_balance(&ctx, recipient).await?; + let faucet_balance_after = account_balance(&ctx, faucet_account_id).await?; let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?; assert_eq!(recipient_balance_after, recipient_balance_before); @@ -413,14 +321,8 @@ async fn cannot_execute_faucet_program() -> Result<()> { let vault_program_id = programs::vault().id(); let recipient_vault_id = vault_core::compute_vault_account_id(vault_program_id, recipient); - let recipient_balance_before = ctx - .sequencer_client() - .get_account_balance(recipient) - .await?; - let faucet_balance_before = ctx - .sequencer_client() - .get_account_balance(faucet_account_id) - .await?; + let recipient_balance_before = account_balance(&ctx, recipient).await?; + let faucet_balance_before = account_balance(&ctx, faucet_account_id).await?; let amount = 1_u128; let message = public_transaction::Message::try_new( @@ -445,14 +347,8 @@ async fn cannot_execute_faucet_program() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let recipient_balance_after = ctx - .sequencer_client() - .get_account_balance(recipient) - .await?; - let faucet_balance_after = ctx - .sequencer_client() - .get_account_balance(faucet_account_id) - .await?; + let recipient_balance_after = account_balance(&ctx, recipient).await?; + let faucet_balance_after = account_balance(&ctx, faucet_account_id).await?; let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?; assert_eq!(recipient_balance_after, recipient_balance_before); @@ -493,28 +389,16 @@ async fn user_tx_that_chain_calls_faucet_is_dropped() -> Result<()> { lee::public_transaction::WitnessSet::from_raw_parts(vec![]), )); - let faucet_balance_before = ctx - .sequencer_client() - .get_account_balance(faucet_account_id) - .await?; - let vault_balance_before = ctx - .sequencer_client() - .get_account_balance(attacker_vault_id) - .await?; + let faucet_balance_before = account_balance(&ctx, faucet_account_id).await?; + let vault_balance_before = account_balance(&ctx, attacker_vault_id).await?; let tx_hash = ctx.sequencer_client().send_transaction(attack_tx).await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let faucet_balance_after = ctx - .sequencer_client() - .get_account_balance(faucet_account_id) - .await?; - let vault_balance_after = ctx - .sequencer_client() - .get_account_balance(attacker_vault_id) - .await?; + let faucet_balance_after = account_balance(&ctx, faucet_account_id).await?; + let vault_balance_after = account_balance(&ctx, attacker_vault_id).await?; let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?; assert_eq!(faucet_balance_after, faucet_balance_before); diff --git a/integration_tests/tests/bridge.rs b/integration_tests/tests/bridge.rs index 9efdd641..9c241ebb 100644 --- a/integration_tests/tests/bridge.rs +++ b/integration_tests/tests/bridge.rs @@ -11,7 +11,8 @@ use borsh::BorshSerialize; use common::transaction::LeeTransaction; use futures::StreamExt as _; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, wait_for_indexer_to_catch_up, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, get_account, + wait_for_indexer_to_catch_up, }; use lee::{ AccountId, execute_and_prove, privacy_preserving_transaction, program::Program, @@ -65,27 +66,15 @@ async fn public_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> { lee::public_transaction::WitnessSet::from_raw_parts(vec![]), )); - let bridge_balance_before = ctx - .sequencer_client() - .get_account_balance(bridge_account_id) - .await?; - let vault_balance_before = ctx - .sequencer_client() - .get_account_balance(recipient_vault_id) - .await?; + let bridge_balance_before = account_balance(&ctx, bridge_account_id).await?; + let vault_balance_before = account_balance(&ctx, recipient_vault_id).await?; let tx_hash = ctx.sequencer_client().send_transaction(attack_tx).await?; tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let bridge_balance_after = ctx - .sequencer_client() - .get_account_balance(bridge_account_id) - .await?; - let vault_balance_after = ctx - .sequencer_client() - .get_account_balance(recipient_vault_id) - .await?; + let bridge_balance_after = account_balance(&ctx, bridge_account_id).await?; + let vault_balance_after = account_balance(&ctx, recipient_vault_id).await?; let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?; assert_eq!(bridge_balance_after, bridge_balance_before); @@ -169,16 +158,12 @@ async fn private_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> { // Get pre-state of bridge and vault accounts let bridge_pre = AccountWithMetadata::new( - ctx.sequencer_client() - .get_account(bridge_account_id) - .await?, + get_account(&ctx, bridge_account_id).await?, false, bridge_account_id, ); let vault_pre = AccountWithMetadata::new( - ctx.sequencer_client() - .get_account(recipient_vault_id) - .await?, + get_account(&ctx, recipient_vault_id).await?, false, recipient_vault_id, ); @@ -229,27 +214,15 @@ async fn private_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> { witness_set, )); - let bridge_balance_before = ctx - .sequencer_client() - .get_account_balance(bridge_account_id) - .await?; - let vault_balance_before = ctx - .sequencer_client() - .get_account_balance(recipient_vault_id) - .await?; + let bridge_balance_before = account_balance(&ctx, bridge_account_id).await?; + let vault_balance_before = account_balance(&ctx, recipient_vault_id).await?; let tx_hash = ctx.sequencer_client().send_transaction(attack_tx).await?; tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let bridge_balance_after = ctx - .sequencer_client() - .get_account_balance(bridge_account_id) - .await?; - let vault_balance_after = ctx - .sequencer_client() - .get_account_balance(recipient_vault_id) - .await?; + let bridge_balance_after = account_balance(&ctx, bridge_account_id).await?; + let vault_balance_after = account_balance(&ctx, recipient_vault_id).await?; let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?; assert_eq!(bridge_balance_after, bridge_balance_before); @@ -421,7 +394,7 @@ async fn wait_for_vault_balance( + Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS); tokio::time::timeout(timeout, async { loop { - let balance = ctx.sequencer_client().get_account_balance(vault_id).await?; + let balance = account_balance(ctx, vault_id).await?; if balance == expected_balance { return Ok(()); } @@ -449,14 +422,8 @@ async fn bedrock_deposit_claim_and_withdraw_round_trip_succeeds() -> anyhow::Res let vault_program_id = programs::vault().id(); let recipient_vault_id = vault_core::compute_vault_account_id(vault_program_id, recipient_id); - let vault_balance_before = ctx - .sequencer_client() - .get_account_balance(recipient_vault_id) - .await?; - let recipient_balance_before = ctx - .sequencer_client() - .get_account_balance(recipient_id) - .await?; + let vault_balance_before = account_balance(&ctx, recipient_vault_id).await?; + let recipient_balance_before = account_balance(&ctx, recipient_id).await?; // Submit deposit to Bedrock submit_bedrock_deposit(ctx.bedrock_addr(), bedrock_account_pk, recipient_id, amount) @@ -507,14 +474,8 @@ async fn bedrock_deposit_claim_and_withdraw_round_trip_succeeds() -> anyhow::Res tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; let claim_on_chain = ctx.sequencer_client().get_transaction(claim_hash).await?; - let vault_balance_after_claim = ctx - .sequencer_client() - .get_account_balance(recipient_vault_id) - .await?; - let recipient_balance_after_claim = ctx - .sequencer_client() - .get_account_balance(recipient_id) - .await?; + let vault_balance_after_claim = account_balance(&ctx, recipient_vault_id).await?; + let recipient_balance_after_claim = account_balance(&ctx, recipient_id).await?; assert!( claim_on_chain.is_some(), @@ -543,7 +504,7 @@ async fn bedrock_deposit_claim_and_withdraw_round_trip_succeeds() -> anyhow::Res account_id.into(), ) .await?; - let sequencer_account = ctx.sequencer_client().get_account(account_id).await?; + let sequencer_account = get_account(&ctx, account_id).await?; assert_eq!( indexer_account, sequencer_account.into(), diff --git a/integration_tests/tests/indexer_state_consistency.rs b/integration_tests/tests/indexer_state_consistency.rs index e87927bc..4ed2fd26 100644 --- a/integration_tests/tests/indexer_state_consistency.rs +++ b/integration_tests/tests/indexer_state_consistency.rs @@ -1,51 +1,36 @@ #![expect( - clippy::shadow_unrelated, clippy::tests_outside_test_module, reason = "We don't care about these in tests" )] use std::time::Duration; -use anyhow::{Context as _, Result}; +use anyhow::Result; use indexer_service_rpc::RpcClient as _; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, - verify_commitment_is_in_state, wait_for_indexer_to_catch_up, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, + assert_private_commitment_in_state, get_account, private_mention, public_mention, send, + wait_for_indexer_to_catch_up, }; use lee::AccountId; use log::info; -use wallet::cli::{Command, programs::native_token_transfer::AuthTransferSubcommand}; #[tokio::test] async fn indexer_state_consistency() -> Result<()> { let mut ctx = TestContext::new().await?; - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(ctx.existing_public_accounts()[0]), - to: Some(public_mention(ctx.existing_public_accounts()[1])), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + let (acc0, acc1) = ( + ctx.existing_public_accounts()[0], + ctx.existing_public_accounts()[1], + ); + send(&mut ctx, public_mention(acc0), public_mention(acc1), 100).await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move"); - let acc_1_balance = sequencer_service_rpc::RpcClient::get_account_balance( - ctx.sequencer_client(), - ctx.existing_public_accounts()[0], - ) - .await?; - let acc_2_balance = sequencer_service_rpc::RpcClient::get_account_balance( - ctx.sequencer_client(), - ctx.existing_public_accounts()[1], - ) - .await?; + let acc_1_balance = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?; + let acc_2_balance = account_balance(&ctx, ctx.existing_public_accounts()[1]).await?; info!("Balance of sender: {acc_1_balance:#?}"); info!("Balance of receiver: {acc_2_balance:#?}"); @@ -56,32 +41,13 @@ async fn indexer_state_consistency() -> Result<()> { let from: AccountId = ctx.existing_private_accounts()[0]; let to: AccountId = ctx.existing_private_accounts()[1]; - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: private_mention(from), - to: Some(private_mention(to)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send(&mut ctx, private_mention(from), private_mention(to), 100).await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let new_commitment1 = ctx - .wallet() - .get_private_account_commitment(from) - .context("Failed to get private account commitment for sender")?; - assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await); - - let new_commitment2 = ctx - .wallet() - .get_private_account_commitment(to) - .context("Failed to get private account commitment for receiver")?; - assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await); + assert_private_commitment_in_state(&ctx, from, "sender").await?; + assert_private_commitment_in_state(&ctx, to, "receiver").await?; info!("Successfully transferred privately to owned account"); @@ -100,16 +66,8 @@ async fn indexer_state_consistency() -> Result<()> { .unwrap(); info!("Checking correct state transition"); - let acc1_seq_state = sequencer_service_rpc::RpcClient::get_account( - ctx.sequencer_client(), - ctx.existing_public_accounts()[0], - ) - .await?; - let acc2_seq_state = sequencer_service_rpc::RpcClient::get_account( - ctx.sequencer_client(), - ctx.existing_public_accounts()[1], - ) - .await?; + let acc1_seq_state = get_account(&ctx, ctx.existing_public_accounts()[0]).await?; + let acc2_seq_state = get_account(&ctx, ctx.existing_public_accounts()[1]).await?; assert_eq!(acc1_ind_state, acc1_seq_state.into()); assert_eq!(acc2_ind_state, acc2_seq_state.into()); diff --git a/integration_tests/tests/indexer_state_consistency_with_labels.rs b/integration_tests/tests/indexer_state_consistency_with_labels.rs index 5f561d6f..219c3ebf 100644 --- a/integration_tests/tests/indexer_state_consistency_with_labels.rs +++ b/integration_tests/tests/indexer_state_consistency_with_labels.rs @@ -9,12 +9,13 @@ use std::time::Duration; use anyhow::Result; use indexer_service_rpc::RpcClient as _; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, public_mention, wait_for_indexer_to_catch_up, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, get_account, public_mention, + send, wait_for_indexer_to_catch_up, }; use log::info; use wallet::{ account::Label, - cli::{CliAccountMention, Command, programs::native_token_transfer::AuthTransferSubcommand}, + cli::{CliAccountMention, Command}, }; #[tokio::test] @@ -38,31 +39,19 @@ async fn indexer_state_consistency_with_labels() -> Result<()> { wallet::cli::execute_subcommand(ctx.wallet_mut(), label_cmd).await?; // Send using labels instead of account IDs - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: CliAccountMention::Label(from_label), - to: Some(CliAccountMention::Label(to_label)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send( + &mut ctx, + CliAccountMention::Label(from_label), + CliAccountMention::Label(to_label), + 100, + ) + .await?; info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let acc_1_balance = sequencer_service_rpc::RpcClient::get_account_balance( - ctx.sequencer_client(), - ctx.existing_public_accounts()[0], - ) - .await?; - let acc_2_balance = sequencer_service_rpc::RpcClient::get_account_balance( - ctx.sequencer_client(), - ctx.existing_public_accounts()[1], - ) - .await?; + let acc_1_balance = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?; + let acc_2_balance = account_balance(&ctx, ctx.existing_public_accounts()[1]).await?; assert_eq!(acc_1_balance, 9900); assert_eq!(acc_2_balance, 20100); @@ -75,11 +64,7 @@ async fn indexer_state_consistency_with_labels() -> Result<()> { .get_account(ctx.existing_public_accounts()[0].into()) .await .unwrap(); - let acc1_seq_state = sequencer_service_rpc::RpcClient::get_account( - ctx.sequencer_client(), - ctx.existing_public_accounts()[0], - ) - .await?; + let acc1_seq_state = get_account(&ctx, ctx.existing_public_accounts()[0]).await?; assert_eq!(acc1_ind_state, acc1_seq_state.into()); diff --git a/integration_tests/tests/keys.rs b/integration_tests/tests/keys.rs index 9fd3b3f1..7fc73c62 100644 --- a/integration_tests/tests/keys.rs +++ b/integration_tests/tests/keys.rs @@ -8,8 +8,9 @@ use std::{str::FromStr as _, time::Duration}; use anyhow::{Context as _, Result}; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx, private_mention, - public_mention, verify_commitment_is_in_state, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, assert_public_account_restored, + fetch_privacy_preserving_tx, new_account, private_mention, public_mention, + restored_private_account, send, verify_commitment_is_in_state, }; use key_protocol::key_management::key_tree::chain_index::ChainIndex; use lee::AccountId; @@ -17,8 +18,7 @@ use log::info; use sequencer_service_rpc::RpcClient as _; use tokio::test; use wallet::cli::{ - Command, SubcommandReturnValue, - account::{AccountSubcommand, NewSubcommand}, + Command, SubcommandReturnValue, account::AccountSubcommand, programs::native_token_transfer::AuthTransferSubcommand, }; @@ -28,35 +28,12 @@ async fn sync_private_account_with_non_zero_chain_index() -> Result<()> { 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, - })); - + // Key Tree shift — create 3 accounts to advance the key index 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"); - }; + new_account(&mut ctx, true, None).await?; } - 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"); - }; + let to_account_id = new_account(&mut ctx, true, None).await?; // Get the keys for the newly created account let to_account = ctx @@ -118,107 +95,47 @@ async fn restore_keys_from_seed() -> Result<()> { let from: AccountId = ctx.existing_private_accounts()[0]; - // Create first private account at root - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: Some(ChainIndex::root()), - label: None, - })); - let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - let SubcommandReturnValue::RegisterAccount { - account_id: to_account_id1, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + // Create private accounts at root and /0 + let to_account_id1 = new_account(&mut ctx, true, Some(ChainIndex::root())).await?; + let to_account_id2 = new_account(&mut ctx, true, Some(ChainIndex::from_str("/0")?)).await?; - // Create second private account at /0 - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: Some(ChainIndex::from_str("/0")?), - label: None, - })); - let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - let SubcommandReturnValue::RegisterAccount { - account_id: to_account_id2, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; - - // Send to first private account - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: private_mention(from), - to: Some(private_mention(to_account_id1)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 100, - }); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - - // Send to second private account - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: private_mention(from), - to: Some(private_mention(to_account_id2)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 101, - }); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + // Send to both private accounts + send( + &mut ctx, + private_mention(from), + private_mention(to_account_id1), + 100, + ) + .await?; + send( + &mut ctx, + private_mention(from), + private_mention(to_account_id2), + 101, + ) + .await?; let from: AccountId = ctx.existing_public_accounts()[0]; - // Create first public account at root - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: Some(ChainIndex::root()), - label: None, - })); - let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - let SubcommandReturnValue::RegisterAccount { - account_id: to_account_id3, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + // Create public accounts at root and /0 + let to_account_id3 = new_account(&mut ctx, false, Some(ChainIndex::root())).await?; + let to_account_id4 = new_account(&mut ctx, false, Some(ChainIndex::from_str("/0")?)).await?; - // Create second public account at /0 - let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: Some(ChainIndex::from_str("/0")?), - label: None, - })); - let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - let SubcommandReturnValue::RegisterAccount { - account_id: to_account_id4, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; - - // Send to first public account - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(from), - to: Some(public_mention(to_account_id3)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 102, - }); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - - // Send to second public account - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(from), - to: Some(public_mention(to_account_id4)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 103, - }); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + // Send to both public accounts + send( + &mut ctx, + public_mention(from), + public_mention(to_account_id3), + 102, + ) + .await?; + send( + &mut ctx, + public_mention(from), + public_mention(to_account_id4), + 103, + ) + .await?; info!("Preparation complete, performing keys restoration"); @@ -226,34 +143,12 @@ async fn restore_keys_from_seed() -> Result<()> { wallet::cli::execute_keys_restoration(ctx.wallet_mut(), 10).await?; // Verify restored private accounts - let acc1 = ctx - .wallet() - .storage() - .key_chain() - .private_account(to_account_id1) - .expect("Acc 1 should be restored"); - - let acc2 = ctx - .wallet() - .storage() - .key_chain() - .private_account(to_account_id2) - .expect("Acc 2 should be restored"); + let acc1 = restored_private_account(&ctx, to_account_id1, "Acc 1"); + let acc2 = restored_private_account(&ctx, to_account_id2, "Acc 2"); // Verify restored public accounts - let _acc3 = ctx - .wallet() - .storage() - .key_chain() - .pub_account_signing_key(to_account_id3) - .expect("Acc 3 should be restored"); - - let _acc4 = ctx - .wallet() - .storage() - .key_chain() - .pub_account_signing_key(to_account_id4) - .expect("Acc 4 should be restored"); + assert_public_account_restored(&ctx, to_account_id3, "Acc 3"); + assert_public_account_restored(&ctx, to_account_id4, "Acc 4"); assert_eq!( acc1.account.program_owner, @@ -270,27 +165,20 @@ async fn restore_keys_from_seed() -> Result<()> { info!("Tree checks passed, testing restored accounts can transact"); // Test that restored accounts can send transactions - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: private_mention(to_account_id1), - to: Some(private_mention(to_account_id2)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 10, - }); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; - - let command = Command::AuthTransfer(AuthTransferSubcommand::Send { - from: public_mention(to_account_id3), - to: Some(public_mention(to_account_id4)), - to_npk: None, - to_vpk: None, - to_keys: None, - to_identifier: Some(0), - amount: 11, - }); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + send( + &mut ctx, + private_mention(to_account_id1), + private_mention(to_account_id2), + 10, + ) + .await?; + send( + &mut ctx, + public_mention(to_account_id3), + public_mention(to_account_id4), + 11, + ) + .await?; tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; diff --git a/integration_tests/tests/pinata.rs b/integration_tests/tests/pinata.rs index fa4c3d98..cea7a707 100644 --- a/integration_tests/tests/pinata.rs +++ b/integration_tests/tests/pinata.rs @@ -8,15 +8,13 @@ use std::time::Duration; use anyhow::{Context as _, Result}; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, - verify_commitment_is_in_state, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, new_account, private_mention, + public_mention, sync_private, verify_commitment_is_in_state, }; use log::info; -use sequencer_service_rpc::RpcClient as _; use tokio::test; use wallet::cli::{ Command, SubcommandReturnValue, - account::{AccountSubcommand, NewSubcommand}, programs::{ native_token_transfer::AuthTransferSubcommand, pinata::PinataProgramAgnosticSubcommand, }, @@ -26,25 +24,9 @@ use wallet::cli::{ async fn claim_pinata_to_uninitialized_public_account_fails_fast() -> Result<()> { let mut ctx = TestContext::new().await?; - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: winner_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let winner_account_id = new_account(&mut ctx, false, None).await?; - let pinata_balance_pre = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?; let claim_result = wallet::cli::execute_subcommand( ctx.wallet_mut(), @@ -64,10 +46,7 @@ async fn claim_pinata_to_uninitialized_public_account_fails_fast() -> Result<()> "Expected init guidance, got: {err}", ); - let pinata_balance_post = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?; assert_eq!(pinata_balance_post, pinata_balance_pre); @@ -78,25 +57,9 @@ async fn claim_pinata_to_uninitialized_public_account_fails_fast() -> Result<()> async fn claim_pinata_to_uninitialized_private_account_fails_fast() -> Result<()> { let mut ctx = TestContext::new().await?; - 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: winner_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let winner_account_id = new_account(&mut ctx, true, None).await?; - let pinata_balance_pre = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?; let claim_result = wallet::cli::execute_subcommand( ctx.wallet_mut(), @@ -116,10 +79,7 @@ async fn claim_pinata_to_uninitialized_private_account_fails_fast() -> Result<() "Expected init guidance, got: {err}", ); - let pinata_balance_post = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?; assert_eq!(pinata_balance_post, pinata_balance_pre); @@ -135,10 +95,7 @@ async fn claim_pinata_to_existing_public_account() -> Result<()> { to: public_mention(ctx.existing_public_accounts()[0]), }); - let pinata_balance_pre = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?; wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; @@ -146,15 +103,9 @@ async fn claim_pinata_to_existing_public_account() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Checking correct balance move"); - let pinata_balance_post = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?; - let winner_balance_post = ctx - .sequencer_client() - .get_account_balance(ctx.existing_public_accounts()[0]) - .await?; + let winner_balance_post = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?; assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize); assert_eq!(winner_balance_post, 10000 + pinata_prize); @@ -173,10 +124,7 @@ async fn claim_pinata_to_existing_private_account() -> Result<()> { to: private_mention(ctx.existing_private_accounts()[0]), }); - let pinata_balance_pre = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?; let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash: _ } = result else { @@ -187,8 +135,7 @@ async fn claim_pinata_to_existing_private_account() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; info!("Syncing private accounts"); - let command = Command::Account(AccountSubcommand::SyncPrivate {}); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + sync_private(&mut ctx).await?; let new_commitment = ctx .wallet() @@ -196,10 +143,7 @@ async fn claim_pinata_to_existing_private_account() -> Result<()> { .context("Failed to get private account commitment")?; assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); - let pinata_balance_post = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?; assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize); @@ -215,20 +159,7 @@ async fn claim_pinata_to_new_private_account() -> Result<()> { let pinata_prize = 150; // Create new private account - 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: winner_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let winner_account_id = new_account(&mut ctx, true, None).await?; // Initialize account under auth transfer program let command = Command::AuthTransfer(AuthTransferSubcommand::Init { @@ -250,10 +181,7 @@ async fn claim_pinata_to_new_private_account() -> Result<()> { to: private_mention(winner_account_id), }); - let pinata_balance_pre = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?; wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; @@ -266,10 +194,7 @@ async fn claim_pinata_to_new_private_account() -> Result<()> { .context("Failed to get private account commitment")?; assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); - let pinata_balance_post = ctx - .sequencer_client() - .get_account_balance(system_accounts::pinata_account_id()) - .await?; + let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?; assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize); diff --git a/integration_tests/tests/private_pda.rs b/integration_tests/tests/private_pda.rs index f3136717..2ccdfad8 100644 --- a/integration_tests/tests/private_pda.rs +++ b/integration_tests/tests/private_pda.rs @@ -9,7 +9,7 @@ use anyhow::{Context as _, Result}; use authenticated_transfer_core::Instruction as AuthTransferInstruction; use common::transaction::LeeTransaction; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, verify_commitment_is_in_state, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, sync_private, verify_commitment_is_in_state, }; use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder; use lee::{ @@ -22,7 +22,7 @@ use lee::{ program::Program, }; use lee_core::{ - EncryptedAccountData, InputAccountIdentity, NullifierPublicKey, + DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, NullifierPublicKey, account::{Account, AccountWithMetadata}, encryption::ViewingPublicKey, program::PdaSeed, @@ -30,10 +30,7 @@ use lee_core::{ use log::info; use sequencer_service_rpc::RpcClient as _; use tokio::test; -use wallet::{ - AccountIdentity, WalletCore, - cli::{Command, account::AccountSubcommand}, -}; +use wallet::{AccountIdentity, WalletCore}; /// Funds a private PDA by calling `auth_transfer` directly. #[expect( @@ -78,6 +75,7 @@ async fn fund_private_pda( npk, ssk, identifier, + commitment_root: DUMMY_COMMITMENT_HASH, seed: Some((seed, authority_program_id)), }, ]; @@ -217,11 +215,7 @@ async fn private_pda_family_members_receive_and_spend() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Sync so alice's wallet discovers and stores both PDAs. - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::SyncPrivate {}), - ) - .await?; + sync_private(&mut ctx).await?; // Both PDAs must be discoverable and have the correct balance. let pda_0_account = ctx @@ -300,11 +294,7 @@ async fn private_pda_family_members_receive_and_spend() -> Result<()> { info!("Waiting for block"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::SyncPrivate {}), - ) - .await?; + sync_private(&mut ctx).await?; // After spending, PDAs should have the remaining balance. let pda_0_spent = ctx diff --git a/integration_tests/tests/program_deployment.rs b/integration_tests/tests/program_deployment.rs index ec01c3c8..2b242549 100644 --- a/integration_tests/tests/program_deployment.rs +++ b/integration_tests/tests/program_deployment.rs @@ -7,14 +7,11 @@ use std::{io::Write as _, time::Duration}; use anyhow::Result; use common::transaction::LeeTransaction; -use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext}; +use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, get_account, new_account}; use log::info; use sequencer_service_rpc::RpcClient as _; use tokio::test; -use wallet::cli::{ - Command, SubcommandReturnValue, - account::{AccountSubcommand, NewSubcommand}, -}; +use wallet::cli::Command; #[test] async fn deploy_and_execute_program() -> Result<()> { @@ -35,17 +32,7 @@ async fn deploy_and_execute_program() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let SubcommandReturnValue::RegisterAccount { account_id } = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await? - else { - panic!("Expected RegisterAccount return value"); - }; + let account_id = new_account(&mut ctx, false, None).await?; let nonces = ctx.wallet().get_accounts_nonces(vec![account_id]).await?; let private_key = ctx @@ -66,7 +53,7 @@ async fn deploy_and_execute_program() -> Result<()> { // block tokio::time::sleep(Duration::from_secs(2 * TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let post_state_account = ctx.sequencer_client().get_account(account_id).await?; + let post_state_account = get_account(&ctx, account_id).await?; let expected_data: &[u8] = &[]; assert_eq!(post_state_account.program_owner, claimer.id()); diff --git a/integration_tests/tests/shared_accounts.rs b/integration_tests/tests/shared_accounts.rs index 39bdd36c..cc6e6e1d 100644 --- a/integration_tests/tests/shared_accounts.rs +++ b/integration_tests/tests/shared_accounts.rs @@ -19,7 +19,7 @@ use std::time::Duration; use anyhow::{Context as _, Result}; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, sync_private, }; use log::info; use tokio::test; @@ -197,8 +197,7 @@ async fn fund_shared_account_from_public() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Sync private accounts - let command = Command::Account(AccountSubcommand::SyncPrivate); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + sync_private(&mut ctx).await?; // Fund from a public account let from_public = ctx.existing_public_accounts()[0]; @@ -216,8 +215,7 @@ async fn fund_shared_account_from_public() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Sync private accounts - let command = Command::Account(AccountSubcommand::SyncPrivate); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + sync_private(&mut ctx).await?; // Verify the shared account was updated let entry = ctx diff --git a/integration_tests/tests/token.rs b/integration_tests/tests/token.rs index 60bd3de8..d7578a22 100644 --- a/integration_tests/tests/token.rs +++ b/integration_tests/tests/token.rs @@ -8,12 +8,11 @@ use std::time::Duration; use anyhow::{Context as _, Result}; use integration_tests::{ - TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, - verify_commitment_is_in_state, + TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, get_account, new_account, private_mention, + public_mention, sync_private, verify_commitment_is_in_state, }; use key_protocol::key_management::key_tree::chain_index::ChainIndex; use log::info; -use sequencer_service_rpc::RpcClient as _; use token_core::{TokenDefinition, TokenHolding}; use tokio::test; use wallet::{ @@ -30,52 +29,13 @@ async fn create_and_transfer_public_token() -> Result<()> { let mut ctx = TestContext::new().await?; // Create new account for the token definition - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, false, None).await?; // Create new account for the token supply holder - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, false, None).await?; // Create new account for receiving a token transaction - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, false, None).await?; // Create new token let name = "A NAME".to_owned(); @@ -92,10 +52,7 @@ async fn create_and_transfer_public_token() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Check the status of the token definition account - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!(definition_acc.program_owner, programs::token().id()); @@ -109,10 +66,7 @@ async fn create_and_transfer_public_token() -> Result<()> { ); // Check the status of the token holding account with the total supply - let supply_acc = ctx - .sequencer_client() - .get_account(supply_account_id) - .await?; + let supply_acc = get_account(&ctx, supply_account_id).await?; // The account must be owned by the token program assert_eq!(supply_acc.program_owner, programs::token().id()); @@ -143,10 +97,7 @@ async fn create_and_transfer_public_token() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Check the status of the supply account after transfer - let supply_acc = ctx - .sequencer_client() - .get_account(supply_account_id) - .await?; + let supply_acc = get_account(&ctx, supply_account_id).await?; assert_eq!(supply_acc.program_owner, programs::token().id()); let token_holding = TokenHolding::try_from(&supply_acc.data)?; assert_eq!( @@ -158,10 +109,7 @@ async fn create_and_transfer_public_token() -> Result<()> { ); // Check the status of the recipient account after transfer - let recipient_acc = ctx - .sequencer_client() - .get_account(recipient_account_id) - .await?; + let recipient_acc = get_account(&ctx, recipient_account_id).await?; assert_eq!(recipient_acc.program_owner, programs::token().id()); let token_holding = TokenHolding::try_from(&recipient_acc.data)?; assert_eq!( @@ -186,10 +134,7 @@ async fn create_and_transfer_public_token() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Check the status of the token definition account after burn - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!( @@ -202,10 +147,7 @@ async fn create_and_transfer_public_token() -> Result<()> { ); // Check the status of the recipient account after burn - let recipient_acc = ctx - .sequencer_client() - .get_account(recipient_account_id) - .await?; + let recipient_acc = get_account(&ctx, recipient_account_id).await?; let token_holding = TokenHolding::try_from(&recipient_acc.data)?; assert_eq!( @@ -234,10 +176,7 @@ async fn create_and_transfer_public_token() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Check the status of the token definition account after mint - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!( @@ -250,10 +189,7 @@ async fn create_and_transfer_public_token() -> Result<()> { ); // Check the status of the recipient account after mint - let recipient_acc = ctx - .sequencer_client() - .get_account(recipient_account_id) - .await?; + let recipient_acc = get_account(&ctx, recipient_account_id).await?; let token_holding = TokenHolding::try_from(&recipient_acc.data)?; assert_eq!( @@ -274,52 +210,13 @@ async fn create_and_transfer_token_with_private_supply() -> Result<()> { let mut ctx = TestContext::new().await?; // Create new account for the token definition (public) - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, false, None).await?; // Create new account for the token supply holder (private) - 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: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, true, None).await?; // Create new account for receiving a token transaction (private) - 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: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, true, None).await?; // Create new token let name = "A NAME".to_owned(); @@ -337,10 +234,7 @@ async fn create_and_transfer_token_with_private_supply() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Check the status of the token definition account - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!(definition_acc.program_owner, programs::token().id()); @@ -402,10 +296,7 @@ async fn create_and_transfer_token_with_private_supply() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Check the token definition account after burn - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!( @@ -448,36 +339,10 @@ async fn create_token_with_private_definition() -> Result<()> { let mut ctx = TestContext::new().await?; // Create token definition account (private) - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Private { - cci: Some(ChainIndex::root()), - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, true, Some(ChainIndex::root())).await?; // Create supply account (public) - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: Some(ChainIndex::root()), - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, false, Some(ChainIndex::root())).await?; // Create token with private definition let name = "A NAME".to_owned(); @@ -502,10 +367,7 @@ async fn create_token_with_private_definition() -> Result<()> { assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await); // Verify supply account - let supply_acc = ctx - .sequencer_client() - .get_account(supply_account_id) - .await?; + let supply_acc = get_account(&ctx, supply_account_id).await?; assert_eq!(supply_acc.program_owner, programs::token().id()); let token_holding = TokenHolding::try_from(&supply_acc.data)?; @@ -518,36 +380,10 @@ async fn create_token_with_private_definition() -> Result<()> { ); // Create private recipient account - 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: recipient_account_id_private, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id_private = new_account(&mut ctx, true, None).await?; // Create public recipient account - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: recipient_account_id_public, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id_public = new_account(&mut ctx, false, None).await?; // Mint to public account let mint_amount_public = 10; @@ -583,10 +419,7 @@ async fn create_token_with_private_definition() -> Result<()> { ); // Verify public recipient received tokens - let recipient_acc = ctx - .sequencer_client() - .get_account(recipient_account_id_public) - .await?; + let recipient_acc = get_account(&ctx, recipient_account_id_public).await?; let token_holding = TokenHolding::try_from(&recipient_acc.data)?; assert_eq!( @@ -646,36 +479,10 @@ async fn create_token_with_private_definition_and_supply() -> Result<()> { let mut ctx = TestContext::new().await?; // Create token definition account (private) - 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: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, true, None).await?; // Create supply account (private) - 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: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, true, None).await?; // Create token with both private definition and supply let name = "A NAME".to_owned(); @@ -722,20 +529,7 @@ async fn create_token_with_private_definition_and_supply() -> Result<()> { ); // Create recipient account - 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: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, true, None).await?; // Transfer tokens let transfer_amount = 7; @@ -804,52 +598,13 @@ async fn shielded_token_transfer() -> Result<()> { let mut ctx = TestContext::new().await?; // Create token definition account (public) - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, false, None).await?; // Create supply account (public) - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, false, None).await?; // Create recipient account (private) for shielded transfer - 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: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, true, None).await?; // Create token let name = "A NAME".to_owned(); @@ -884,10 +639,7 @@ async fn shielded_token_transfer() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Verify supply account balance - let supply_acc = ctx - .sequencer_client() - .get_account(supply_account_id) - .await?; + let supply_acc = get_account(&ctx, supply_account_id).await?; let token_holding = TokenHolding::try_from(&supply_acc.data)?; assert_eq!( token_holding, @@ -928,52 +680,13 @@ async fn deshielded_token_transfer() -> Result<()> { let mut ctx = TestContext::new().await?; // Create token definition account (public) - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, false, None).await?; // Create supply account (private) - 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: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, true, None).await?; // Create recipient account (public) for deshielded transfer - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, false, None).await?; // Create token with private supply let name = "A NAME".to_owned(); @@ -1029,10 +742,7 @@ async fn deshielded_token_transfer() -> Result<()> { ); // Verify recipient balance - let recipient_acc = ctx - .sequencer_client() - .get_account(recipient_account_id) - .await?; + let recipient_acc = get_account(&ctx, recipient_account_id).await?; let token_holding = TokenHolding::try_from(&recipient_acc.data)?; assert_eq!( token_holding, @@ -1052,36 +762,10 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> { let mut ctx = TestContext::new().await?; // Create token definition account (private) - 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: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, true, None).await?; // Create supply account (private) - 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: supply_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let supply_account_id = new_account(&mut ctx, true, None).await?; // Create token let name = "A NAME".to_owned(); @@ -1099,20 +783,7 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Create new private account for claiming path - 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: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, true, None).await?; // Get keys for foreign mint (claiming path) let holder = ctx @@ -1143,8 +814,7 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> { tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; // Sync to claim the account - let command = Command::Account(AccountSubcommand::SyncPrivate {}); - wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?; + sync_private(&mut ctx).await?; // Verify commitment exists let recipient_commitment = ctx @@ -1224,10 +894,7 @@ async fn create_token_using_labels() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let definition_acc = ctx - .sequencer_client() - .get_account(definition_account_id) - .await?; + let definition_acc = get_account(&ctx, definition_account_id).await?; let token_definition = TokenDefinition::try_from(&definition_acc.data)?; assert_eq!(definition_acc.program_owner, programs::token().id()); @@ -1240,10 +907,7 @@ async fn create_token_using_labels() -> Result<()> { } ); - let supply_acc = ctx - .sequencer_client() - .get_account(supply_account_id) - .await?; + let supply_acc = get_account(&ctx, supply_account_id).await?; let token_holding = TokenHolding::try_from(&supply_acc.data)?; assert_eq!( token_holding, @@ -1263,20 +927,7 @@ async fn transfer_token_using_from_label() -> Result<()> { let mut ctx = TestContext::new().await?; // Create definition account - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: definition_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let definition_account_id = new_account(&mut ctx, false, None).await?; // Create supply account with a label let supply_label = Label::new("token-supply-sender"); @@ -1296,20 +947,7 @@ async fn transfer_token_using_from_label() -> Result<()> { }; // Create recipient account - let result = wallet::cli::execute_subcommand( - ctx.wallet_mut(), - Command::Account(AccountSubcommand::New(NewSubcommand::Public { - cci: None, - label: None, - })), - ) - .await?; - let SubcommandReturnValue::RegisterAccount { - account_id: recipient_account_id, - } = result - else { - anyhow::bail!("Expected RegisterAccount return value"); - }; + let recipient_account_id = new_account(&mut ctx, false, None).await?; // Create token let total_supply = 50; @@ -1340,10 +978,7 @@ async fn transfer_token_using_from_label() -> Result<()> { info!("Waiting for next block creation"); tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await; - let recipient_acc = ctx - .sequencer_client() - .get_account(recipient_account_id) - .await?; + let recipient_acc = get_account(&ctx, recipient_account_id).await?; let token_holding = TokenHolding::try_from(&recipient_acc.data)?; assert_eq!( token_holding, diff --git a/integration_tests/tests/tps.rs b/integration_tests/tests/tps.rs index a11668a8..0b9a1002 100644 --- a/integration_tests/tests/tps.rs +++ b/integration_tests/tests/tps.rs @@ -23,7 +23,8 @@ use lee::{ public_transaction as putx, }; use lee_core::{ - EncryptedAccountData, InputAccountIdentity, MembershipProof, NullifierPublicKey, + DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, MembershipProof, + NullifierPublicKey, account::{AccountWithMetadata, Nonce, data::Data}, encryption::ViewingPublicKey, }; @@ -314,6 +315,7 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction { npk: recipient_npk, ssk: recipient_ss, identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, }, ], &program.into(), diff --git a/integration_tests/tests/wallet_ffi.rs b/integration_tests/tests/wallet_ffi.rs index 0ef53592..57f1afa4 100644 --- a/integration_tests/tests/wallet_ffi.rs +++ b/integration_tests/tests/wallet_ffi.rs @@ -16,6 +16,7 @@ use std::{ ffi::{CStr, CString, c_char}, io::Write as _, path::Path, + str::FromStr as _, time::Duration, }; @@ -30,9 +31,11 @@ use log::info; use tempfile::tempdir; use wallet::{account::HumanReadableAccount, program_facades::vault::Vault}; use wallet_ffi::{ - FfiAccount, FfiAccountIdentity, FfiAccountList, FfiBytes32, FfiPrivateAccountKeys, - FfiProgramId, FfiPublicAccountKey, FfiTransferResult, FfiU128, WalletHandle, error, + FfiAccount, FfiAccountIdWithPrivacy, FfiAccountIdentity, FfiAccountList, FfiBytes32, + FfiPrivateAccountKeys, FfiProgramId, FfiPublicAccountKey, FfiTransferResult, FfiU128, + WalletHandle, error, generic_transaction::{FfiProgramWithDependencies, FfiTransactionResult}, + label::{AccountIdResolvedFromLabel, LabelAvailability, LabelList}, wallet::FfiCreateWalletOutput, }; @@ -257,6 +260,29 @@ unsafe extern "C" { fn wallet_ffi_free_transaction_result(result: *mut FfiTransactionResult); fn wallet_ffi_free_account_identity(account_identity: *mut FfiAccountIdentity); + + fn wallet_ffi_check_label_available( + handle: *mut WalletHandle, + label: *const c_char, + ) -> LabelAvailability; + + fn wallet_ffi_add_label( + handle: *mut WalletHandle, + label: *const c_char, + account_id_with_privacy: FfiAccountIdWithPrivacy, + ) -> error::WalletFfiError; + + fn wallet_ffi_resolve_label( + handle: *mut WalletHandle, + label: *const c_char, + ) -> AccountIdResolvedFromLabel; + + fn wallet_ffi_get_all_labels_for_account( + handle: *mut WalletHandle, + account_id_with_privacy: FfiAccountIdWithPrivacy, + ) -> LabelList; + + fn wallet_ffi_free_label_list(label_list: *mut LabelList) -> error::WalletFfiError; } fn new_wallet_ffi_with_test_context_config( @@ -1926,3 +1952,125 @@ fn test_wallet_ffi_vault_balance_and_claim_private() -> Result<()> { Ok(()) } + +#[test] +fn test_wallet_ffi_single_label() -> Result<()> { + let ctx = BlockingTestContext::new()?; + let home = tempfile::tempdir()?; + let FfiCreateWalletOutput { + wallet: wallet_ffi_handle, + mnemonic: _, + } = new_wallet_ffi_with_test_context_config(&ctx, home.path())?; + + let mut out_account_id_1 = FfiBytes32::from_bytes([0; 32]); + unsafe { + wallet_ffi_create_account_public(wallet_ffi_handle, &raw mut out_account_id_1).unwrap(); + } + + info!("Waiting for next block creation"); + std::thread::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)); + + let lab_1 = CString::from_str("LABEL1").unwrap().into_raw(); + + let lab_1_availability = unsafe { wallet_ffi_check_label_available(wallet_ffi_handle, lab_1) }; + + assert_eq!(lab_1_availability.error, error::WalletFfiError::Success); + assert!(lab_1_availability.is_available); + + let acc_1_id_with_privacy = FfiAccountIdWithPrivacy { + account_id: out_account_id_1, + is_private: false, + }; + + let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_1, acc_1_id_with_privacy) }; + + assert_eq!(err, error::WalletFfiError::Success); + + let lab_1_availability = unsafe { wallet_ffi_check_label_available(wallet_ffi_handle, lab_1) }; + + assert!(!lab_1_availability.is_available); + + let acc_resolved = unsafe { wallet_ffi_resolve_label(wallet_ffi_handle, lab_1) }; + + assert_eq!(acc_resolved.account_id, acc_1_id_with_privacy); + + unsafe { + wallet_ffi_free_string(lab_1); + wallet_ffi_destroy(wallet_ffi_handle); + } + + Ok(()) +} + +#[test] +fn test_wallet_ffi_more_labels() -> Result<()> { + let ctx = BlockingTestContext::new()?; + let home = tempfile::tempdir()?; + let FfiCreateWalletOutput { + wallet: wallet_ffi_handle, + mnemonic: _, + } = new_wallet_ffi_with_test_context_config(&ctx, home.path())?; + + let mut out_account_id_1 = FfiBytes32::from_bytes([0; 32]); + unsafe { + wallet_ffi_create_account_public(wallet_ffi_handle, &raw mut out_account_id_1).unwrap(); + } + + info!("Waiting for next block creation"); + std::thread::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)); + + let lab_1 = CString::from_str("LABEL1").unwrap().into_raw(); + let lab_2 = CString::from_str("LABEL2").unwrap().into_raw(); + let lab_3 = CString::from_str("LABEL3").unwrap().into_raw(); + + let acc_1_id_with_privacy = FfiAccountIdWithPrivacy { + account_id: out_account_id_1, + is_private: false, + }; + + let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_1, acc_1_id_with_privacy) }; + + assert_eq!(err, error::WalletFfiError::Success); + + let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_2, acc_1_id_with_privacy) }; + + assert_eq!(err, error::WalletFfiError::Success); + + let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_3, acc_1_id_with_privacy) }; + + assert_eq!(err, error::WalletFfiError::Success); + + let mut label_list_for_out_acc = + unsafe { wallet_ffi_get_all_labels_for_account(wallet_ffi_handle, acc_1_id_with_privacy) }; + + assert_eq!(label_list_for_out_acc.error, error::WalletFfiError::Success); + assert_eq!(label_list_for_out_acc.labels_size, 3); + + let lab_ref_1 = unsafe { &*label_list_for_out_acc.labels_data.add(0) }; + let lab_ref_c_str_1 = unsafe { CStr::from_ptr(*lab_ref_1) }; + + assert_eq!(lab_ref_c_str_1.to_str().unwrap(), "LABEL1"); + + let lab_ref_2 = unsafe { &*label_list_for_out_acc.labels_data.add(1) }; + let lab_ref_c_str_2 = unsafe { CStr::from_ptr(*lab_ref_2) }; + + assert_eq!(lab_ref_c_str_2.to_str().unwrap(), "LABEL2"); + + let lab_ref_3 = unsafe { &*label_list_for_out_acc.labels_data.add(2) }; + let lab_ref_c_str_3 = unsafe { CStr::from_ptr(*lab_ref_3) }; + + assert_eq!(lab_ref_c_str_3.to_str().unwrap(), "LABEL3"); + + let err = unsafe { wallet_ffi_free_label_list(&raw mut label_list_for_out_acc) }; + + assert_eq!(err, error::WalletFfiError::Success); + + unsafe { + wallet_ffi_free_string(lab_1); + wallet_ffi_free_string(lab_2); + wallet_ffi_free_string(lab_3); + wallet_ffi_destroy(wallet_ffi_handle); + } + + Ok(()) +} diff --git a/lee/key_protocol/src/key_management/ephemeral_key_holder.rs b/lee/key_protocol/src/key_management/ephemeral_key_holder.rs index a53ae47c..9bc81391 100644 --- a/lee/key_protocol/src/key_management/ephemeral_key_holder.rs +++ b/lee/key_protocol/src/key_management/ephemeral_key_holder.rs @@ -48,14 +48,3 @@ impl EphemeralKeyHolder { self.shared_secret } } - -/// Encapsulates a fresh shared secret toward `vpk` and returns `(shared_secret, ciphertext)`. -/// -/// Used when the local side is acting as an "ephemeral receiver" — i.e. generating a -/// one-sided encryption that only the holder of the VSK can decrypt. -#[must_use] -pub fn produce_one_sided_shared_secret_receiver( - vpk: &ViewingPublicKey, -) -> (SharedSecretKey, EphemeralPublicKey) { - SharedSecretKey::encapsulate(vpk) -} diff --git a/lee/key_protocol/src/key_management/group_key_holder.rs b/lee/key_protocol/src/key_management/group_key_holder.rs index b18960cf..1efc49cb 100644 --- a/lee/key_protocol/src/key_management/group_key_holder.rs +++ b/lee/key_protocol/src/key_management/group_key_holder.rs @@ -1,7 +1,7 @@ use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _}; use lee_core::{ Identifier, SharedSecretKey, - encryption::{EphemeralPublicKey, ViewingPublicKey}, + encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey}, program::{PdaSeed, ProgramId}, }; use rand::{RngCore as _, rngs::OsRng}; @@ -167,7 +167,7 @@ impl GroupKeyHolder { /// /// Uses ML-KEM-768 encapsulation to derive a shared secret, then AES-256-GCM to encrypt /// the payload. The returned bytes are - /// `kem_ciphertext (1088) || nonce (12) || ciphertext+tag (48)` = 1148 bytes. + /// `kem_ciphertext (ML_KEM_768_CIPHERTEXT_LEN) || nonce (12) || ciphertext+tag (48)`. /// /// Each call generates a fresh KEM encapsulation, so two seals of the same holder produce /// different ciphertexts. @@ -187,7 +187,7 @@ impl GroupKeyHolder { .encrypt(&nonce, self.gms.as_ref()) .expect("AES-GCM encryption should not fail with valid key/nonce"); - let capacity = 1088_usize + let capacity = ML_KEM_768_CIPHERTEXT_LEN .checked_add(12) .and_then(|n| n.checked_add(ciphertext.len())) .expect("seal capacity overflow"); @@ -203,21 +203,21 @@ impl GroupKeyHolder { /// Returns `Err` if the ciphertext is too short or the AES-GCM authentication tag /// doesn't verify (wrong key or tampered data). pub fn unseal(sealed: &[u8], own_key: &SealingSecretKey) -> Result { - // kem_ciphertext (1088) + nonce (12) = header, then AES-GCM tag (16) minimum. - const KEM_CT_LEN: usize = 1088; - const HEADER_LEN: usize = KEM_CT_LEN + 12; + // kem_ciphertext (ML_KEM_768_CIPHERTEXT_LEN) + nonce (12) = header, then AES-GCM tag (16) + // minimum. + const HEADER_LEN: usize = ML_KEM_768_CIPHERTEXT_LEN + 12; const MIN_LEN: usize = HEADER_LEN + 16; if sealed.len() < MIN_LEN { return Err(SealError::TooShort); } - let kem_ct = EphemeralPublicKey(sealed[..KEM_CT_LEN].to_vec()); - let nonce = aes_gcm::Nonce::from_slice(&sealed[KEM_CT_LEN..HEADER_LEN]); + let kem_ct = EphemeralPublicKey(sealed[..ML_KEM_768_CIPHERTEXT_LEN].to_vec()); + let nonce = aes_gcm::Nonce::from_slice(&sealed[ML_KEM_768_CIPHERTEXT_LEN..HEADER_LEN]); let ciphertext = &sealed[HEADER_LEN..]; let shared = SharedSecretKey::decapsulate(&kem_ct, &own_key.d, &own_key.z) - .expect("key_protocol::group_key_holder::GroupKeyHolder::unseal: KEM_CT_LEN guarantees exactly 1088 bytes"); + .expect("key_protocol::group_key_holder::GroupKeyHolder::unseal: ML_KEM_768_CIPHERTEXT_LEN guarantees exactly 1088 bytes"); let aes_key = Self::seal_kdf(&shared); let cipher = Aes256Gcm::new(&aes_key.into()); diff --git a/lee/key_protocol/src/key_management/key_tree/chain_index.rs b/lee/key_protocol/src/key_management/key_tree/chain_index.rs index b22dc779..6ea2e8a1 100644 --- a/lee/key_protocol/src/key_management/key_tree/chain_index.rs +++ b/lee/key_protocol/src/key_management/key_tree/chain_index.rs @@ -139,7 +139,7 @@ impl ChainIndex { .map(|item| Self(item.into_iter().copied().collect())) } - pub fn chain_ids_at_depth(depth: usize) -> impl Iterator { + fn collect_chain_ids_at_depth(depth: usize) -> Vec { let mut stack = vec![Self(vec![0; depth])]; let mut cumulative_stack = vec![Self(vec![0; depth])]; @@ -152,23 +152,18 @@ impl ChainIndex { } } - cumulative_stack.into_iter().unique() + cumulative_stack + } + + pub fn chain_ids_at_depth(depth: usize) -> impl Iterator { + Self::collect_chain_ids_at_depth(depth).into_iter().unique() } pub fn chain_ids_at_depth_rev(depth: usize) -> impl Iterator { - let mut stack = vec![Self(vec![0; depth])]; - let mut cumulative_stack = vec![Self(vec![0; depth])]; - - while let Some(top_id) = stack.pop() { - if let Some(collapsed_id) = top_id.collapse_back() { - for id in collapsed_id.shuffle_iter() { - stack.push(id.clone()); - cumulative_stack.push(id); - } - } - } - - cumulative_stack.into_iter().rev().unique() + Self::collect_chain_ids_at_depth(depth) + .into_iter() + .rev() + .unique() } } diff --git a/lee/key_protocol/src/key_management/key_tree/keys_private.rs b/lee/key_protocol/src/key_management/key_tree/keys_private.rs index 5a27be79..48b47975 100644 --- a/lee/key_protocol/src/key_management/key_tree/keys_private.rs +++ b/lee/key_protocol/src/key_management/key_tree/keys_private.rs @@ -6,7 +6,7 @@ use sha2::Digest as _; use crate::key_management::{ KeyChain, - key_tree::traits::KeyTreeNode, + key_tree::{split_hash, traits::KeyTreeNode}, secret_holders::{PrivateKeyHolder, SecretSpendingKey}, }; @@ -23,38 +23,11 @@ impl ChildKeysPrivate { #[must_use] pub fn root(seed: [u8; 64]) -> Self { let hash_value = hmac_sha512::HMAC::mac(seed, b"LEE_master_priv"); + let (first, ccc) = split_hash(&hash_value); - let ssk = SecretSpendingKey( - *hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - ); - let ccc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); + let ssk = SecretSpendingKey(first); - let nsk = ssk.generate_nullifier_secret_key(None); - let vsk = ssk.generate_viewing_secret_seed_key(None); - - let npk = NullifierPublicKey::from(&nsk); - let vpk = ViewingPublicKey::from(&vsk); - - Self { - value: ( - KeyChain { - secret_spending_key: ssk, - nullifier_public_key: npk, - viewing_public_key: vpk, - private_key_holder: PrivateKeyHolder { - nullifier_secret_key: nsk, - viewing_secret_key: vsk, - }, - }, - BTreeMap::from_iter([(PrivateAccountKind::Regular(0), lee::Account::default())]), - ), - ccc, - cci: None, - } + Self::from_ssk_and_ccc(ssk, ccc, None) } #[must_use] @@ -77,18 +50,16 @@ impl ChildKeysPrivate { input.extend_from_slice(&cci.to_be_bytes()); let hash_value = hmac_sha512::HMAC::mac(input, self.ccc); + let (first, ccc) = split_hash(&hash_value); - let ssk = SecretSpendingKey( - *hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - ); - let ccc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); + let ssk = SecretSpendingKey(first); - let nsk = ssk.generate_nullifier_secret_key(Some(cci)); - let vsk = ssk.generate_viewing_secret_seed_key(Some(cci)); + Self::from_ssk_and_ccc(ssk, ccc, Some(cci)) + } + + fn from_ssk_and_ccc(ssk: SecretSpendingKey, ccc: [u8; 32], cci: Option) -> Self { + let nsk = ssk.generate_nullifier_secret_key(cci); + let vsk = ssk.generate_viewing_secret_seed_key(cci); let npk = NullifierPublicKey::from(&nsk); let vpk = ViewingPublicKey::from(&vsk); @@ -107,7 +78,7 @@ impl ChildKeysPrivate { BTreeMap::from_iter([(PrivateAccountKind::Regular(0), lee::Account::default())]), ), ccc, - cci: Some(cci), + cci, } } } @@ -137,16 +108,16 @@ mod tests { use super::*; use crate::key_management::{self, secret_holders::ViewingSecretKey}; + const SEED: [u8; 64] = [ + 252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, 49, + 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, 129, 223, + 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, 114, 39, 38, + 118, 197, 205, 225, + ]; + #[test] fn master_key_generation() { - let seed: [u8; 64] = [ - 252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, - 49, 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, - 129, 223, 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, - 114, 39, 38, 118, 197, 205, 225, - ]; - - let keys = ChildKeysPrivate::root(seed); + let keys = ChildKeysPrivate::root(SEED); let expected_ssk = key_management::secret_holders::SecretSpendingKey([ 246, 79, 26, 124, 135, 95, 52, 51, 201, 27, 48, 194, 2, 144, 51, 219, 245, 128, 139, @@ -179,6 +150,7 @@ mod tests { ], ); + // Length matches MlKem768EncapsulationKey::LEN. let expected_vpk: [u8; 1184] = [ 127, 229, 162, 212, 104, 117, 4, 150, 192, 103, 122, 195, 14, 35, 12, 60, 52, 23, 220, 150, 100, 203, 34, 34, 127, 232, 156, 43, 218, 109, 6, 160, 67, 35, 210, 194, 25, 181, @@ -253,14 +225,7 @@ mod tests { #[test] fn child_keys_generation() { - let seed: [u8; 64] = [ - 252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, - 49, 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, - 129, 223, 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, - 114, 39, 38, 118, 197, 205, 225, - ]; - - let root_node = ChildKeysPrivate::root(seed); + let root_node = ChildKeysPrivate::root(SEED); let child_node = ChildKeysPrivate::nth_child(&root_node, 42_u32); let expected_ssk = key_management::secret_holders::SecretSpendingKey([ @@ -293,6 +258,7 @@ mod tests { ], ); + // Length matches MlKem768EncapsulationKey::LEN. let expected_vpk: [u8; 1184] = [ 215, 229, 207, 120, 148, 177, 148, 197, 72, 222, 134, 3, 231, 146, 123, 226, 36, 84, 232, 179, 205, 16, 241, 142, 9, 81, 58, 54, 12, 115, 148, 182, 19, 245, 22, 203, 57, diff --git a/lee/key_protocol/src/key_management/key_tree/keys_public.rs b/lee/key_protocol/src/key_management/key_tree/keys_public.rs index 947fb83c..4caad0e7 100644 --- a/lee/key_protocol/src/key_management/key_tree/keys_public.rs +++ b/lee/key_protocol/src/key_management/key_tree/keys_public.rs @@ -1,7 +1,7 @@ use k256::elliptic_curve::PrimeField as _; use serde::{Deserialize, Serialize}; -use crate::key_management::key_tree::traits::KeyTreeNode; +use crate::key_management::key_tree::{split_hash, traits::KeyTreeNode}; #[derive(Debug, Serialize, Deserialize, Clone)] #[cfg_attr(any(test, feature = "test_utils"), derive(PartialEq, Eq))] @@ -21,52 +21,32 @@ impl ChildKeysPublic { #[must_use] pub fn root(seed: [u8; 64]) -> Self { let hash_value = hmac_sha512::HMAC::mac(seed, "LEE_master_pub"); + let (first, cc) = split_hash(&hash_value); - let sk = lee::PrivateKey::try_new( - *hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - ) - .expect("Expect a valid Private Key"); - let ssk = lee::PrivateKey::tweak(sk.value()).expect("`key_protocol::key_management::keys_public::root()`: Invalid private key produced from `tweak`"); + let sk = lee::PrivateKey::try_new(first).expect("Expect a valid Private Key"); - let cc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); - let pk = lee::PublicKey::new_from_private_key(&ssk); - - Self { - sk, - ssk, - pk, - cc, - cci: None, - } + Self::from_sk_and_cc(sk, cc, None) } #[must_use] pub fn nth_child(&self, cci: u32) -> Self { let hash_value = self.compute_hash_value(cci); + let (first, cc) = split_hash(&hash_value); - let lhs = k256::Scalar::from_repr( - (*hash_value - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32")) - .into(), - ) - .expect("Expect a valid k256 scalar"); + let lhs = k256::Scalar::from_repr(first.into()).expect("Expect a valid k256 scalar"); let rhs = k256::Scalar::from_repr((*self.sk.value()).into()).expect("Expect a valid k256 scalar"); let sk = lee::PrivateKey::try_new(lhs.add(&rhs).to_bytes().into()) .expect("Expect a valid private key"); - let ssk = lee::PrivateKey::tweak(sk.value()).expect("`key_protocol::key_management::keys_public::nth_child()`: Invalid private key produced from `tweak`"); - - let cc = *hash_value - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"); + Self::from_sk_and_cc(sk, cc, Some(cci)) + } + fn from_sk_and_cc(sk: lee::PrivateKey, cc: [u8; 32], cci: Option) -> Self { + let ssk = lee::PrivateKey::tweak(sk.value()).expect( + "`key_protocol::key_management::keys_public::ChildKeysPublic`: Invalid private key produced from `tweak`", + ); let pk = lee::PublicKey::new_from_private_key(&ssk); Self { @@ -74,7 +54,7 @@ impl ChildKeysPublic { ssk, pk, cc, - cci: Some(cci), + cci, } } @@ -128,15 +108,16 @@ mod tests { use super::*; + const SEED: [u8; 64] = [ + 88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173, + 134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, 22, + 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, 187, + 148, 92, 44, 253, 210, 37, + ]; + #[test] fn master_keys_generation() { - let seed = [ - 88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173, - 134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, - 22, 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, - 187, 148, 92, 44, 253, 210, 37, - ]; - let keys = ChildKeysPublic::root(seed); + let keys = ChildKeysPublic::root(SEED); let expected_cc = [ 238, 94, 84, 154, 56, 224, 80, 218, 133, 249, 179, 222, 9, 24, 17, 252, 120, 127, 222, @@ -169,13 +150,7 @@ mod tests { #[test] fn child_keys_generation() { - let seed = [ - 88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173, - 134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, - 22, 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, - 187, 148, 92, 44, 253, 210, 37, - ]; - let root_keys = ChildKeysPublic::root(seed); + let root_keys = ChildKeysPublic::root(SEED); let cci = (2_u32).pow(31) + 13; let child_keys = ChildKeysPublic::nth_child(&root_keys, cci); diff --git a/lee/key_protocol/src/key_management/key_tree/mod.rs b/lee/key_protocol/src/key_management/key_tree/mod.rs index c15c09a5..7e4b432c 100644 --- a/lee/key_protocol/src/key_management/key_tree/mod.rs +++ b/lee/key_protocol/src/key_management/key_tree/mod.rs @@ -39,16 +39,7 @@ impl KeyTree { .try_into() .expect("SeedHolder seed is 64 bytes long"); - let root_keys = N::from_seed(seed_fit); - let account_id_map = root_keys - .account_ids() - .map(|id| (id, ChainIndex::root())) - .collect(); - - Self { - key_map: BTreeMap::from_iter([(ChainIndex::root(), root_keys)]), - account_id_map, - } + Self::new_from_root(N::from_seed(seed_fit)) } pub fn new_from_root(root: N) -> Self { @@ -63,6 +54,15 @@ impl KeyTree { } } + fn insert_child(&mut self, child_keys: N, chain_index: ChainIndex) -> ChainIndex { + for account_id in child_keys.account_ids() { + self.account_id_map.insert(account_id, chain_index.clone()); + } + self.key_map.insert(chain_index.clone(), child_keys); + + chain_index + } + pub fn generate_new_node(&mut self, parent_cci: &ChainIndex) -> Option { let parent_keys = self.key_map.get(parent_cci)?; let next_child_id = self @@ -71,14 +71,8 @@ impl KeyTree { let next_cci = parent_cci.nth_child(next_child_id); let child_keys = parent_keys.derive_child(next_child_id); - let account_ids = child_keys.account_ids(); - for account_id in account_ids { - self.account_id_map.insert(account_id, next_cci.clone()); - } - self.key_map.insert(next_cci.clone(), child_keys); - - Some(next_cci) + Some(self.insert_child(child_keys, next_cci)) } pub fn fill_node(&mut self, chain_index: &ChainIndex) -> Option { @@ -86,14 +80,8 @@ impl KeyTree { let child_id = *chain_index.chain().last()?; let child_keys = parent_keys.derive_child(child_id); - let account_ids = child_keys.account_ids(); - for account_id in account_ids { - self.account_id_map.insert(account_id, chain_index.clone()); - } - self.key_map.insert(chain_index.clone(), child_keys); - - Some(chain_index.clone()) + Some(self.insert_child(child_keys, chain_index.clone())) } #[must_use] @@ -200,24 +188,27 @@ impl KeyTree { } impl KeyTree { + /// Pairs `cci` with the account ID of the node stored at it. + fn account_id_for_cci(&self, cci: ChainIndex) -> Option<(lee::AccountId, ChainIndex)> { + let node = self.key_map.get(&cci)?; + let account_id = node.account_ids().next()?; + Some((account_id, cci)) + } + /// Generate a new public key node, returning the account ID and chain index. pub fn generate_new_public_node( &mut self, parent_cci: &ChainIndex, ) -> Option<(lee::AccountId, ChainIndex)> { let cci = self.generate_new_node(parent_cci)?; - let node = self.key_map.get(&cci)?; - let account_id = node.account_ids().next()?; - Some((account_id, cci)) + self.account_id_for_cci(cci) } /// Generate a new public key node using layered placement, returning the account ID and chain /// index. pub fn generate_new_public_node_layered(&mut self) -> Option<(lee::AccountId, ChainIndex)> { let cci = self.generate_new_node_layered()?; - let node = self.key_map.get(&cci)?; - let account_id = node.account_ids().next()?; - Some((account_id, cci)) + self.account_id_for_cci(cci) } /// Cleanup of non-initialized accounts in a public tree. @@ -322,6 +313,16 @@ impl KeyTree { } } +const fn split_hash(hash_value: &[u8; 64]) -> ([u8; 32], [u8; 32]) { + let first = *hash_value + .first_chunk::<32>() + .expect("hash_value is 64 bytes, must be safe to get first 32"); + let last = *hash_value + .last_chunk::<32>() + .expect("hash_value is 64 bytes, must be safe to get last 32"); + (first, last) +} + #[cfg(test)] mod tests { #![expect(clippy::shadow_unrelated, reason = "We don't care about this in tests")] @@ -339,6 +340,18 @@ mod tests { } } + #[test] + fn split_hash_splits_into_first_and_last_32_bytes() { + let mut hash_value = [0_u8; 64]; + hash_value[..32].fill(0xAA); + hash_value[32..].fill(0xBB); + + let (first, last) = split_hash(&hash_value); + + assert_eq!(first, [0xAA; 32]); + assert_eq!(last, [0xBB; 32]); + } + #[test] fn simple_key_tree() { let seed_holder = seed_holder_for_tests(); diff --git a/lee/key_protocol/src/key_management/mod.rs b/lee/key_protocol/src/key_management/mod.rs index 459badf0..3a066fd9 100644 --- a/lee/key_protocol/src/key_management/mod.rs +++ b/lee/key_protocol/src/key_management/mod.rs @@ -25,8 +25,22 @@ impl KeyChain { #[must_use] pub fn new_os_random() -> Self { // Currently dropping SeedHolder at the end of initialization. - // Now entirely sure if we need it in the future. + // Not entirely sure if we need it in the future. let seed_holder = SeedHolder::new_os_random(); + + Self::from_seed_holder(&seed_holder) + } + + #[must_use] + pub fn new_mnemonic(passphrase: &str) -> (Self, bip39::Mnemonic) { + // Currently dropping SeedHolder at the end of initialization. + // Not entirely sure if we need it in the future. + let (seed_holder, mnemonic) = SeedHolder::new_mnemonic(passphrase); + + (Self::from_seed_holder(&seed_holder), mnemonic) + } + + fn from_seed_holder(seed_holder: &SeedHolder) -> Self { let secret_spending_key = seed_holder.produce_top_secret_key_holder(); let private_key_holder = secret_spending_key.produce_private_key_holder(None); @@ -42,29 +56,6 @@ impl KeyChain { } } - #[must_use] - pub fn new_mnemonic(passphrase: &str) -> (Self, bip39::Mnemonic) { - // Currently dropping SeedHolder at the end of initialization. - // Not entirely sure if we need it in the future. - let (seed_holder, mnemonic) = SeedHolder::new_mnemonic(passphrase); - let secret_spending_key = seed_holder.produce_top_secret_key_holder(); - - let private_key_holder = secret_spending_key.produce_private_key_holder(None); - - let nullifier_public_key = private_key_holder.generate_nullifier_public_key(); - let viewing_public_key = private_key_holder.generate_viewing_public_key(); - - ( - Self { - secret_spending_key, - private_key_holder, - nullifier_public_key, - viewing_public_key, - }, - mnemonic, - ) - } - #[must_use] pub fn calculate_shared_secret_receiver( &self, diff --git a/lee/key_protocol/src/key_management/secret_holders.rs b/lee/key_protocol/src/key_management/secret_holders.rs index 7bda4ffb..b8225a4b 100644 --- a/lee/key_protocol/src/key_management/secret_holders.rs +++ b/lee/key_protocol/src/key_management/secret_holders.rs @@ -43,16 +43,7 @@ pub struct PrivateKeyHolder { impl SeedHolder { #[must_use] pub fn new_os_random() -> Self { - let mut enthopy_bytes: [u8; 32] = [0; 32]; - OsRng.fill_bytes(&mut enthopy_bytes); - - let mnemonic = Mnemonic::from_entropy(&enthopy_bytes) - .expect("Enthropy must be a multiple of 32 bytes"); - let seed_wide = mnemonic.to_seed("mnemonic"); - - Self { - seed: seed_wide.to_vec(), - } + Self::new_mnemonic("mnemonic").0 } #[must_use] @@ -62,14 +53,8 @@ impl SeedHolder { let mnemonic = Mnemonic::from_entropy(&entropy_bytes).expect("Entropy must be a multiple of 32 bytes"); - let seed_wide = mnemonic.to_seed(passphrase); - ( - Self { - seed: seed_wide.to_vec(), - }, - mnemonic, - ) + (Self::from_mnemonic(&mnemonic, passphrase), mnemonic) } #[must_use] @@ -107,10 +92,7 @@ impl SecretSpendingKey { const SUFFIX_1: &[u8; 1] = &[1]; const SUFFIX_2: &[u8; 19] = &[0; 19]; - let index = match index { - None => 0_u32, - _ => index.expect("Expect a valid u32"), - }; + let index = index.unwrap_or(0); let mut hasher = sha2::Sha256::new(); hasher.update(PREFIX); @@ -129,10 +111,7 @@ impl SecretSpendingKey { const SUFFIX_1: &[u8; 1] = &[2]; const SUFFIX_2: &[u8; 19] = &[0; 19]; - let index = match index { - None => 0_u32, - _ => index.expect("Expect a valid u32"), - }; + let index = index.unwrap_or(0); let mut bytes: Vec = Vec::with_capacity(64); bytes.extend_from_slice(PREFIX); @@ -146,14 +125,7 @@ impl SecretSpendingKey { let full_seed = hmac_sha512::HMAC::mac(bytes, b"LEE_viewing_seed"); - ViewingSecretKey::new( - *full_seed - .first_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get first 32"), - *full_seed - .last_chunk::<32>() - .expect("hash_value is 64 bytes, must be safe to get last 32"), - ) + Self::generate_viewing_secret_key(full_seed) } #[must_use] @@ -181,7 +153,7 @@ impl From<&ViewingSecretKey> for ViewingPublicKey { seed_bytes[32..].copy_from_slice(&sk.z); let dk = ::DecapsulationKey::from_seed(Seed::from(seed_bytes)); Self::from_bytes(dk.encapsulation_key().to_bytes().to_vec()) - .expect("key_protocol::secret_holders::From<&ViewingSecretKey>: ML-KEM-768 encapsulation key is always 1184 bytes") + .expect("key_protocol::secret_holders::From<&ViewingSecretKey>: ML-KEM-768 encapsulation key is always ViewingPublicKey::LEN bytes") } } @@ -201,7 +173,6 @@ impl PrivateKeyHolder { mod tests { use super::*; - // TODO? are these necessary? #[test] fn seed_generation_test() { let seed_holder = SeedHolder::new_os_random(); diff --git a/lee/privacy_preserving_circuit/src/execution_state.rs b/lee/privacy_preserving_circuit/src/execution_state.rs index 8d920068..4e81f9ce 100644 --- a/lee/privacy_preserving_circuit/src/execution_state.rs +++ b/lee/privacy_preserving_circuit/src/execution_state.rs @@ -59,46 +59,9 @@ impl ExecutionState { program_id: ProgramId, program_outputs: Vec, ) -> Self { - // Build position → (npk, identifier) map for private-PDA pre_states, indexed by position - // in `account_identities`. The vec is documented as 1:1 with the program's pre_state - // order, so position here matches `pre_state_position` used downstream in - // `validate_and_sync_states`. - let mut private_pda_npk_by_position: HashMap = - HashMap::new(); - for (pos, account_identity) in account_identities.iter().enumerate() { - if let Some((npk, identifier)) = account_identity.npk_if_private_pda() { - private_pda_npk_by_position.insert(pos, (npk, identifier)); - } - } - - let block_valid_from = program_outputs - .iter() - .filter_map(|output| output.block_validity_window.start()) - .max(); - let block_valid_until = program_outputs - .iter() - .filter_map(|output| output.block_validity_window.end()) - .min(); - let ts_valid_from = program_outputs - .iter() - .filter_map(|output| output.timestamp_validity_window.start()) - .max(); - let ts_valid_until = program_outputs - .iter() - .filter_map(|output| output.timestamp_validity_window.end()) - .min(); - - let block_validity_window: BlockValidityWindow = (block_valid_from, block_valid_until) - .try_into() - .expect( - "There should be non empty intersection in the program output block validity windows", - ); - let timestamp_validity_window: TimestampValidityWindow = - (ts_valid_from, ts_valid_until) - .try_into() - .expect( - "There should be non empty intersection in the program output timestamp validity windows", - ); + let private_pda_npk_by_position = build_private_pda_npk_map(account_identities); + let (block_validity_window, timestamp_validity_window) = + intersect_validity_windows(&program_outputs); let mut execution_state = Self { pre_states: Vec::new(), @@ -136,49 +99,7 @@ impl ExecutionState { panic!("Insufficient program outputs for chained calls"); }; - // Check that instruction data in chained call is the instruction data in program output - assert_eq!( - chained_call.instruction_data, program_output.instruction_data, - "Mismatched instruction data between chained call and program output" - ); - - // Check that `program_output` is consistent with the execution of the corresponding - // program. - let program_output_words = - &to_vec(&program_output).expect("program_output must be serializable"); - env::verify(chained_call.program_id, program_output_words).unwrap_or_else( - |_: Infallible| unreachable!("Infallible error is never constructed"), - ); - - // Verify that the program output's self_program_id matches the expected program ID. - // This ensures the proof commits to which program produced the output. - assert_eq!( - program_output.self_program_id, chained_call.program_id, - "Program output self_program_id does not match chained call program_id" - ); - - // Verify that the program output's caller_program_id matches the actual caller. - // This prevents a malicious user from privately executing an internal function - // by spoofing caller_program_id (e.g. passing caller_program_id = self_program_id - // to bypass access control checks). - assert_eq!( - program_output.caller_program_id, caller_program_id, - "Program output caller_program_id does not match actual caller" - ); - - // Check that the program is well behaved. - // See the # Programs section for the definition of the `validate_execution` method. - let validated_execution = validate_execution( - &program_output.pre_states, - &program_output.post_states, - chained_call.program_id, - ); - if let Err(err) = validated_execution { - panic!( - "Invalid program behavior in program {:?}: {err}", - chained_call.program_id - ); - } + verify_program_output(&chained_call, caller_program_id, &program_output); for next_call in program_output.chained_calls.iter().rev() { chained_calls.push_front((next_call.clone(), Some(chained_call.program_id))); @@ -202,41 +123,8 @@ impl ExecutionState { "Inner call without a chained call found", ); - // Every private-PDA pre_state must have had its npk bound to its account_id, either via - // a `Claim::Pda(seed)` in some program's post_state or via a caller's `pda_seeds` - // matching the private derivation. An unbound private-PDA pre_state has no - // cryptographic link between the supplied npk and the account_id, and must be rejected. - for (pos, account_identity) in account_identities.iter().enumerate() { - if account_identity.is_private_pda() { - assert!( - execution_state - .private_pda_bound_positions - .contains_key(&pos), - "private PDA pre_state at position {pos} has no proven (seed, npk) binding via Claim::Pda or caller pda_seeds" - ); - } - } - - // Check that all modified uninitialized accounts were claimed - for (account_id, post) in execution_state - .pre_states - .iter() - .filter(|a| a.account.program_owner == DEFAULT_PROGRAM_ID) - .map(|a| { - let post = execution_state - .post_states - .get(&a.account_id) - .expect("Post state must exist for pre state"); - (a, post) - }) - .filter(|(pre_default, post)| pre_default.account != **post) - .map(|(pre, post)| (pre.account_id, post)) - { - assert_ne!( - post.program_owner, DEFAULT_PROGRAM_ID, - "Account {account_id} was modified but not claimed" - ); - } + execution_state.assert_all_pda_positions_bound(account_identities); + execution_state.assert_modified_accounts_claimed(); execution_state } @@ -254,200 +142,185 @@ impl ExecutionState { for (pre, mut post) in output_pre_states.into_iter().zip(output_post_states) { let pre_account_id = pre.account_id; let pre_is_authorized = pre.is_authorized; - let post_states_entry = self.post_states.entry(pre.account_id); - match &post_states_entry { - Entry::Occupied(occupied) => { - #[expect( - clippy::shadow_unrelated, - reason = "Shadowing is intentional to use all fields" - )] - let AccountWithMetadata { - account: pre_account, - account_id: pre_account_id, - is_authorized: pre_is_authorized, - } = pre; - // Ensure that new pre state is the same as known post state - assert_eq!( - occupied.get(), - &pre_account, - "Inconsistent pre state for account {pre_account_id}", + if let Some(existing) = self.post_states.get(&pre.account_id) { + #[expect( + clippy::shadow_unrelated, + reason = "Shadowing is intentional to use all fields" + )] + let AccountWithMetadata { + account: pre_account, + account_id: pre_account_id, + is_authorized: pre_is_authorized, + } = pre; + + assert_eq!( + existing, &pre_account, + "Inconsistent pre state for account {pre_account_id}", + ); + + let (previous_is_authorized, pre_state_position) = self + .pre_states + .iter() + .enumerate() + .find(|(_, acc)| acc.account_id == pre_account_id) + .map_or_else( + || panic!( + "Pre state must exist in execution state for account {pre_account_id}", + ), + |(pos, acc)| (acc.is_authorized, pos), ); - let (previous_is_authorized, pre_state_position) = self - .pre_states - .iter() - .enumerate() - .find(|(_, acc)| acc.account_id == pre_account_id) - .map_or_else( - || panic!( - "Pre state must exist in execution state for account {pre_account_id}", - ), - |(pos, acc)| (acc.is_authorized, pos) - ); + let is_authorized = resolve_authorization_and_record_bindings( + &mut self.pda_family_binding, + &mut self.private_pda_bound_positions, + &self.private_pda_npk_by_position, + &mut self.authorized_accounts, + pre_account_id, + pre_state_position, + caller_program_id, + caller_pda_seeds, + previous_is_authorized, + ); - let is_authorized = resolve_authorization_and_record_bindings( - &mut self.pda_family_binding, - &mut self.private_pda_bound_positions, - &self.private_pda_npk_by_position, - &mut self.authorized_accounts, - pre_account_id, - pre_state_position, - caller_program_id, - caller_pda_seeds, - previous_is_authorized, - ); - - assert_eq!( - pre_is_authorized, is_authorized, - "Inconsistent authorization for account {pre_account_id}", - ); - } - Entry::Vacant(_) => { - // Pre state for the initial call - let pre_state_position = self.pre_states.len(); - let external_seed = match account_identities.get(pre_state_position) { - Some(InputAccountIdentity::PrivatePdaInit { - npk, - identifier, - seed: Some((seed, authority_program_id)), - .. - }) => { - let expected = AccountId::for_private_pda( - authority_program_id, - seed, - npk, - *identifier, - ); - assert_eq!( - pre_account_id, expected, - "External seed mismatch for PrivatePdaInit at position {pre_state_position}" - ); - Some((*seed, *authority_program_id)) - } - Some(InputAccountIdentity::PrivatePdaUpdate { - nsk, - identifier, - seed: Some((seed, authority_program_id)), - .. - }) => { - let npk = NullifierPublicKey::from(nsk); - let expected = AccountId::for_private_pda( - authority_program_id, - seed, - &npk, - *identifier, - ); - assert_eq!( - pre_account_id, expected, - "External seed mismatch for PrivatePdaUpdate at position {pre_state_position}" - ); - Some((*seed, *authority_program_id)) - } - _ => None, - }; - // External seed is only consulted the first time the account is seen. - // Subsequent calls need no re-check because the entry is already recorded on - // private_pda_bound_positions. - if let Some((seed, authority_program_id)) = external_seed { - assert!( - !pre.is_authorized, - "Private PDA with externally-provided seed must not be authorized at position {pre_state_position}" - ); - bind_private_pda_position( - &mut self.private_pda_bound_positions, - pre_state_position, - authority_program_id, - seed, - ); - assert_family_binding( - &mut self.pda_family_binding, - authority_program_id, - seed, - pre_account_id, - ); - } - self.pre_states.push(pre); - } + assert_eq!( + pre_is_authorized, is_authorized, + "Inconsistent authorization for account {pre_account_id}", + ); + } else { + let pre_state_position = self.pre_states.len(); + resolve_external_seed( + account_identities, + pre_state_position, + pre_account_id, + pre.is_authorized, + &mut self.private_pda_bound_positions, + &mut self.pda_family_binding, + ); + self.pre_states.push(pre); } if let Some(claim) = post.required_claim() { - // The invoked program can only claim accounts with default program id. - assert_eq!( - post.account().program_owner, - DEFAULT_PROGRAM_ID, - "Cannot claim an initialized account {pre_account_id}" + self.process_claim( + account_identities, + &mut post, + pre_account_id, + pre_is_authorized, + program_id, + claim, ); - - let pre_state_position = self - .pre_states - .iter() - .position(|acc| acc.account_id == pre_account_id) - .expect("Pre state must exist at this point"); - - let account_identity = &account_identities[pre_state_position]; - if account_identity.is_public() { - match claim { - Claim::Authorized => { - // Note: no need to check authorized pdas because we have already - // checked consistency of authorization above. - assert!( - pre_is_authorized, - "Cannot claim unauthorized account {pre_account_id}" - ); - } - Claim::Pda(seed) => { - let pda = AccountId::for_public_pda(&program_id, &seed); - assert_eq!( - pre_account_id, pda, - "Invalid PDA claim for account {pre_account_id} which does not match derived PDA {pda}" - ); - assert_family_binding( - &mut self.pda_family_binding, - program_id, - seed, - pre_account_id, - ); - } - } - } else { - // Private accounts: don't enforce the claim semantics. Unauthorized private - // claiming is intentionally allowed - match claim { - Claim::Authorized => {} - Claim::Pda(seed) => { - let (npk, identifier) = self - .private_pda_npk_by_position - .get(&pre_state_position) - .expect( - "private PDA pre_state must have an npk in the position map", - ); - let pda = - AccountId::for_private_pda(&program_id, &seed, npk, *identifier); - assert_eq!( - pre_account_id, pda, - "Invalid private PDA claim for account {pre_account_id}" - ); - bind_private_pda_position( - &mut self.private_pda_bound_positions, - pre_state_position, - program_id, - seed, - ); - assert_family_binding( - &mut self.pda_family_binding, - program_id, - seed, - pre_account_id, - ); - } - } - } - - post.account_mut().program_owner = program_id; } - post_states_entry.insert_entry(post.into_account()); + self.post_states.insert(pre_account_id, post.into_account()); + } + } + + fn process_claim( + &mut self, + account_identities: &[InputAccountIdentity], + post: &mut AccountPostState, + pre_account_id: AccountId, + pre_is_authorized: bool, + program_id: ProgramId, + claim: Claim, + ) { + assert_eq!( + post.account().program_owner, + DEFAULT_PROGRAM_ID, + "Cannot claim an initialized account {pre_account_id}" + ); + + let pre_state_position = self + .pre_states + .iter() + .position(|acc| acc.account_id == pre_account_id) + .expect("Pre state must exist at this point"); + + let account_identity = &account_identities[pre_state_position]; + if account_identity.is_public() { + match claim { + Claim::Authorized => { + assert!( + pre_is_authorized, + "Cannot claim unauthorized account {pre_account_id}" + ); + } + Claim::Pda(seed) => { + let pda = AccountId::for_public_pda(&program_id, &seed); + assert_eq!( + pre_account_id, pda, + "Invalid PDA claim for account {pre_account_id} which does not match derived PDA {pda}" + ); + assert_family_binding( + &mut self.pda_family_binding, + program_id, + seed, + pre_account_id, + ); + } + } + } else { + match claim { + Claim::Authorized => {} + Claim::Pda(seed) => { + let (npk, identifier) = self + .private_pda_npk_by_position + .get(&pre_state_position) + .expect("private PDA pre_state must have an npk in the position map"); + let pda = AccountId::for_private_pda(&program_id, &seed, npk, *identifier); + assert_eq!( + pre_account_id, pda, + "Invalid private PDA claim for account {pre_account_id}" + ); + bind_private_pda_position( + &mut self.private_pda_bound_positions, + pre_state_position, + program_id, + seed, + ); + assert_family_binding( + &mut self.pda_family_binding, + program_id, + seed, + pre_account_id, + ); + } + } + } + + post.account_mut().program_owner = program_id; + } + + fn assert_all_pda_positions_bound(&self, account_identities: &[InputAccountIdentity]) { + for (pos, account_identity) in account_identities.iter().enumerate() { + if account_identity.is_private_pda() { + assert!( + self.private_pda_bound_positions.contains_key(&pos), + "private PDA pre_state at position {pos} has no proven (seed, npk) binding via Claim::Pda or caller pda_seeds" + ); + } + } + } + + fn assert_modified_accounts_claimed(&self) { + for (account_id, post) in self + .pre_states + .iter() + .filter(|a| a.account.program_owner == DEFAULT_PROGRAM_ID) + .map(|a| { + let post = self + .post_states + .get(&a.account_id) + .expect("Post state must exist for pre state"); + (a, post) + }) + .filter(|(pre_default, post)| pre_default.account != **post) + .map(|(pre, post)| (pre.account_id, post)) + { + assert_ne!( + post.program_owner, DEFAULT_PROGRAM_ID, + "Account {account_id} was modified but not claimed" + ); } } @@ -486,6 +359,150 @@ impl ExecutionState { } } +fn verify_program_output( + chained_call: &ChainedCall, + caller_program_id: Option, + program_output: &ProgramOutput, +) { + assert_eq!( + chained_call.instruction_data, program_output.instruction_data, + "Mismatched instruction data between chained call and program output" + ); + + let program_output_words = + &to_vec(program_output).expect("program_output must be serializable"); + env::verify(chained_call.program_id, program_output_words) + .unwrap_or_else(|_: Infallible| unreachable!("Infallible error is never constructed")); + + assert_eq!( + program_output.self_program_id, chained_call.program_id, + "Program output self_program_id does not match chained call program_id" + ); + + assert_eq!( + program_output.caller_program_id, caller_program_id, + "Program output caller_program_id does not match actual caller" + ); + + if let Err(err) = validate_execution( + &program_output.pre_states, + &program_output.post_states, + chained_call.program_id, + ) { + panic!( + "Invalid program behavior in program {:?}: {err}", + chained_call.program_id + ); + } +} + +fn build_private_pda_npk_map( + account_identities: &[InputAccountIdentity], +) -> HashMap { + account_identities + .iter() + .enumerate() + .filter_map(|(pos, identity)| { + identity + .npk_if_private_pda() + .map(|(npk, identifier)| (pos, (npk, identifier))) + }) + .collect() +} + +fn intersect_validity_windows( + program_outputs: &[ProgramOutput], +) -> (BlockValidityWindow, TimestampValidityWindow) { + let block_valid_from = program_outputs + .iter() + .filter_map(|output| output.block_validity_window.start()) + .max(); + let block_valid_until = program_outputs + .iter() + .filter_map(|output| output.block_validity_window.end()) + .min(); + let ts_valid_from = program_outputs + .iter() + .filter_map(|output| output.timestamp_validity_window.start()) + .max(); + let ts_valid_until = program_outputs + .iter() + .filter_map(|output| output.timestamp_validity_window.end()) + .min(); + + let block_validity_window: BlockValidityWindow = + (block_valid_from, block_valid_until).try_into().expect( + "There should be non empty intersection in the program output block validity windows", + ); + let timestamp_validity_window: TimestampValidityWindow = (ts_valid_from, ts_valid_until) + .try_into() + .expect( + "There should be non empty intersection in the program output timestamp validity windows", + ); + + (block_validity_window, timestamp_validity_window) +} + +fn resolve_external_seed( + account_identities: &[InputAccountIdentity], + pre_state_position: usize, + pre_account_id: AccountId, + is_authorized: bool, + private_pda_bound_positions: &mut HashMap, + pda_family_binding: &mut HashMap<(ProgramId, PdaSeed), AccountId>, +) { + let external_seed = match account_identities.get(pre_state_position) { + Some(InputAccountIdentity::PrivatePdaInit { + npk, + identifier, + seed: Some((seed, authority_program_id)), + .. + }) => { + let expected = AccountId::for_private_pda(authority_program_id, seed, npk, *identifier); + assert_eq!( + pre_account_id, expected, + "External seed mismatch for PrivatePdaInit at position {pre_state_position}" + ); + Some((*seed, *authority_program_id)) + } + Some(InputAccountIdentity::PrivatePdaUpdate { + nsk, + identifier, + seed: Some((seed, authority_program_id)), + .. + }) => { + let npk = NullifierPublicKey::from(nsk); + let expected = + AccountId::for_private_pda(authority_program_id, seed, &npk, *identifier); + assert_eq!( + pre_account_id, expected, + "External seed mismatch for PrivatePdaUpdate at position {pre_state_position}" + ); + Some((*seed, *authority_program_id)) + } + _ => None, + }; + + if let Some((seed, authority_program_id)) = external_seed { + assert!( + !is_authorized, + "Private PDA with externally-provided seed must not be authorized at position {pre_state_position}" + ); + bind_private_pda_position( + private_pda_bound_positions, + pre_state_position, + authority_program_id, + seed, + ); + assert_family_binding( + pda_family_binding, + authority_program_id, + seed, + pre_account_id, + ); + } +} + /// Record or re-verify the `(program_id, seed) → account_id` family binding for the /// transaction. Any claim or caller-seed authorization that resolves a `pre_state` under /// `(program_id, seed)` must agree with every prior resolution of the same pair; otherwise a diff --git a/lee/privacy_preserving_circuit/src/output.rs b/lee/privacy_preserving_circuit/src/output.rs index 8c8ec2a4..74c1cccc 100644 --- a/lee/privacy_preserving_circuit/src/output.rs +++ b/lee/privacy_preserving_circuit/src/output.rs @@ -1,13 +1,244 @@ use lee_core::{ - Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme, - EphemeralPublicKey, InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey, - NullifierSecretKey, PrivacyPreservingCircuitOutput, PrivateAccountKind, SharedSecretKey, + Commitment, CommitmentSetDigest, EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, + InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey, + PrivacyPreservingCircuitOutput, PrivateAccountKind, SharedSecretKey, account::{Account, AccountId, Nonce}, compute_digest_for_path, }; use crate::execution_state::ExecutionState; +struct PrivateOutputHandler<'ctx> { + output: &'ctx mut PrivacyPreservingCircuitOutput, + output_index: &'ctx mut u32, + pre_state: &'ctx lee_core::account::AccountWithMetadata, + post_state: Account, + epk: &'ctx EphemeralPublicKey, + view_tag: u8, + ssk: &'ctx SharedSecretKey, + identifier: u128, +} + +impl PrivateOutputHandler<'_> { + fn authorized_init(self, nsk: &NullifierSecretKey, commitment_root: &CommitmentSetDigest) { + let npk = NullifierPublicKey::from(nsk); + let account_id = + derive_and_verify_account_id(&npk, self.identifier, self.pre_state.account_id); + + assert!( + self.pre_state.is_authorized, + "Pre-state not authorized for authenticated private account" + ); + assert_eq!( + self.pre_state.account, + Account::default(), + "Found new private account with non default values" + ); + + let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id, commitment_root); + let kind = PrivateAccountKind::Regular(self.identifier); + + self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce); + } + + fn authorized_update(self, nsk: &NullifierSecretKey, membership_proof: &MembershipProof) { + let npk = NullifierPublicKey::from(nsk); + let account_id = + derive_and_verify_account_id(&npk, self.identifier, self.pre_state.account_id); + + assert!( + self.pre_state.is_authorized, + "Pre-state not authorized for authenticated private account" + ); + + let new_nullifier = compute_update_nullifier_and_set_digest( + membership_proof, + &self.pre_state.account, + &account_id, + nsk, + ); + let new_nonce = self + .pre_state + .account + .nonce + .private_account_nonce_increment(nsk); + let kind = PrivateAccountKind::Regular(self.identifier); + + self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce); + } + + fn unauthorized(self, npk: &NullifierPublicKey, commitment_root: &CommitmentSetDigest) { + let account_id = + derive_and_verify_account_id(npk, self.identifier, self.pre_state.account_id); + + assert_eq!( + self.pre_state.account, + Account::default(), + "Found new private account with non default values", + ); + assert!( + !self.pre_state.is_authorized, + "Found new private account marked as authorized." + ); + + let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id, commitment_root); + let kind = PrivateAccountKind::Regular(self.identifier); + + self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce); + } + + fn pda_init( + self, + commitment_root: &CommitmentSetDigest, + pos: usize, + pda_seed_by_position: &std::collections::HashMap< + usize, + (lee_core::program::ProgramId, lee_core::program::PdaSeed), + >, + ) { + // The npk-to-account_id binding is established upstream in + // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds` + // match. Here we only enforce the init pre-conditions. + assert!( + !self.pre_state.is_authorized, + "PrivatePdaInit requires unauthorized pre_state" + ); + assert_eq!( + self.pre_state.account, + Account::default(), + "New private PDA must be default" + ); + + let (new_nullifier, new_nonce) = + init_nullifier_and_nonce(&self.pre_state.account_id, commitment_root); + + let account_id = self.pre_state.account_id; + let (authority_program_id, seed) = pda_seed_by_position + .get(&pos) + .expect("PrivatePdaInit position must be in pda_seed_by_position"); + let kind = PrivateAccountKind::Pda { + program_id: *authority_program_id, + seed: *seed, + identifier: self.identifier, + }; + + self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce); + } + + fn pda_update( + self, + nsk: &NullifierSecretKey, + membership_proof: &MembershipProof, + external_seed: Option<&(lee_core::program::PdaSeed, lee_core::program::ProgramId)>, + pos: usize, + pda_seed_by_position: &std::collections::HashMap< + usize, + (lee_core::program::ProgramId, lee_core::program::PdaSeed), + >, + ) { + // With an external seed the binding comes from the circuit input and the + // pre_state is intentionally unauthorized; without one the binding comes from + // a Claim or caller pda_seeds, so the pre_state must already be authorized. + assert!( + self.pre_state.is_authorized ^ external_seed.is_some(), + "PrivatePdaUpdate requires authorized pre_state or external seed" + ); + + let new_nullifier = compute_update_nullifier_and_set_digest( + membership_proof, + &self.pre_state.account, + &self.pre_state.account_id, + nsk, + ); + let new_nonce = self + .pre_state + .account + .nonce + .private_account_nonce_increment(nsk); + + let account_id = self.pre_state.account_id; + let (authority_program_id, seed) = pda_seed_by_position + .get(&pos) + .expect("PrivatePdaUpdate position must be in pda_seed_by_position"); + let kind = PrivateAccountKind::Pda { + program_id: *authority_program_id, + seed: *seed, + identifier: self.identifier, + }; + + self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce); + } + + fn emit_private_output( + self, + account_id: &AccountId, + kind: &PrivateAccountKind, + new_nullifier: (Nullifier, CommitmentSetDigest), + new_nonce: Nonce, + ) { + self.output.new_nullifiers.push(new_nullifier); + + let mut post_with_updated_nonce = self.post_state; + post_with_updated_nonce.nonce = new_nonce; + + let commitment_post = Commitment::new(account_id, &post_with_updated_nonce); + let encrypted_account = EncryptionScheme::encrypt( + &post_with_updated_nonce, + kind, + self.ssk, + &commitment_post, + *self.output_index, + ); + + self.output.new_commitments.push(commitment_post); + self.output + .encrypted_private_post_states + .push(EncryptedAccountData { + ciphertext: encrypted_account, + epk: self.epk.clone(), + view_tag: self.view_tag, + }); + *self.output_index = self + .output_index + .checked_add(1) + .unwrap_or_else(|| panic!("Too many private accounts, output index overflow")); + } +} + +fn init_nullifier_and_nonce( + account_id: &AccountId, + commitment_root: &CommitmentSetDigest, +) -> ((Nullifier, CommitmentSetDigest), Nonce) { + let nullifier = ( + Nullifier::for_account_initialization(account_id), + *commitment_root, + ); + let nonce = Nonce::private_account_nonce_init(account_id); + (nullifier, nonce) +} + +fn derive_and_verify_account_id( + npk: &NullifierPublicKey, + identifier: u128, + pre_state_account_id: AccountId, +) -> AccountId { + let account_id = AccountId::for_regular_private_account(npk, identifier); + assert_eq!(account_id, pre_state_account_id, "AccountId mismatch"); + account_id +} + +fn compute_update_nullifier_and_set_digest( + membership_proof: &MembershipProof, + pre_account: &Account, + account_id: &AccountId, + nsk: &NullifierSecretKey, +) -> (Nullifier, CommitmentSetDigest) { + let commitment_pre = Commitment::new(account_id, pre_account); + let set_digest = compute_digest_for_path(&commitment_pre, membership_proof); + let nullifier = Nullifier::for_account_update(&commitment_pre, nsk); + (nullifier, set_digest) +} + pub fn compute_circuit_output( execution_state: ExecutionState, account_identities: &[InputAccountIdentity], @@ -45,40 +276,18 @@ pub fn compute_circuit_output( ssk, nsk, identifier, - } => { - let npk = NullifierPublicKey::from(nsk); - let account_id = AccountId::for_regular_private_account(&npk, *identifier); - - assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); - assert!( - pre_state.is_authorized, - "Pre-state not authorized for authenticated private account" - ); - assert_eq!( - pre_state.account, - Account::default(), - "Found new private account with non default values" - ); - - let new_nullifier = ( - Nullifier::for_account_initialization(&account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = Nonce::private_account_nonce_init(&account_id); - - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - &PrivateAccountKind::Regular(*identifier), - ssk, - epk, - *view_tag, - new_nullifier, - new_nonce, - ); + commitment_root, + } => PrivateOutputHandler { + output: &mut output, + output_index: &mut output_index, + pre_state: &pre_state, + post_state, + epk, + view_tag: *view_tag, + ssk, + identifier: *identifier, } + .authorized_init(nsk, commitment_root), InputAccountIdentity::PrivateAuthorizedUpdate { epk, view_tag, @@ -86,127 +295,54 @@ pub fn compute_circuit_output( nsk, membership_proof, identifier, - } => { - let npk = NullifierPublicKey::from(nsk); - let account_id = AccountId::for_regular_private_account(&npk, *identifier); - - assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); - assert!( - pre_state.is_authorized, - "Pre-state not authorized for authenticated private account" - ); - - let new_nullifier = compute_update_nullifier_and_set_digest( - membership_proof, - &pre_state.account, - &account_id, - nsk, - ); - let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); - - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - &PrivateAccountKind::Regular(*identifier), - ssk, - epk, - *view_tag, - new_nullifier, - new_nonce, - ); + } => PrivateOutputHandler { + output: &mut output, + output_index: &mut output_index, + pre_state: &pre_state, + post_state, + epk, + view_tag: *view_tag, + ssk, + identifier: *identifier, } + .authorized_update(nsk, membership_proof), InputAccountIdentity::PrivateUnauthorized { epk, view_tag, npk, ssk, identifier, - } => { - let account_id = AccountId::for_regular_private_account(npk, *identifier); - - assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); - assert_eq!( - pre_state.account, - Account::default(), - "Found new private account with non default values", - ); - assert!( - !pre_state.is_authorized, - "Found new private account marked as authorized." - ); - - let new_nullifier = ( - Nullifier::for_account_initialization(&account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = Nonce::private_account_nonce_init(&account_id); - - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - &PrivateAccountKind::Regular(*identifier), - ssk, - epk, - *view_tag, - new_nullifier, - new_nonce, - ); + commitment_root, + } => PrivateOutputHandler { + output: &mut output, + output_index: &mut output_index, + pre_state: &pre_state, + post_state, + epk, + view_tag: *view_tag, + ssk, + identifier: *identifier, } + .unauthorized(npk, commitment_root), InputAccountIdentity::PrivatePdaInit { epk, view_tag, npk: _, ssk, identifier, + commitment_root, seed: _, - } => { - // The npk-to-account_id binding is established upstream in - // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds` - // match. Here we only enforce the init pre-conditions. The supplied npk on - // the variant has been recorded into `private_pda_npk_by_position` and used - // for the binding check; we use `pre_state.account_id` directly for nullifier - // and commitment derivation. - assert!( - !pre_state.is_authorized, - "PrivatePdaInit requires unauthorized pre_state" - ); - assert_eq!( - pre_state.account, - Account::default(), - "New private PDA must be default" - ); - - let new_nullifier = ( - Nullifier::for_account_initialization(&pre_state.account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id); - - let account_id = pre_state.account_id; - let (authority_program_id, seed) = pda_seed_by_position - .get(&pos) - .expect("PrivatePdaInit position must be in pda_seed_by_position"); - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - &PrivateAccountKind::Pda { - program_id: *authority_program_id, - seed: *seed, - identifier: *identifier, - }, - ssk, - epk, - *view_tag, - new_nullifier, - new_nonce, - ); + } => PrivateOutputHandler { + output: &mut output, + output_index: &mut output_index, + pre_state: &pre_state, + post_state, + epk, + view_tag: *view_tag, + ssk, + identifier: *identifier, } + .pda_init(commitment_root, pos, &pda_seed_by_position), InputAccountIdentity::PrivatePdaUpdate { epk, view_tag, @@ -215,103 +351,25 @@ pub fn compute_circuit_output( membership_proof, identifier, seed: external_seed, - } => { - // With an external seed the binding comes from the circuit input and the - // pre_state is intentionally unauthorized; without one the binding comes from - // a Claim or caller pda_seeds, so the pre_state must already be authorized. - // When `external_seed` is `Some`, execution_state already asserted - // `!pre_state.is_authorized`. - assert!( - pre_state.is_authorized ^ external_seed.is_some(), - "PrivatePdaUpdate requires authorized pre_state or external seed" - ); - - let new_nullifier = compute_update_nullifier_and_set_digest( - membership_proof, - &pre_state.account, - &pre_state.account_id, - nsk, - ); - let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); - - let account_id = pre_state.account_id; - let (authority_program_id, seed) = pda_seed_by_position - .get(&pos) - .expect("PrivatePdaUpdate position must be in pda_seed_by_position"); - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - &PrivateAccountKind::Pda { - program_id: *authority_program_id, - seed: *seed, - identifier: *identifier, - }, - ssk, - epk, - *view_tag, - new_nullifier, - new_nonce, - ); + } => PrivateOutputHandler { + output: &mut output, + output_index: &mut output_index, + pre_state: &pre_state, + post_state, + epk, + view_tag: *view_tag, + ssk, + identifier: *identifier, } + .pda_update( + nsk, + membership_proof, + external_seed.as_ref(), + pos, + &pda_seed_by_position, + ), } } output } - -#[expect( - clippy::too_many_arguments, - reason = "Inputs are distinct concerns from the variant arms; bundling would be artificial" -)] -fn emit_private_output( - output: &mut PrivacyPreservingCircuitOutput, - output_index: &mut u32, - post_state: Account, - account_id: &AccountId, - kind: &PrivateAccountKind, - shared_secret: &SharedSecretKey, - epk: &EphemeralPublicKey, - view_tag: u8, - new_nullifier: (Nullifier, CommitmentSetDigest), - new_nonce: Nonce, -) { - output.new_nullifiers.push(new_nullifier); - - let mut post_with_updated_nonce = post_state; - post_with_updated_nonce.nonce = new_nonce; - - let commitment_post = Commitment::new(account_id, &post_with_updated_nonce); - let encrypted_account = EncryptionScheme::encrypt( - &post_with_updated_nonce, - kind, - shared_secret, - &commitment_post, - *output_index, - ); - - output.new_commitments.push(commitment_post); - output - .encrypted_private_post_states - .push(EncryptedAccountData { - ciphertext: encrypted_account, - epk: epk.clone(), - view_tag, - }); - *output_index = output_index - .checked_add(1) - .unwrap_or_else(|| panic!("Too many private accounts, output index overflow")); -} - -fn compute_update_nullifier_and_set_digest( - membership_proof: &MembershipProof, - pre_account: &Account, - account_id: &AccountId, - nsk: &NullifierSecretKey, -) -> (Nullifier, CommitmentSetDigest) { - let commitment_pre = Commitment::new(account_id, pre_account); - let set_digest = compute_digest_for_path(&commitment_pre, membership_proof); - let nullifier = Nullifier::for_account_update(&commitment_pre, nsk); - (nullifier, set_digest) -} diff --git a/lee/state_machine/core/src/circuit_io.rs b/lee/state_machine/core/src/circuit_io.rs index 78bfa24f..88b214d4 100644 --- a/lee/state_machine/core/src/circuit_io.rs +++ b/lee/state_machine/core/src/circuit_io.rs @@ -38,6 +38,7 @@ pub enum InputAccountIdentity { ssk: SharedSecretKey, nsk: NullifierSecretKey, identifier: Identifier, + commitment_root: CommitmentSetDigest, }, /// Update of an authorized standalone private account: existing on-chain commitment, with /// membership proof. @@ -57,6 +58,7 @@ pub enum InputAccountIdentity { npk: NullifierPublicKey, ssk: SharedSecretKey, identifier: Identifier, + commitment_root: CommitmentSetDigest, }, /// Init of a private PDA, unauthorized. The npk-to-account_id binding is proven upstream /// via `Claim::Pda(seed)` or a caller's `pda_seeds` match. The identifier diversifies the @@ -68,6 +70,7 @@ pub enum InputAccountIdentity { npk: NullifierPublicKey, ssk: SharedSecretKey, identifier: Identifier, + commitment_root: CommitmentSetDigest, /// When `Some((seed, authority_program_id))`, the circuit binds this position via the /// external derivation check /// `AccountId::for_private_pda(authority_program_id, seed, npk, identifier) == diff --git a/lee/state_machine/core/src/encoding.rs b/lee/state_machine/core/src/encoding.rs index 59df4b06..0fe3e08f 100644 --- a/lee/state_machine/core/src/encoding.rs +++ b/lee/state_machine/core/src/encoding.rs @@ -7,7 +7,7 @@ use std::io::Read as _; #[cfg(feature = "host")] use crate::Nullifier; #[cfg(feature = "host")] -use crate::encryption::EphemeralPublicKey; +use crate::encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN}; #[cfg(feature = "host")] use crate::error::LeeCoreError; use crate::{ @@ -168,7 +168,7 @@ impl EphemeralPublicKey { /// Deserializes an ML-KEM-768 ciphertext from a cursor. /// Reads exactly 1088 bytes — the fixed ciphertext size for ML-KEM-768. pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result { - let mut value = vec![0_u8; 1088]; + let mut value = vec![0_u8; ML_KEM_768_CIPHERTEXT_LEN]; cursor.read_exact(&mut value)?; Ok(Self(value)) } diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs index 5fa80b60..47e0995d 100644 --- a/lee/state_machine/core/src/encryption/mod.rs +++ b/lee/state_machine/core/src/encryption/mod.rs @@ -12,13 +12,16 @@ use crate::{Commitment, account::Account, program::PrivateAccountKind}; #[cfg(feature = "host")] pub mod shared_key_derivation; +/// Length in bytes of an ML-KEM-768 ciphertext (the `EphemeralPublicKey` payload). +pub const ML_KEM_768_CIPHERTEXT_LEN: usize = 1088; + pub type Scalar = [u8; 32]; #[derive(Serialize, Deserialize, Clone, Copy)] pub struct SharedSecretKey(pub [u8; 32]); /// The ML-KEM-768 ciphertext produced during encapsulation; transmitted on-wire in place of the -/// former ECDH ephemeral public key. Always 1088 bytes for ML-KEM-768. +/// former ECDH ephemeral public key. Always `ML_KEM_768_CIPHERTEXT_LEN` (1088) bytes. #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct EphemeralPublicKey(pub Vec); diff --git a/lee/state_machine/core/src/encryption/shared_key_derivation.rs b/lee/state_machine/core/src/encryption/shared_key_derivation.rs index 5c982c6f..4b46165a 100644 --- a/lee/state_machine/core/src/encryption/shared_key_derivation.rs +++ b/lee/state_machine/core/src/encryption/shared_key_derivation.rs @@ -64,7 +64,7 @@ impl SharedSecretKey { let ek_bytes: ml_kem::kem::Key = ek.0.as_slice() .try_into() - .expect("MlKem768EncapsulationKey must be 1184 bytes"); + .expect("MlKem768EncapsulationKey must be MlKem768EncapsulationKey::LEN bytes"); let ek_obj = ml_kem::EncapsulationKey768::new(&ek_bytes).expect( "MlKem768EncapsulationKey bytes must encode a valid ML-KEM-768 encapsulation key", ); @@ -104,7 +104,7 @@ impl SharedSecretKey { let ek_bytes: ml_kem::kem::Key = ek.0.as_slice() .try_into() - .expect("MlKem768EncapsulationKey must be 1184 bytes"); + .expect("MlKem768EncapsulationKey must be MlKem768EncapsulationKey::LEN bytes"); let ek_obj = ml_kem::EncapsulationKey768::new(&ek_bytes).expect( "MlKem768EncapsulationKey bytes must encode a valid ML-KEM-768 encapsulation key", ); @@ -118,8 +118,9 @@ impl SharedSecretKey { /// Receiver: decapsulate the shared secret from a KEM ciphertext. /// - /// Returns `None` if the `EphemeralPublicKey` is not exactly 1088 bytes — callers on - /// the wallet scan path should skip the output rather than panic on malformed chain data. + /// Returns `None` if the `EphemeralPublicKey` is not exactly [`ML_KEM_768_CIPHERTEXT_LEN`] + /// bytes — callers on the wallet scan path should skip the output rather than panic on + /// malformed chain data. /// /// `d` and `z` are the two 32-byte halves of the FIPS 203 `ViewingSecretKey` seed. #[must_use] @@ -146,6 +147,7 @@ mod tests { use ml_kem::KeyExport as _; use super::*; + use crate::ML_KEM_768_CIPHERTEXT_LEN; #[test] fn encapsulate_decapsulate_round_trip() { @@ -164,11 +166,15 @@ mod tests { let receiver_ss = SharedSecretKey::decapsulate(&epk, &d, &z).unwrap(); assert_eq!(sender_ss.0, receiver_ss.0, "shared secrets must match"); - assert_eq!(epk.0.len(), 1088, "ML-KEM-768 ciphertext is 1088 bytes"); + assert_eq!( + epk.0.len(), + ML_KEM_768_CIPHERTEXT_LEN, + "ML-KEM-768 ciphertext length" + ); assert_eq!( ek.0.len(), - 1184, - "ML-KEM-768 encapsulation key is 1184 bytes" + MlKem768EncapsulationKey::LEN, + "ML-KEM-768 encapsulation key length" ); } @@ -177,15 +183,15 @@ mod tests { let d = [1_u8; 32]; let z = [2_u8; 32]; - // Too short — 100 bytes instead of 1088. + // Too short — 100 bytes instead of ML_KEM_768_CIPHERTEXT_LEN. let short_epk = EphemeralPublicKey(vec![42_u8; 100]); assert!( SharedSecretKey::decapsulate(&short_epk, &d, &z).is_none(), "short EphemeralPublicKey must return None" ); - // Too long — 1089 bytes instead of 1088. - let long_epk = EphemeralPublicKey(vec![42_u8; 1089]); + // Too long — ML_KEM_768_CIPHERTEXT_LEN + 1. + let long_epk = EphemeralPublicKey(vec![42_u8; ML_KEM_768_CIPHERTEXT_LEN + 1]); assert!( SharedSecretKey::decapsulate(&long_epk, &d, &z).is_none(), "long EphemeralPublicKey must return None" diff --git a/lee/state_machine/core/src/lib.rs b/lee/state_machine/core/src/lib.rs index 9ad2858e..900de667 100644 --- a/lee/state_machine/core/src/lib.rs +++ b/lee/state_machine/core/src/lib.rs @@ -11,7 +11,8 @@ pub use commitment::{ compute_digest_for_path, }; pub use encryption::{ - EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, SharedSecretKey, ViewTag, + EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, + SharedSecretKey, ViewTag, }; pub use nullifier::{Identifier, Nullifier, NullifierPublicKey, NullifierSecretKey}; pub use program::PrivateAccountKind; diff --git a/lee/state_machine/core/src/program.rs b/lee/state_machine/core/src/program/mod.rs similarity index 67% rename from lee/state_machine/core/src/program.rs rename to lee/state_machine/core/src/program/mod.rs index f7791e85..287751d2 100644 --- a/lee/state_machine/core/src/program.rs +++ b/lee/state_machine/core/src/program/mod.rs @@ -230,7 +230,7 @@ impl ChainedCall { /// Represents the final state of an `Account` after a program execution. /// /// A post state may optionally request that the executing program -/// becomes the owner of the account (a “claim”). This is used to signal +/// becomes the owner of the account (a "claim"). This is used to signal /// that the program intends to take ownership of the account. #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(any(feature = "host", test), derive(PartialEq, Eq))] @@ -766,338 +766,4 @@ fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> boo } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn validity_window_unbounded_accepts_any_value() { - let w: ValidityWindow = ValidityWindow::new_unbounded(); - assert!(w.is_valid_for(0)); - assert!(w.is_valid_for(u64::MAX)); - } - - #[test] - fn validity_window_bounded_range_includes_from_excludes_to() { - let w: ValidityWindow = (Some(5), Some(10)).try_into().unwrap(); - assert!(!w.is_valid_for(4)); - assert!(w.is_valid_for(5)); - assert!(w.is_valid_for(9)); - assert!(!w.is_valid_for(10)); - } - - #[test] - fn validity_window_only_from_bound() { - let w: ValidityWindow = (Some(5), None).try_into().unwrap(); - assert!(!w.is_valid_for(4)); - assert!(w.is_valid_for(5)); - assert!(w.is_valid_for(u64::MAX)); - } - - #[test] - fn validity_window_only_to_bound() { - let w: ValidityWindow = (None, Some(5)).try_into().unwrap(); - assert!(w.is_valid_for(0)); - assert!(w.is_valid_for(4)); - assert!(!w.is_valid_for(5)); - } - - #[test] - fn validity_window_adjacent_bounds_are_invalid() { - // [5, 5) is an empty range — from == to - assert!(ValidityWindow::::try_from((Some(5), Some(5))).is_err()); - } - - #[test] - fn validity_window_inverted_bounds_are_invalid() { - assert!(ValidityWindow::::try_from((Some(10), Some(5))).is_err()); - } - - #[test] - fn validity_window_getters_match_construction() { - let w: ValidityWindow = (Some(3), Some(7)).try_into().unwrap(); - assert_eq!(w.start(), Some(3)); - assert_eq!(w.end(), Some(7)); - } - - #[test] - fn validity_window_getters_for_unbounded() { - let w: ValidityWindow = ValidityWindow::new_unbounded(); - assert_eq!(w.start(), None); - assert_eq!(w.end(), None); - } - - #[test] - fn validity_window_from_range() { - let w: ValidityWindow = ValidityWindow::try_from(5_u64..10).unwrap(); - assert_eq!(w.start(), Some(5)); - assert_eq!(w.end(), Some(10)); - } - - #[test] - fn validity_window_from_range_empty_is_invalid() { - assert!(ValidityWindow::::try_from(5_u64..5).is_err()); - } - - #[test] - fn validity_window_from_range_inverted_is_invalid() { - let from = 10_u64; - let to = 5_u64; - assert!(ValidityWindow::::try_from(from..to).is_err()); - } - - #[test] - fn validity_window_from_range_from() { - let w: ValidityWindow = (5_u64..).into(); - assert_eq!(w.start(), Some(5)); - assert_eq!(w.end(), None); - } - - #[test] - fn validity_window_from_range_to() { - let w: ValidityWindow = (..10_u64).into(); - assert_eq!(w.start(), None); - assert_eq!(w.end(), Some(10)); - } - - #[test] - fn validity_window_from_range_full() { - let w: ValidityWindow = (..).into(); - assert_eq!(w.start(), None); - assert_eq!(w.end(), None); - } - - #[test] - fn program_output_try_with_block_validity_window_range() { - let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) - .try_with_block_validity_window(10_u64..100) - .unwrap(); - assert_eq!(output.block_validity_window.start(), Some(10)); - assert_eq!(output.block_validity_window.end(), Some(100)); - } - - #[test] - fn program_output_with_block_validity_window_range_from() { - let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) - .with_block_validity_window(10_u64..); - assert_eq!(output.block_validity_window.start(), Some(10)); - assert_eq!(output.block_validity_window.end(), None); - } - - #[test] - fn program_output_with_block_validity_window_range_to() { - let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) - .with_block_validity_window(..100_u64); - assert_eq!(output.block_validity_window.start(), None); - assert_eq!(output.block_validity_window.end(), Some(100)); - } - - #[test] - fn program_output_try_with_block_validity_window_empty_range_fails() { - let result = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) - .try_with_block_validity_window(5_u64..5); - assert!(result.is_err()); - } - - #[test] - fn post_state_new_with_claim_constructor() { - let account = Account { - program_owner: [1, 2, 3, 4, 5, 6, 7, 8], - balance: 1337, - data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), - nonce: 10_u128.into(), - }; - - let account_post_state = AccountPostState::new_claimed(account.clone(), Claim::Authorized); - - assert_eq!(account, account_post_state.account); - assert_eq!(account_post_state.required_claim(), Some(Claim::Authorized)); - } - - #[test] - fn post_state_new_without_claim_constructor() { - let account = Account { - program_owner: [1, 2, 3, 4, 5, 6, 7, 8], - balance: 1337, - data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), - nonce: 10_u128.into(), - }; - - let account_post_state = AccountPostState::new(account.clone()); - - assert_eq!(account, account_post_state.account); - assert!(account_post_state.required_claim().is_none()); - } - - #[test] - fn post_state_account_getter() { - let mut account = Account { - program_owner: [1, 2, 3, 4, 5, 6, 7, 8], - balance: 1337, - data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), - nonce: 10_u128.into(), - }; - - let mut account_post_state = AccountPostState::new(account.clone()); - - assert_eq!(account_post_state.account(), &account); - assert_eq!(account_post_state.account_mut(), &mut account); - } - - // ---- AccountId::for_private_pda tests ---- - - /// Pins `AccountId::for_private_pda` against a hardcoded expected output for a specific - /// `(program_id, seed, npk, identifier)` tuple. Any change to `PRIVATE_PDA_PREFIX`, byte - /// ordering, or the underlying hash breaks this test. - #[test] - fn for_private_pda_matches_pinned_value() { - let program_id: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let npk = NullifierPublicKey([3; 32]); - let identifier: Identifier = u128::MAX; - let expected = AccountId::new([ - 59, 239, 182, 97, 14, 220, 96, 115, 238, 133, 143, 33, 234, 82, 237, 255, 148, 110, 54, - 124, 98, 159, 245, 101, 146, 182, 150, 54, 37, 62, 25, 17, - ]); - assert_eq!( - AccountId::for_private_pda(&program_id, &seed, &npk, identifier), - expected - ); - } - - /// Two groups with different viewing keys at the same (program, seed) get different addresses. - #[test] - fn for_private_pda_differs_for_different_npk() { - let program_id: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let npk_a = NullifierPublicKey([3; 32]); - let npk_b = NullifierPublicKey([4; 32]); - assert_ne!( - AccountId::for_private_pda(&program_id, &seed, &npk_a, u128::MAX), - AccountId::for_private_pda(&program_id, &seed, &npk_b, u128::MAX), - ); - } - - /// Different seeds produce different addresses, even with the same program and npk. - #[test] - fn for_private_pda_differs_for_different_seed() { - let program_id: ProgramId = [1; 8]; - let seed_a = PdaSeed::new([2; 32]); - let seed_b = PdaSeed::new([5; 32]); - let npk = NullifierPublicKey([3; 32]); - assert_ne!( - AccountId::for_private_pda(&program_id, &seed_a, &npk, u128::MAX), - AccountId::for_private_pda(&program_id, &seed_b, &npk, u128::MAX), - ); - } - - /// Different programs produce different addresses, even with the same seed and npk. - #[test] - fn for_private_pda_differs_for_different_program_id() { - let program_id_a: ProgramId = [1; 8]; - let program_id_b: ProgramId = [9; 8]; - let seed = PdaSeed::new([2; 32]); - let npk = NullifierPublicKey([3; 32]); - assert_ne!( - AccountId::for_private_pda(&program_id_a, &seed, &npk, u128::MAX), - AccountId::for_private_pda(&program_id_b, &seed, &npk, u128::MAX), - ); - } - - /// Different identifiers produce different addresses for the same `(program_id, seed, npk)`, - /// confirming that each `(program_id, seed, npk)` tuple controls a family of 2^128 addresses. - #[test] - fn for_private_pda_differs_for_different_identifier() { - let program_id: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let npk = NullifierPublicKey([3; 32]); - assert_ne!( - AccountId::for_private_pda(&program_id, &seed, &npk, 0), - AccountId::for_private_pda(&program_id, &seed, &npk, 1), - ); - assert_ne!( - AccountId::for_private_pda(&program_id, &seed, &npk, 0), - AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX), - ); - } - - /// A private PDA at the same (program, seed) has a different address than a public PDA, - /// because the private formula uses a different prefix and includes npk. - #[test] - fn for_private_pda_differs_from_public_pda() { - let program_id: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let npk = NullifierPublicKey([3; 32]); - let private_id = AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX); - let public_id = AccountId::for_public_pda(&program_id, &seed); - assert_ne!(private_id, public_id); - } - - #[cfg(feature = "host")] - #[test] - fn private_account_kind_header_round_trips() { - let regular = PrivateAccountKind::Regular(42); - let pda = PrivateAccountKind::Pda { - program_id: [1_u32; 8], - seed: PdaSeed::new([2_u8; 32]), - identifier: u128::MAX, - }; - assert_eq!( - PrivateAccountKind::from_header_bytes(®ular.to_header_bytes()), - Some(regular) - ); - assert_eq!( - PrivateAccountKind::from_header_bytes(&pda.to_header_bytes()), - Some(pda) - ); - } - - #[cfg(feature = "host")] - #[test] - fn private_account_kind_unknown_discriminant_returns_none() { - let mut bytes = [0_u8; PrivateAccountKind::HEADER_LEN]; - bytes[0] = 0xFF; - assert_eq!(PrivateAccountKind::from_header_bytes(&bytes), None); - } - - #[test] - fn for_private_account_dispatches_correctly() { - let program_id: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let npk = NullifierPublicKey([3; 32]); - let identifier: Identifier = 77; - - assert_eq!( - AccountId::for_private_account(&npk, &PrivateAccountKind::Regular(identifier)), - AccountId::for_regular_private_account(&npk, identifier), - ); - assert_eq!( - AccountId::for_private_account( - &npk, - &PrivateAccountKind::Pda { - program_id, - seed, - identifier - } - ), - AccountId::for_private_pda(&program_id, &seed, &npk, identifier), - ); - } - - #[test] - fn compute_public_authorized_pdas_with_seeds() { - let caller: ProgramId = [1; 8]; - let seed = PdaSeed::new([2; 32]); - let result = compute_public_authorized_pdas(Some(caller), &[seed]); - let expected = AccountId::for_public_pda(&caller, &seed); - assert!(result.contains(&expected)); - assert_eq!(result.len(), 1); - } - - /// With no caller (top-level call), the result is always empty. - #[test] - fn compute_public_authorized_pdas_no_caller_returns_empty() { - let seed = PdaSeed::new([2; 32]); - let result = compute_public_authorized_pdas(None, &[seed]); - assert!(result.is_empty()); - } -} +mod tests; diff --git a/lee/state_machine/core/src/program/tests.rs b/lee/state_machine/core/src/program/tests.rs new file mode 100644 index 00000000..df9f5126 --- /dev/null +++ b/lee/state_machine/core/src/program/tests.rs @@ -0,0 +1,333 @@ +use super::*; + +#[test] +fn validity_window_unbounded_accepts_any_value() { + let w: ValidityWindow = ValidityWindow::new_unbounded(); + assert!(w.is_valid_for(0)); + assert!(w.is_valid_for(u64::MAX)); +} + +#[test] +fn validity_window_bounded_range_includes_from_excludes_to() { + let w: ValidityWindow = (Some(5), Some(10)).try_into().unwrap(); + assert!(!w.is_valid_for(4)); + assert!(w.is_valid_for(5)); + assert!(w.is_valid_for(9)); + assert!(!w.is_valid_for(10)); +} + +#[test] +fn validity_window_only_from_bound() { + let w: ValidityWindow = (Some(5), None).try_into().unwrap(); + assert!(!w.is_valid_for(4)); + assert!(w.is_valid_for(5)); + assert!(w.is_valid_for(u64::MAX)); +} + +#[test] +fn validity_window_only_to_bound() { + let w: ValidityWindow = (None, Some(5)).try_into().unwrap(); + assert!(w.is_valid_for(0)); + assert!(w.is_valid_for(4)); + assert!(!w.is_valid_for(5)); +} + +#[test] +fn validity_window_adjacent_bounds_are_invalid() { + // [5, 5) is an empty range — from == to + assert!(ValidityWindow::::try_from((Some(5), Some(5))).is_err()); +} + +#[test] +fn validity_window_inverted_bounds_are_invalid() { + assert!(ValidityWindow::::try_from((Some(10), Some(5))).is_err()); +} + +#[test] +fn validity_window_getters_match_construction() { + let w: ValidityWindow = (Some(3), Some(7)).try_into().unwrap(); + assert_eq!(w.start(), Some(3)); + assert_eq!(w.end(), Some(7)); +} + +#[test] +fn validity_window_getters_for_unbounded() { + let w: ValidityWindow = ValidityWindow::new_unbounded(); + assert_eq!(w.start(), None); + assert_eq!(w.end(), None); +} + +#[test] +fn validity_window_from_range() { + let w: ValidityWindow = ValidityWindow::try_from(5_u64..10).unwrap(); + assert_eq!(w.start(), Some(5)); + assert_eq!(w.end(), Some(10)); +} + +#[test] +fn validity_window_from_range_empty_is_invalid() { + assert!(ValidityWindow::::try_from(5_u64..5).is_err()); +} + +#[test] +fn validity_window_from_range_inverted_is_invalid() { + let from = 10_u64; + let to = 5_u64; + assert!(ValidityWindow::::try_from(from..to).is_err()); +} + +#[test] +fn validity_window_from_range_from() { + let w: ValidityWindow = (5_u64..).into(); + assert_eq!(w.start(), Some(5)); + assert_eq!(w.end(), None); +} + +#[test] +fn validity_window_from_range_to() { + let w: ValidityWindow = (..10_u64).into(); + assert_eq!(w.start(), None); + assert_eq!(w.end(), Some(10)); +} + +#[test] +fn validity_window_from_range_full() { + let w: ValidityWindow = (..).into(); + assert_eq!(w.start(), None); + assert_eq!(w.end(), None); +} + +#[test] +fn program_output_try_with_block_validity_window_range() { + let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) + .try_with_block_validity_window(10_u64..100) + .unwrap(); + assert_eq!(output.block_validity_window.start(), Some(10)); + assert_eq!(output.block_validity_window.end(), Some(100)); +} + +#[test] +fn program_output_with_block_validity_window_range_from() { + let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) + .with_block_validity_window(10_u64..); + assert_eq!(output.block_validity_window.start(), Some(10)); + assert_eq!(output.block_validity_window.end(), None); +} + +#[test] +fn program_output_with_block_validity_window_range_to() { + let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) + .with_block_validity_window(..100_u64); + assert_eq!(output.block_validity_window.start(), None); + assert_eq!(output.block_validity_window.end(), Some(100)); +} + +#[test] +fn program_output_try_with_block_validity_window_empty_range_fails() { + let result = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![]) + .try_with_block_validity_window(5_u64..5); + assert!(result.is_err()); +} + +#[test] +fn post_state_new_with_claim_constructor() { + let account = Account { + program_owner: [1, 2, 3, 4, 5, 6, 7, 8], + balance: 1337, + data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), + nonce: 10_u128.into(), + }; + + let account_post_state = AccountPostState::new_claimed(account.clone(), Claim::Authorized); + + assert_eq!(account, account_post_state.account); + assert_eq!(account_post_state.required_claim(), Some(Claim::Authorized)); +} + +#[test] +fn post_state_new_without_claim_constructor() { + let account = Account { + program_owner: [1, 2, 3, 4, 5, 6, 7, 8], + balance: 1337, + data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), + nonce: 10_u128.into(), + }; + + let account_post_state = AccountPostState::new(account.clone()); + + assert_eq!(account, account_post_state.account); + assert!(account_post_state.required_claim().is_none()); +} + +#[test] +fn post_state_account_getter() { + let mut account = Account { + program_owner: [1, 2, 3, 4, 5, 6, 7, 8], + balance: 1337, + data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(), + nonce: 10_u128.into(), + }; + + let mut account_post_state = AccountPostState::new(account.clone()); + + assert_eq!(account_post_state.account(), &account); + assert_eq!(account_post_state.account_mut(), &mut account); +} + +// ---- AccountId::for_private_pda tests ---- + +/// Pins `AccountId::for_private_pda` against a hardcoded expected output for a specific +/// `(program_id, seed, npk, identifier)` tuple. Any change to `PRIVATE_PDA_PREFIX`, byte +/// ordering, or the underlying hash breaks this test. +#[test] +fn for_private_pda_matches_pinned_value() { + let program_id: ProgramId = [1; 8]; + let seed = PdaSeed::new([2; 32]); + let npk = NullifierPublicKey([3; 32]); + let identifier: Identifier = u128::MAX; + let expected = AccountId::new([ + 59, 239, 182, 97, 14, 220, 96, 115, 238, 133, 143, 33, 234, 82, 237, 255, 148, 110, 54, + 124, 98, 159, 245, 101, 146, 182, 150, 54, 37, 62, 25, 17, + ]); + assert_eq!( + AccountId::for_private_pda(&program_id, &seed, &npk, identifier), + expected + ); +} + +/// Two groups with different viewing keys at the same (program, seed) get different addresses. +#[test] +fn for_private_pda_differs_for_different_npk() { + let program_id: ProgramId = [1; 8]; + let seed = PdaSeed::new([2; 32]); + let npk_a = NullifierPublicKey([3; 32]); + let npk_b = NullifierPublicKey([4; 32]); + assert_ne!( + AccountId::for_private_pda(&program_id, &seed, &npk_a, u128::MAX), + AccountId::for_private_pda(&program_id, &seed, &npk_b, u128::MAX), + ); +} + +/// Different seeds produce different addresses, even with the same program and npk. +#[test] +fn for_private_pda_differs_for_different_seed() { + let program_id: ProgramId = [1; 8]; + let seed_a = PdaSeed::new([2; 32]); + let seed_b = PdaSeed::new([5; 32]); + let npk = NullifierPublicKey([3; 32]); + assert_ne!( + AccountId::for_private_pda(&program_id, &seed_a, &npk, u128::MAX), + AccountId::for_private_pda(&program_id, &seed_b, &npk, u128::MAX), + ); +} + +/// Different programs produce different addresses, even with the same seed and npk. +#[test] +fn for_private_pda_differs_for_different_program_id() { + let program_id_a: ProgramId = [1; 8]; + let program_id_b: ProgramId = [9; 8]; + let seed = PdaSeed::new([2; 32]); + let npk = NullifierPublicKey([3; 32]); + assert_ne!( + AccountId::for_private_pda(&program_id_a, &seed, &npk, u128::MAX), + AccountId::for_private_pda(&program_id_b, &seed, &npk, u128::MAX), + ); +} + +/// Different identifiers produce different addresses for the same `(program_id, seed, npk)`, +/// confirming that each `(program_id, seed, npk)` tuple controls a family of 2^128 addresses. +#[test] +fn for_private_pda_differs_for_different_identifier() { + let program_id: ProgramId = [1; 8]; + let seed = PdaSeed::new([2; 32]); + let npk = NullifierPublicKey([3; 32]); + assert_ne!( + AccountId::for_private_pda(&program_id, &seed, &npk, 0), + AccountId::for_private_pda(&program_id, &seed, &npk, 1), + ); + assert_ne!( + AccountId::for_private_pda(&program_id, &seed, &npk, 0), + AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX), + ); +} + +/// A private PDA at the same (program, seed) has a different address than a public PDA, +/// because the private formula uses a different prefix and includes npk. +#[test] +fn for_private_pda_differs_from_public_pda() { + let program_id: ProgramId = [1; 8]; + let seed = PdaSeed::new([2; 32]); + let npk = NullifierPublicKey([3; 32]); + let private_id = AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX); + let public_id = AccountId::for_public_pda(&program_id, &seed); + assert_ne!(private_id, public_id); +} + +#[cfg(feature = "host")] +#[test] +fn private_account_kind_header_round_trips() { + let regular = PrivateAccountKind::Regular(42); + let pda = PrivateAccountKind::Pda { + program_id: [1_u32; 8], + seed: PdaSeed::new([2_u8; 32]), + identifier: u128::MAX, + }; + assert_eq!( + PrivateAccountKind::from_header_bytes(®ular.to_header_bytes()), + Some(regular) + ); + assert_eq!( + PrivateAccountKind::from_header_bytes(&pda.to_header_bytes()), + Some(pda) + ); +} + +#[cfg(feature = "host")] +#[test] +fn private_account_kind_unknown_discriminant_returns_none() { + let mut bytes = [0_u8; PrivateAccountKind::HEADER_LEN]; + bytes[0] = 0xFF; + assert_eq!(PrivateAccountKind::from_header_bytes(&bytes), None); +} + +#[test] +fn for_private_account_dispatches_correctly() { + let program_id: ProgramId = [1; 8]; + let seed = PdaSeed::new([2; 32]); + let npk = NullifierPublicKey([3; 32]); + let identifier: Identifier = 77; + + assert_eq!( + AccountId::for_private_account(&npk, &PrivateAccountKind::Regular(identifier)), + AccountId::for_regular_private_account(&npk, identifier), + ); + assert_eq!( + AccountId::for_private_account( + &npk, + &PrivateAccountKind::Pda { + program_id, + seed, + identifier + } + ), + AccountId::for_private_pda(&program_id, &seed, &npk, identifier), + ); +} + +#[test] +fn compute_public_authorized_pdas_with_seeds() { + let caller: ProgramId = [1; 8]; + let seed = PdaSeed::new([2; 32]); + let result = compute_public_authorized_pdas(Some(caller), &[seed]); + let expected = AccountId::for_public_pda(&caller, &seed); + assert!(result.contains(&expected)); + assert_eq!(result.len(), 1); +} + +/// With no caller (top-level call), the result is always empty. +#[test] +fn compute_public_authorized_pdas_no_caller_returns_empty() { + let seed = PdaSeed::new([2; 32]); + let result = compute_public_authorized_pdas(None, &[seed]); + assert!(result.is_empty()); +} diff --git a/lee/state_machine/src/merkle_tree/mod.rs b/lee/state_machine/src/merkle_tree/mod.rs index e439d092..ee8106c5 100644 --- a/lee/state_machine/src/merkle_tree/mod.rs +++ b/lee/state_machine/src/merkle_tree/mod.rs @@ -164,398 +164,4 @@ const fn prev_power_of_two(x: usize) -> usize { } #[cfg(test)] -mod tests { - use hex_literal::hex; - - use super::*; - - impl MerkleTree { - pub fn new(values: &[Value]) -> Self { - let mut this = Self::with_capacity(values.len()); - for value in values.iter().copied() { - this.insert(value); - } - this - } - } - - #[test] - fn empty_merkle_tree() { - let tree = MerkleTree::with_capacity(4); - let expected_root = - hex!("0000000000000000000000000000000000000000000000000000000000000000"); - assert_eq!(tree.root(), expected_root); - assert_eq!(tree.capacity, 4); - assert_eq!(tree.length, 0); - } - - #[test] - fn merkle_tree_0() { - let values = [[0; 32]]; - let tree = MerkleTree::new(&values); - assert_eq!(tree.root(), hash_value(&[0; 32])); - assert_eq!(tree.capacity, 1); - assert_eq!(tree.length, 1); - } - - #[test] - fn merkle_tree_1() { - let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; - let tree = MerkleTree::new(&values); - let expected_root = - hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"); - assert_eq!(tree.root(), expected_root); - assert_eq!(tree.capacity, 4); - assert_eq!(tree.length, 4); - } - - #[test] - fn merkle_tree_2() { - let values = [[1; 32], [2; 32], [3; 32], [0; 32]]; - let tree = MerkleTree::new(&values); - let expected_root = - hex!("c9bbb83096df85157a146e7d770455a98412dee0633187ee86fee6c8a45b831a"); - assert_eq!(tree.root(), expected_root); - assert_eq!(tree.capacity, 4); - assert_eq!(tree.length, 4); - } - - #[test] - fn merkle_tree_3() { - let values = [[1; 32], [2; 32], [3; 32]]; - let tree = MerkleTree::new(&values); - let expected_root = - hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568"); - assert_eq!(tree.root(), expected_root); - assert_eq!(tree.capacity, 4); - assert_eq!(tree.length, 3); - } - - #[test] - fn merkle_tree_4() { - let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]]; - let tree = MerkleTree::new(&values); - let expected_root = - hex!("ef418aed5aa20702d4d94c92da79a4012f2e36f1008bfdb3cd1e38749dca2499"); - - assert_eq!(tree.root(), expected_root); - assert_eq!(tree.capacity, 8); - assert_eq!(tree.length, 5); - } - - #[test] - fn merkle_tree_5() { - let values = [ - [11; 32], [12; 32], [12; 32], [13; 32], [14; 32], [15; 32], [15; 32], [13; 32], - [13; 32], [15; 32], [11; 32], - ]; - let tree = MerkleTree::new(&values); - let expected_root = - hex!("3f72d2ff55921a86c48e5988ec3e19ee9d0d5aa3e23197842970a903508ed767"); - assert_eq!(tree.root(), expected_root); - assert_eq!(tree.capacity, 16); - assert_eq!(tree.length, 11); - } - - #[test] - fn merkle_tree_6() { - let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; - let tree = MerkleTree::new(&values); - let expected_root = - hex!("069cb8259a06fe6edb3fa7ff7933a6dd7dca6fca299314379794a688926c3792"); - assert_eq!(tree.root(), expected_root); - } - - #[test] - fn with_capacity_4() { - let tree = MerkleTree::with_capacity(4); - - assert_eq!(tree.length, 0); - assert_eq!(tree.nodes.len(), 7); - for i in 3..7 { - assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0], "{i}"); - } - for i in 1..3 { - assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1], "{i}"); - } - assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[2]); - } - - #[test] - fn with_capacity_5() { - let tree = MerkleTree::with_capacity(5); - - assert_eq!(tree.length, 0); - assert_eq!(tree.nodes.len(), 15); - for i in 7..15 { - assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0]); - } - for i in 3..7 { - assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1]); - } - for i in 1..3 { - assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[2]); - } - assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[3]); - } - - #[test] - fn with_capacity_6() { - let mut tree = MerkleTree::with_capacity(100); - - let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; - - let expected_root = - hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"); - - assert_eq!(0, tree.insert(values[0])); - assert_eq!(1, tree.insert(values[1])); - assert_eq!(2, tree.insert(values[2])); - assert_eq!(3, tree.insert(values[3])); - - assert_eq!(tree.root(), expected_root); - } - - #[test] - fn with_capacity_7() { - let mut tree = MerkleTree::with_capacity(599); - - let values = [[1; 32], [2; 32], [3; 32]]; - - let expected_root = - hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568"); - - assert_eq!(0, tree.insert(values[0])); - assert_eq!(1, tree.insert(values[1])); - assert_eq!(2, tree.insert(values[2])); - - assert_eq!(tree.root(), expected_root); - } - - #[test] - fn with_capacity_8() { - let mut tree = MerkleTree::with_capacity(1); - - let values = [[1; 32], [2; 32], [3; 32]]; - - let expected_root = - hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568"); - - assert_eq!(0, tree.insert(values[0])); - assert_eq!(1, tree.insert(values[1])); - assert_eq!(2, tree.insert(values[2])); - - assert_eq!(tree.root(), expected_root); - } - - #[test] - fn insert_value_1() { - let mut tree = MerkleTree::with_capacity(1); - - let values = [[1; 32], [2; 32], [3; 32]]; - let expected_tree = MerkleTree::new(&values); - - assert_eq!(0, tree.insert(values[0])); - assert_eq!(1, tree.insert(values[1])); - assert_eq!(2, tree.insert(values[2])); - - assert_eq!(expected_tree, tree); - } - - #[test] - fn insert_value_2() { - let mut tree = MerkleTree::with_capacity(1); - - let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; - let expected_tree = MerkleTree::new(&values); - - assert_eq!(0, tree.insert(values[0])); - assert_eq!(1, tree.insert(values[1])); - assert_eq!(2, tree.insert(values[2])); - assert_eq!(3, tree.insert(values[3])); - - assert_eq!(expected_tree, tree); - } - - #[test] - fn insert_value_3() { - let mut tree = MerkleTree::with_capacity(1); - - let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]]; - let expected_tree = MerkleTree::new(&values); - - tree.insert(values[0]); - tree.insert(values[1]); - tree.insert(values[2]); - tree.insert(values[3]); - tree.insert(values[4]); - - assert_eq!(expected_tree, tree); - } - - // Reference implementation - fn verify_authentication_path(value: &Value, index: usize, path: &[Node], root: &Node) -> bool { - let mut result = hash_value(value); - let mut level_index = index; - for node in path { - let is_left_child = level_index & 1 == 0; - if is_left_child { - result = hash_two(&result, node); - } else { - result = hash_two(node, &result); - } - level_index >>= 1; - } - &result == root - } - - #[test] - fn authentication_path_1() { - let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; - let tree = MerkleTree::new(&values); - let expected_authentication_path = vec![ - hex!("9f4fb68f3e1dac82202f9aa581ce0bbf1f765df0e9ac3c8c57e20f685abab8ed"), - hex!("50a27d4746f357cb700cbe9d4883b77fb64f0128828a3489dc6a6f21ddbf2414"), - ]; - - let authentication_path = tree.get_authentication_path_for(2).unwrap(); - assert_eq!(authentication_path, expected_authentication_path); - } - - #[test] - fn authentication_path_2() { - let values = [[1; 32], [2; 32], [3; 32]]; - let tree = MerkleTree::new(&values); - let expected_authentication_path = vec![ - hex!("75877bb41d393b5fb8455ce60ecd8dda001d06316496b14dfa7f895656eeca4a"), - hex!("a41b855d2db4de9052cd7be5ec67d6586629cb9f6e3246a4afa5ba313f07a9c5"), - ]; - - let authentication_path = tree.get_authentication_path_for(0).unwrap(); - assert_eq!(authentication_path, expected_authentication_path); - } - - #[test] - fn authentication_path_3() { - let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; - let tree = MerkleTree::new(&values); - let expected_authentication_path = vec![ - hex!("0000000000000000000000000000000000000000000000000000000000000000"), - hex!("f5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b"), - hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"), - ]; - - let authentication_path = tree.get_authentication_path_for(4).unwrap(); - assert_eq!(authentication_path, expected_authentication_path); - } - - #[test] - fn authentication_path_4() { - let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; - let tree = MerkleTree::new(&values); - assert!(tree.get_authentication_path_for(5).is_none()); - } - - #[test] - fn authentication_path_5() { - let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; - let tree = MerkleTree::new(&values); - let index = 4; - let value = values[index]; - let path = tree.get_authentication_path_for(index).unwrap(); - assert!(verify_authentication_path( - &value, - index, - &path, - &tree.root() - )); - } - - #[test] - fn tree_with_63_insertions() { - let values = [ - hex!("cd00acab0f45736e6c6311f1953becc0b69a062e7c2a7310875d28bdf9ef9c5b"), - hex!("0df5a6afbcc7bf126caf7084acfc593593ab512e6ca433c61c1a922be40a04ea"), - hex!("23c1258620266c7bedb6d1ee32f6da9413e4010ace975239dccb34e727e07c40"), - hex!("f33ccc3a11476b0ef62326ca5ec292056759b05e6a28023d2d1ce66165611353"), - hex!("77f914ab016b8049f6bea7704000e413a393865918a3824f9285c3db0aacff23"), - hex!("910a1c23188e54d57fd167ddb0f8bf68c6b70ed9ec76ef56c4b7f2632f82ca7f"), - hex!("047ee85526197d1e7403a559cf6d2f22c1926c8ad59481a2e2f1b697af45e40b"), - hex!("9d355cf89fb382ae34bf80566b28489278d10f2cebb5b0ea42fab1bac5adae0c"), - hex!("604018b95232596b2685a9bc737b6cccb53b10e483d2d9a2f4a755410b02a188"), - hex!("a16708ef7b6bf1796063addaf57d6a566b6f87b0bbe42af43a4590d05f1684cb"), - hex!("820f2dfa271cd2fd41e1452406d5dad552c85c1223c45d45dbd7446759fdc6b8"), - hex!("680b6912d7e219f8805d4d28adb4428dd78fea0dc1b8cdb2412645c4b1962c88"), - hex!("14d5471ce6c45506753982b17cac5790ac7bc29e6f388f31052d7dfd62b294e5"), - hex!("8b364200172b777d4aa16d2098b5eb98ac3dd4a1b9597e5c2bf6f6930031f230"), - hex!("9bb45b910711874339dda8a21a9aad73822286f5e52d7d3de0ed78dfbba329a5"), - hex!("d6806d5df5cb25ce5d531042f09b3cb34fb9e47c61182b63cccd9d44392f6027"), - hex!("b8cfa90ebc8fd09c04682d93a08fddd3e8e57715174dcc92451edd191264a58b"), - hex!("3463c7f81d00f809b3dfa83195447c927fb4045b3913dac6f45bee6c4010d7ed"), - hex!("1d6ad7f7d677905feb506c58f4b404a79370ebc567296abea3a368b61d5a8239"), - hex!("a58085ecf00963cb22da23c901b9b3ddc56462bb96ff03c923d67708e10dd29c"), - hex!("c3319f4a65fb5bbb8447137b0972c03cbd84ebf7d9da194e0fcbd68c2d4d5bdb"), - hex!("4aa31e90e0090faf3648d05e5d5499df2c78ebed4d6e6c23d8147de5d67dae73"), - hex!("9f33b1d2c8bc7bd265336de1033ede6344bc41260313bdcb43f1108b83b9be92"), - hex!("6500d4ad93d41c16ec81eaa5e70f173194aabe5c1072ac263b5727296f5b7cac"), - hex!("3584f5d260003669fad98786e13171376b0f19410cb232ce65606cbff79e6768"), - hex!("c8410946ebf56f13141c894a34ced85a5230088af70dcea581e44f52847830ac"), - hex!("71dd90281cdebb70422f2d04ae446d5d2d5ea64b803c16128d37e3fcd5d1a4cc"), - hex!("c05acf8d77ab4d659a538bd35af590864a7ad9c055ff5d6cda9d5aecfccecba3"), - hex!("f1df98822ea084cce9021aa9cc81b1746cd1e84a75690da63e10fd877633ed77"), - hex!("2ca822bc8f67bceb0a71a0d06fea7349036ef3e5ec21795a851e4182bd35ce01"), - hex!("7fd2179abc3bcf89b4d8092988ba8c23952b3bbd3d7caea6b5ea0c13cf19f68b"), - hex!("91b6ad516e017f6aa5a2e95776538bd3a3e933c1b1d32bb5e0f00a9db63c9c24"), - hex!("cd31a8b5eef5ca0be5ef1cb261d0bf0a74d774a3152bb99739cfd296a1d0b85e"), - hex!("3fb16f48b2bf93f3815979e6638f975d7f935088ec37db0be0f07965fbc78339"), - hex!("c60c61b99bf486af5f4bf780a69860dafcd35c1474306a8575666fb5449bcec0"), - hex!("8048d0d7e14091251f3f6c6b10bf6b5880a014b513f9f8c2395501dbffa6192a"), - hex!("778b5af10b9dbe80b60a8e4f0bb91caf4476bcb812801099760754ae623fbd84"), - hex!("d3ac25467920a4e08998b7a3226b8b54bfe66ac58cfedc71f15b2402fee0054a"), - hex!("029aa94598fae2961a0d43937b8a9a3138bcfeae99a7cb15f77fac7c506f8432"), - hex!("2eee5ef52fe669cb6882a68c893abdc1262dcf4424e4ba7a479da7cf1c10171d"), - hex!("de3fb3d070e3a90f0eed8b5e65088a8dc0e4e3c342b9c0bf33bab714eae5dfec"), - hex!("14d40177e833ab45bbfdc5f2b11fba7efaebb3f69facc554f24b549a2efe8538"), - hex!("5734355069702448774fb2df95f1d562e1b9fe1514aeb6b922554ee9d2d01068"), - hex!("8a273d49ac110343cec2cf3359d16eb2906b446bd9ec9833e2a640cebc8d5155"), - hex!("e3fa984dd3cbeb9a7e827ed32d3d4e6a6ba643a55d82be97d9ddb06ee809fa3e"), - hex!("90b1d5a364e17c8b7965396b06ec6e13749b5fc16500731518ad8fc30ae33e77"), - hex!("7517376541b2e8ec83cbab04522b54a26610908a9872feb663451385aea58eb1"), - hex!("5cba2e4cf7448e526d161133c4b2ea7c919ac4813a7308612595f46f11dea6cd"), - hex!("c721911b300bec0691c8a2dfaabfef1d66b7b6258918914d3c3ad690729f05b7"), - hex!("d0d0a70d8ae0d27806fa0b711c507290c260a89cbca0436d339d1dccdd087d62"), - hex!("2a625c28ea763c5e82dd0a93ecfca7ec371ccbb363cd42be359c2c875f58009d"), - hex!("174ef0119932ed890397d9f3837dd85f9100558b6fc9085d4af947ae8cf74bbc"), - hex!("b497bc267151e8efa3c6daa461e6804b01a3f05f44f1f4d5b41d5f0d3f5219b1"), - hex!("e987e91f5734630ddd7e6b58733b4fcdbc316ee9e8cac0e94c36c91cf58e59cc"), - hex!("55019ad8bbe656c51eb042190c1c8da53f42baf43fd2350ebea38fc7cca2fae3"), - hex!("c45a638edd18a6d9f5ad20b870c81b8626459bcb22dae7d58add7a6b6c6a84a8"), - hex!("d42d3a5fb2ad50b2027fe5a36d59dd71e49a63e4b1b299073c96bbf7ba5d68a1"), - hex!("9599e561054bcd3f647eb018ab0b069d3176497d42be9c4466551cbb959be47c"), - hex!("42f33b23775327ff71aea6569548255f3cc9929da73373cc9bb1743d417f7cda"), - hex!("ab24294f44fc6fdbeb96e0f6e93c4f6d97d035b73b9a337c353e18c6d0603bdd"), - hex!("33954ec63520334f99b640a2982ac966b68c363fed383d621a1ab573934f1d33"), - hex!("5e2a1f7df963d1fd8f50a285387cfbb5df581426619b325563e20bf7886c62b7"), - hex!("13ffde471d4e27c473254e766fd1328ad80c42cab4d4955cffeae43d866f86e5"), - ]; - - let expected_root = - hex!("1cf9b214217d7823f9de51b8f6cb34d0a99436a3a1bb762f90b815672a6afcc0"); - - let mut tree_less_capacity = MerkleTree::with_capacity(1); - let mut tree_exact_capacity = MerkleTree::with_capacity(64); - let mut tree_more_capacity = MerkleTree::with_capacity(128); - - for value in &values { - tree_less_capacity.insert(*value); - tree_exact_capacity.insert(*value); - tree_more_capacity.insert(*value); - } - - assert_eq!(tree_more_capacity.root(), expected_root); - assert_eq!(tree_less_capacity.root(), expected_root); - assert_eq!(tree_exact_capacity.root(), expected_root); - } -} - -// +mod tests; diff --git a/lee/state_machine/src/merkle_tree/tests.rs b/lee/state_machine/src/merkle_tree/tests.rs new file mode 100644 index 00000000..756fd45f --- /dev/null +++ b/lee/state_machine/src/merkle_tree/tests.rs @@ -0,0 +1,380 @@ +use hex_literal::hex; + +use super::*; + +impl MerkleTree { + pub fn new(values: &[Value]) -> Self { + let mut this = Self::with_capacity(values.len()); + for value in values.iter().copied() { + this.insert(value); + } + this + } +} + +#[test] +fn empty_merkle_tree() { + let tree = MerkleTree::with_capacity(4); + let expected_root = hex!("0000000000000000000000000000000000000000000000000000000000000000"); + assert_eq!(tree.root(), expected_root); + assert_eq!(tree.capacity, 4); + assert_eq!(tree.length, 0); +} + +#[test] +fn merkle_tree_0() { + let values = [[0; 32]]; + let tree = MerkleTree::new(&values); + assert_eq!(tree.root(), hash_value(&[0; 32])); + assert_eq!(tree.capacity, 1); + assert_eq!(tree.length, 1); +} + +#[test] +fn merkle_tree_1() { + let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; + let tree = MerkleTree::new(&values); + let expected_root = hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"); + assert_eq!(tree.root(), expected_root); + assert_eq!(tree.capacity, 4); + assert_eq!(tree.length, 4); +} + +#[test] +fn merkle_tree_2() { + let values = [[1; 32], [2; 32], [3; 32], [0; 32]]; + let tree = MerkleTree::new(&values); + let expected_root = hex!("c9bbb83096df85157a146e7d770455a98412dee0633187ee86fee6c8a45b831a"); + assert_eq!(tree.root(), expected_root); + assert_eq!(tree.capacity, 4); + assert_eq!(tree.length, 4); +} + +#[test] +fn merkle_tree_3() { + let values = [[1; 32], [2; 32], [3; 32]]; + let tree = MerkleTree::new(&values); + let expected_root = hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568"); + assert_eq!(tree.root(), expected_root); + assert_eq!(tree.capacity, 4); + assert_eq!(tree.length, 3); +} + +#[test] +fn merkle_tree_4() { + let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]]; + let tree = MerkleTree::new(&values); + let expected_root = hex!("ef418aed5aa20702d4d94c92da79a4012f2e36f1008bfdb3cd1e38749dca2499"); + + assert_eq!(tree.root(), expected_root); + assert_eq!(tree.capacity, 8); + assert_eq!(tree.length, 5); +} + +#[test] +fn merkle_tree_5() { + let values = [ + [11; 32], [12; 32], [12; 32], [13; 32], [14; 32], [15; 32], [15; 32], [13; 32], [13; 32], + [15; 32], [11; 32], + ]; + let tree = MerkleTree::new(&values); + let expected_root = hex!("3f72d2ff55921a86c48e5988ec3e19ee9d0d5aa3e23197842970a903508ed767"); + assert_eq!(tree.root(), expected_root); + assert_eq!(tree.capacity, 16); + assert_eq!(tree.length, 11); +} + +#[test] +fn merkle_tree_6() { + let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; + let tree = MerkleTree::new(&values); + let expected_root = hex!("069cb8259a06fe6edb3fa7ff7933a6dd7dca6fca299314379794a688926c3792"); + assert_eq!(tree.root(), expected_root); +} + +#[test] +fn with_capacity_4() { + let tree = MerkleTree::with_capacity(4); + + assert_eq!(tree.length, 0); + assert_eq!(tree.nodes.len(), 7); + for i in 3..7 { + assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0], "{i}"); + } + for i in 1..3 { + assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1], "{i}"); + } + assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[2]); +} + +#[test] +fn with_capacity_5() { + let tree = MerkleTree::with_capacity(5); + + assert_eq!(tree.length, 0); + assert_eq!(tree.nodes.len(), 15); + for i in 7..15 { + assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0]); + } + for i in 3..7 { + assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1]); + } + for i in 1..3 { + assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[2]); + } + assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[3]); +} + +#[test] +fn with_capacity_6() { + let mut tree = MerkleTree::with_capacity(100); + + let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; + + let expected_root = hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"); + + assert_eq!(0, tree.insert(values[0])); + assert_eq!(1, tree.insert(values[1])); + assert_eq!(2, tree.insert(values[2])); + assert_eq!(3, tree.insert(values[3])); + + assert_eq!(tree.root(), expected_root); +} + +#[test] +fn with_capacity_7() { + let mut tree = MerkleTree::with_capacity(599); + + let values = [[1; 32], [2; 32], [3; 32]]; + + let expected_root = hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568"); + + assert_eq!(0, tree.insert(values[0])); + assert_eq!(1, tree.insert(values[1])); + assert_eq!(2, tree.insert(values[2])); + + assert_eq!(tree.root(), expected_root); +} + +#[test] +fn with_capacity_8() { + let mut tree = MerkleTree::with_capacity(1); + + let values = [[1; 32], [2; 32], [3; 32]]; + + let expected_root = hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568"); + + assert_eq!(0, tree.insert(values[0])); + assert_eq!(1, tree.insert(values[1])); + assert_eq!(2, tree.insert(values[2])); + + assert_eq!(tree.root(), expected_root); +} + +#[test] +fn insert_value_1() { + let mut tree = MerkleTree::with_capacity(1); + + let values = [[1; 32], [2; 32], [3; 32]]; + let expected_tree = MerkleTree::new(&values); + + assert_eq!(0, tree.insert(values[0])); + assert_eq!(1, tree.insert(values[1])); + assert_eq!(2, tree.insert(values[2])); + + assert_eq!(expected_tree, tree); +} + +#[test] +fn insert_value_2() { + let mut tree = MerkleTree::with_capacity(1); + + let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; + let expected_tree = MerkleTree::new(&values); + + assert_eq!(0, tree.insert(values[0])); + assert_eq!(1, tree.insert(values[1])); + assert_eq!(2, tree.insert(values[2])); + assert_eq!(3, tree.insert(values[3])); + + assert_eq!(expected_tree, tree); +} + +#[test] +fn insert_value_3() { + let mut tree = MerkleTree::with_capacity(1); + + let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]]; + let expected_tree = MerkleTree::new(&values); + + tree.insert(values[0]); + tree.insert(values[1]); + tree.insert(values[2]); + tree.insert(values[3]); + tree.insert(values[4]); + + assert_eq!(expected_tree, tree); +} + +// Reference implementation +fn verify_authentication_path(value: &Value, index: usize, path: &[Node], root: &Node) -> bool { + let mut result = hash_value(value); + let mut level_index = index; + for node in path { + let is_left_child = level_index & 1 == 0; + if is_left_child { + result = hash_two(&result, node); + } else { + result = hash_two(node, &result); + } + level_index >>= 1; + } + &result == root +} + +#[test] +fn authentication_path_1() { + let values = [[1; 32], [2; 32], [3; 32], [4; 32]]; + let tree = MerkleTree::new(&values); + let expected_authentication_path = vec![ + hex!("9f4fb68f3e1dac82202f9aa581ce0bbf1f765df0e9ac3c8c57e20f685abab8ed"), + hex!("50a27d4746f357cb700cbe9d4883b77fb64f0128828a3489dc6a6f21ddbf2414"), + ]; + + let authentication_path = tree.get_authentication_path_for(2).unwrap(); + assert_eq!(authentication_path, expected_authentication_path); +} + +#[test] +fn authentication_path_2() { + let values = [[1; 32], [2; 32], [3; 32]]; + let tree = MerkleTree::new(&values); + let expected_authentication_path = vec![ + hex!("75877bb41d393b5fb8455ce60ecd8dda001d06316496b14dfa7f895656eeca4a"), + hex!("a41b855d2db4de9052cd7be5ec67d6586629cb9f6e3246a4afa5ba313f07a9c5"), + ]; + + let authentication_path = tree.get_authentication_path_for(0).unwrap(); + assert_eq!(authentication_path, expected_authentication_path); +} + +#[test] +fn authentication_path_3() { + let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; + let tree = MerkleTree::new(&values); + let expected_authentication_path = vec![ + hex!("0000000000000000000000000000000000000000000000000000000000000000"), + hex!("f5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b"), + hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"), + ]; + + let authentication_path = tree.get_authentication_path_for(4).unwrap(); + assert_eq!(authentication_path, expected_authentication_path); +} + +#[test] +fn authentication_path_4() { + let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; + let tree = MerkleTree::new(&values); + assert!(tree.get_authentication_path_for(5).is_none()); +} + +#[test] +fn authentication_path_5() { + let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; + let tree = MerkleTree::new(&values); + let index = 4; + let value = values[index]; + let path = tree.get_authentication_path_for(index).unwrap(); + assert!(verify_authentication_path( + &value, + index, + &path, + &tree.root() + )); +} + +#[test] +fn tree_with_63_insertions() { + let values = [ + hex!("cd00acab0f45736e6c6311f1953becc0b69a062e7c2a7310875d28bdf9ef9c5b"), + hex!("0df5a6afbcc7bf126caf7084acfc593593ab512e6ca433c61c1a922be40a04ea"), + hex!("23c1258620266c7bedb6d1ee32f6da9413e4010ace975239dccb34e727e07c40"), + hex!("f33ccc3a11476b0ef62326ca5ec292056759b05e6a28023d2d1ce66165611353"), + hex!("77f914ab016b8049f6bea7704000e413a393865918a3824f9285c3db0aacff23"), + hex!("910a1c23188e54d57fd167ddb0f8bf68c6b70ed9ec76ef56c4b7f2632f82ca7f"), + hex!("047ee85526197d1e7403a559cf6d2f22c1926c8ad59481a2e2f1b697af45e40b"), + hex!("9d355cf89fb382ae34bf80566b28489278d10f2cebb5b0ea42fab1bac5adae0c"), + hex!("604018b95232596b2685a9bc737b6cccb53b10e483d2d9a2f4a755410b02a188"), + hex!("a16708ef7b6bf1796063addaf57d6a566b6f87b0bbe42af43a4590d05f1684cb"), + hex!("820f2dfa271cd2fd41e1452406d5dad552c85c1223c45d45dbd7446759fdc6b8"), + hex!("680b6912d7e219f8805d4d28adb4428dd78fea0dc1b8cdb2412645c4b1962c88"), + hex!("14d5471ce6c45506753982b17cac5790ac7bc29e6f388f31052d7dfd62b294e5"), + hex!("8b364200172b777d4aa16d2098b5eb98ac3dd4a1b9597e5c2bf6f6930031f230"), + hex!("9bb45b910711874339dda8a21a9aad73822286f5e52d7d3de0ed78dfbba329a5"), + hex!("d6806d5df5cb25ce5d531042f09b3cb34fb9e47c61182b63cccd9d44392f6027"), + hex!("b8cfa90ebc8fd09c04682d93a08fddd3e8e57715174dcc92451edd191264a58b"), + hex!("3463c7f81d00f809b3dfa83195447c927fb4045b3913dac6f45bee6c4010d7ed"), + hex!("1d6ad7f7d677905feb506c58f4b404a79370ebc567296abea3a368b61d5a8239"), + hex!("a58085ecf00963cb22da23c901b9b3ddc56462bb96ff03c923d67708e10dd29c"), + hex!("c3319f4a65fb5bbb8447137b0972c03cbd84ebf7d9da194e0fcbd68c2d4d5bdb"), + hex!("4aa31e90e0090faf3648d05e5d5499df2c78ebed4d6e6c23d8147de5d67dae73"), + hex!("9f33b1d2c8bc7bd265336de1033ede6344bc41260313bdcb43f1108b83b9be92"), + hex!("6500d4ad93d41c16ec81eaa5e70f173194aabe5c1072ac263b5727296f5b7cac"), + hex!("3584f5d260003669fad98786e13171376b0f19410cb232ce65606cbff79e6768"), + hex!("c8410946ebf56f13141c894a34ced85a5230088af70dcea581e44f52847830ac"), + hex!("71dd90281cdebb70422f2d04ae446d5d2d5ea64b803c16128d37e3fcd5d1a4cc"), + hex!("c05acf8d77ab4d659a538bd35af590864a7ad9c055ff5d6cda9d5aecfccecba3"), + hex!("f1df98822ea084cce9021aa9cc81b1746cd1e84a75690da63e10fd877633ed77"), + hex!("2ca822bc8f67bceb0a71a0d06fea7349036ef3e5ec21795a851e4182bd35ce01"), + hex!("7fd2179abc3bcf89b4d8092988ba8c23952b3bbd3d7caea6b5ea0c13cf19f68b"), + hex!("91b6ad516e017f6aa5a2e95776538bd3a3e933c1b1d32bb5e0f00a9db63c9c24"), + hex!("cd31a8b5eef5ca0be5ef1cb261d0bf0a74d774a3152bb99739cfd296a1d0b85e"), + hex!("3fb16f48b2bf93f3815979e6638f975d7f935088ec37db0be0f07965fbc78339"), + hex!("c60c61b99bf486af5f4bf780a69860dafcd35c1474306a8575666fb5449bcec0"), + hex!("8048d0d7e14091251f3f6c6b10bf6b5880a014b513f9f8c2395501dbffa6192a"), + hex!("778b5af10b9dbe80b60a8e4f0bb91caf4476bcb812801099760754ae623fbd84"), + hex!("d3ac25467920a4e08998b7a3226b8b54bfe66ac58cfedc71f15b2402fee0054a"), + hex!("029aa94598fae2961a0d43937b8a9a3138bcfeae99a7cb15f77fac7c506f8432"), + hex!("2eee5ef52fe669cb6882a68c893abdc1262dcf4424e4ba7a479da7cf1c10171d"), + hex!("de3fb3d070e3a90f0eed8b5e65088a8dc0e4e3c342b9c0bf33bab714eae5dfec"), + hex!("14d40177e833ab45bbfdc5f2b11fba7efaebb3f69facc554f24b549a2efe8538"), + hex!("5734355069702448774fb2df95f1d562e1b9fe1514aeb6b922554ee9d2d01068"), + hex!("8a273d49ac110343cec2cf3359d16eb2906b446bd9ec9833e2a640cebc8d5155"), + hex!("e3fa984dd3cbeb9a7e827ed32d3d4e6a6ba643a55d82be97d9ddb06ee809fa3e"), + hex!("90b1d5a364e17c8b7965396b06ec6e13749b5fc16500731518ad8fc30ae33e77"), + hex!("7517376541b2e8ec83cbab04522b54a26610908a9872feb663451385aea58eb1"), + hex!("5cba2e4cf7448e526d161133c4b2ea7c919ac4813a7308612595f46f11dea6cd"), + hex!("c721911b300bec0691c8a2dfaabfef1d66b7b6258918914d3c3ad690729f05b7"), + hex!("d0d0a70d8ae0d27806fa0b711c507290c260a89cbca0436d339d1dccdd087d62"), + hex!("2a625c28ea763c5e82dd0a93ecfca7ec371ccbb363cd42be359c2c875f58009d"), + hex!("174ef0119932ed890397d9f3837dd85f9100558b6fc9085d4af947ae8cf74bbc"), + hex!("b497bc267151e8efa3c6daa461e6804b01a3f05f44f1f4d5b41d5f0d3f5219b1"), + hex!("e987e91f5734630ddd7e6b58733b4fcdbc316ee9e8cac0e94c36c91cf58e59cc"), + hex!("55019ad8bbe656c51eb042190c1c8da53f42baf43fd2350ebea38fc7cca2fae3"), + hex!("c45a638edd18a6d9f5ad20b870c81b8626459bcb22dae7d58add7a6b6c6a84a8"), + hex!("d42d3a5fb2ad50b2027fe5a36d59dd71e49a63e4b1b299073c96bbf7ba5d68a1"), + hex!("9599e561054bcd3f647eb018ab0b069d3176497d42be9c4466551cbb959be47c"), + hex!("42f33b23775327ff71aea6569548255f3cc9929da73373cc9bb1743d417f7cda"), + hex!("ab24294f44fc6fdbeb96e0f6e93c4f6d97d035b73b9a337c353e18c6d0603bdd"), + hex!("33954ec63520334f99b640a2982ac966b68c363fed383d621a1ab573934f1d33"), + hex!("5e2a1f7df963d1fd8f50a285387cfbb5df581426619b325563e20bf7886c62b7"), + hex!("13ffde471d4e27c473254e766fd1328ad80c42cab4d4955cffeae43d866f86e5"), + ]; + + let expected_root = hex!("1cf9b214217d7823f9de51b8f6cb34d0a99436a3a1bb762f90b815672a6afcc0"); + + let mut tree_less_capacity = MerkleTree::with_capacity(1); + let mut tree_exact_capacity = MerkleTree::with_capacity(64); + let mut tree_more_capacity = MerkleTree::with_capacity(128); + + for value in &values { + tree_less_capacity.insert(*value); + tree_exact_capacity.insert(*value); + tree_more_capacity.insert(*value); + } + + assert_eq!(tree_more_capacity.root(), expected_root); + assert_eq!(tree_less_capacity.root(), expected_root); + assert_eq!(tree_exact_capacity.root(), expected_root); +} diff --git a/lee/state_machine/src/privacy_preserving_transaction/circuit.rs b/lee/state_machine/src/privacy_preserving_transaction/circuit.rs deleted file mode 100644 index 489ee373..00000000 --- a/lee/state_machine/src/privacy_preserving_transaction/circuit.rs +++ /dev/null @@ -1,906 +0,0 @@ -use std::collections::{HashMap, VecDeque}; - -use borsh::{BorshDeserialize, BorshSerialize}; -use lee_core::{ - InputAccountIdentity, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, - account::AccountWithMetadata, - program::{ChainedCall, InstructionData, ProgramId, ProgramOutput}, -}; -use risc0_zkvm::{ExecutorEnv, InnerReceipt, ProverOpts, Receipt, default_prover}; - -use crate::{ - PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID, - error::{InvalidProgramBehaviorError, LeeError}, - program::Program, - state::MAX_NUMBER_CHAINED_CALLS, -}; - -/// Proof of the privacy preserving execution circuit. -#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] -pub struct Proof(pub(crate) Vec); - -impl Proof { - #[must_use] - pub fn into_inner(self) -> Vec { - self.0 - } - - #[must_use] - pub const fn from_inner(inner: Vec) -> Self { - Self(inner) - } - - pub(crate) fn is_valid_for(&self, circuit_output: &PrivacyPreservingCircuitOutput) -> bool { - let Ok(inner) = borsh::from_slice::(&self.0) else { - return false; - }; - let receipt = Receipt::new(inner, circuit_output.to_bytes()); - receipt.verify(PRIVACY_PRESERVING_CIRCUIT_ID).is_ok() - } -} - -#[derive(Clone)] -pub struct ProgramWithDependencies { - pub program: Program, - // TODO: avoid having a copy of the bytecode of each dependency. - pub dependencies: HashMap, -} - -impl ProgramWithDependencies { - #[must_use] - pub const fn new(program: Program, dependencies: HashMap) -> Self { - Self { - program, - dependencies, - } - } -} - -impl From for ProgramWithDependencies { - fn from(program: Program) -> Self { - Self::new(program, HashMap::new()) - } -} - -/// Generates a proof of the execution of a LEE program inside the privacy preserving execution -/// circuit. -pub fn execute_and_prove( - pre_states: Vec, - instruction_data: InstructionData, - account_identities: Vec, - program_with_dependencies: &ProgramWithDependencies, -) -> Result<(PrivacyPreservingCircuitOutput, Proof), LeeError> { - let ProgramWithDependencies { - program: initial_program, - dependencies, - } = program_with_dependencies; - let mut env_builder = ExecutorEnv::builder(); - let mut program_outputs = Vec::new(); - - let initial_call = ChainedCall { - program_id: initial_program.id(), - instruction_data, - pre_states, - pda_seeds: vec![], - }; - - let mut chained_calls = VecDeque::from_iter([(initial_call, initial_program, None)]); - let mut chain_calls_counter = 0; - while let Some((chained_call, program, caller_program_id)) = chained_calls.pop_front() { - if chain_calls_counter >= MAX_NUMBER_CHAINED_CALLS { - return Err(LeeError::MaxChainedCallsDepthExceeded); - } - - let inner_receipt = execute_and_prove_program( - program, - caller_program_id, - &chained_call.pre_states, - &chained_call.instruction_data, - )?; - - let program_output: ProgramOutput = inner_receipt - .journal - .decode() - .map_err(|e| LeeError::ProgramOutputDeserializationError(e.to_string()))?; - - // TODO: remove clone - program_outputs.push(program_output.clone()); - - // Prove circuit. - env_builder.add_assumption(inner_receipt); - - for new_call in program_output.chained_calls.into_iter().rev() { - let next_program = dependencies.get(&new_call.program_id).ok_or( - InvalidProgramBehaviorError::UndeclaredProgramDependency { - program_id: new_call.program_id, - }, - )?; - chained_calls.push_front((new_call, next_program, Some(chained_call.program_id))); - } - - chain_calls_counter = chain_calls_counter - .checked_add(1) - .expect("we check the max depth at the beginning of the loop"); - } - - let circuit_input = PrivacyPreservingCircuitInput { - program_outputs, - account_identities, - program_id: program_with_dependencies.program.id(), - }; - - env_builder.write(&circuit_input).unwrap(); - let env = env_builder.build().unwrap(); - let prover = default_prover(); - let opts = ProverOpts::succinct(); - let prove_info = prover - .prove_with_opts(env, PRIVACY_PRESERVING_CIRCUIT_ELF, &opts) - .map_err(|e| LeeError::CircuitProvingError(e.to_string()))?; - - let proof = Proof(borsh::to_vec(&prove_info.receipt.inner)?); - - let circuit_output: PrivacyPreservingCircuitOutput = prove_info - .receipt - .journal - .decode() - .map_err(|e| LeeError::CircuitOutputDeserializationError(e.to_string()))?; - - Ok((circuit_output, proof)) -} - -fn execute_and_prove_program( - program: &Program, - caller_program_id: Option, - pre_states: &[AccountWithMetadata], - instruction_data: &InstructionData, -) -> Result { - // Write inputs to the program - let mut env_builder = ExecutorEnv::builder(); - Program::write_inputs( - program.id(), - caller_program_id, - pre_states, - instruction_data, - &mut env_builder, - )?; - let env = env_builder.build().unwrap(); - - // Prove the program - let prover = default_prover(); - Ok(prover - .prove(env, program.elf()) - .map_err(|e| LeeError::ProgramProveFailed(e.to_string()))? - .receipt) -} - -#[cfg(test)] -mod tests { - #![expect(clippy::shadow_unrelated, reason = "We don't care about it in tests")] - - use lee_core::{ - Commitment, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme, - EphemeralPublicKey, Nullifier, PrivacyPreservingCircuitOutput, SharedSecretKey, - account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data}, - program::{PdaSeed, PrivateAccountKind}, - }; - - use super::*; - use crate::{ - error::LeeError, - privacy_preserving_transaction::circuit::execute_and_prove, - program::Program, - state::{ - CommitmentSet, - tests::{test_private_account_keys_1, test_private_account_keys_2}, - }, - }; - - fn decrypt_kind( - output: &PrivacyPreservingCircuitOutput, - ssk: &SharedSecretKey, - idx: usize, - ) -> PrivateAccountKind { - let (kind, _) = EncryptionScheme::decrypt( - &output.encrypted_private_post_states[idx].ciphertext, - ssk, - &output.new_commitments[idx], - u32::try_from(idx).expect("idx fits in u32"), - ) - .unwrap(); - kind - } - - #[test] - fn proof_inner_roundtrip() { - // `Proof::from_inner(b).into_inner()` must return exactly `b`. Catches - // mutations of `into_inner` returning `vec![]`, `vec![0]`, or `vec![1]`, - // and of `from_inner` discarding its argument. - let bytes = vec![0xDE_u8, 0xAD, 0xBE, 0xEF]; - assert_eq!(Proof::from_inner(bytes.clone()).into_inner(), bytes); - assert!(Proof::from_inner(vec![]).into_inner().is_empty()); - assert_eq!(Proof::from_inner(vec![0xFF]).into_inner(), vec![0xFF_u8]); - } - - #[test] - fn prove_privacy_preserving_execution_circuit_public_and_private_pre_accounts() { - let recipient_keys = test_private_account_keys_1(); - let program = crate::test_methods::simple_balance_transfer(); - let sender = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); - let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id); - - let balance_to_move: u128 = 37; - - let expected_sender_post = Account { - program_owner: program.id(), - balance: 100 - balance_to_move, - nonce: Nonce::default(), - data: Data::default(), - }; - - let expected_recipient_post = Account { - program_owner: program.id(), - balance: balance_to_move, - nonce: Nonce::private_account_nonce_init(&recipient_account_id), - data: Data::default(), - }; - - let expected_sender_pre = sender.clone(); - - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0; - - let (output, proof) = execute_and_prove( - vec![sender, recipient], - Program::serialize_instruction(balance_to_move).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: shared_secret, - identifier: 0, - }, - ], - &crate::test_methods::simple_balance_transfer().into(), - ) - .unwrap(); - - assert!(proof.is_valid_for(&output)); - - let [sender_pre] = output.public_pre_states.try_into().unwrap(); - let [sender_post] = output.public_post_states.try_into().unwrap(); - assert_eq!(sender_pre, expected_sender_pre); - assert_eq!(sender_post, expected_sender_post); - assert_eq!(output.new_commitments.len(), 1); - assert_eq!(output.new_nullifiers.len(), 1); - assert_eq!(output.encrypted_private_post_states.len(), 1); - - let (_identifier, recipient_post) = EncryptionScheme::decrypt( - &output.encrypted_private_post_states[0].ciphertext, - &shared_secret, - &output.new_commitments[0], - 0, - ) - .unwrap(); - assert_eq!(recipient_post, expected_recipient_post); - } - - #[test] - fn prove_privacy_preserving_execution_circuit_fully_private() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - - let sender_nonce = Nonce(0xdead_beef); - let sender_pre = AccountWithMetadata::new( - Account { - balance: 100, - nonce: sender_nonce, - program_owner: program.id(), - data: Data::default(), - }, - true, - AccountId::for_regular_private_account(&sender_keys.npk(), 0), - ); - let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); - let commitment_sender = Commitment::new(&sender_account_id, &sender_pre.account); - - let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); - let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id); - let balance_to_move: u128 = 37; - - let mut commitment_set = CommitmentSet::with_capacity(2); - commitment_set.extend(std::slice::from_ref(&commitment_sender)); - let expected_new_nullifiers = vec![ - ( - Nullifier::for_account_update(&commitment_sender, &sender_keys.nsk), - commitment_set.digest(), - ), - ( - Nullifier::for_account_initialization(&recipient_account_id), - DUMMY_COMMITMENT_HASH, - ), - ]; - - let program = crate::test_methods::simple_balance_transfer(); - - let expected_private_account_1 = Account { - program_owner: program.id(), - balance: 100 - balance_to_move, - nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk), - ..Default::default() - }; - let expected_private_account_2 = Account { - program_owner: program.id(), - balance: balance_to_move, - nonce: Nonce::private_account_nonce_init(&recipient_account_id), - ..Default::default() - }; - let expected_new_commitments = vec![ - Commitment::new(&sender_account_id, &expected_private_account_1), - Commitment::new(&recipient_account_id, &expected_private_account_2), - ]; - - let shared_secret_1 = - SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0; - - let shared_secret_2 = - SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1).0; - - let (output, proof) = execute_and_prove( - vec![sender_pre, recipient], - Program::serialize_instruction(balance_to_move).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: shared_secret_1, - nsk: sender_keys.nsk, - membership_proof: commitment_set - .get_proof_for(&commitment_sender) - .expect("sender's commitment must be in the set"), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: shared_secret_2, - identifier: 0, - }, - ], - &program.into(), - ) - .unwrap(); - - assert!(proof.is_valid_for(&output)); - assert!(output.public_pre_states.is_empty()); - assert!(output.public_post_states.is_empty()); - assert_eq!(output.new_commitments, expected_new_commitments); - assert_eq!(output.new_nullifiers, expected_new_nullifiers); - assert_eq!(output.encrypted_private_post_states.len(), 2); - - let (_identifier, sender_post) = EncryptionScheme::decrypt( - &output.encrypted_private_post_states[0].ciphertext, - &shared_secret_1, - &expected_new_commitments[0], - 0, - ) - .unwrap(); - assert_eq!(sender_post, expected_private_account_1); - - let (_identifier, recipient_post) = EncryptionScheme::decrypt( - &output.encrypted_private_post_states[1].ciphertext, - &shared_secret_2, - &expected_new_commitments[1], - 1, - ) - .unwrap(); - assert_eq!(recipient_post, expected_private_account_2); - } - - #[test] - fn circuit_fails_when_chained_validity_windows_have_empty_intersection() { - let account_keys = test_private_account_keys_1(); - let pre = AccountWithMetadata::new( - Account::default(), - false, - AccountId::for_regular_private_account(&account_keys.npk(), 0), - ); - - let validity_window_chain_caller = crate::test_methods::validity_window_chain_caller(); - let validity_window = crate::test_methods::validity_window(); - - let instruction = Program::serialize_instruction(( - Some(1_u64), - Some(4_u64), - validity_window.id(), - Some(4_u64), - Some(7_u64), - )) - .unwrap(); - - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0).0; - - let program_with_deps = ProgramWithDependencies::new( - validity_window_chain_caller, - [(validity_window.id(), validity_window)].into(), - ); - - let result = execute_and_prove( - vec![pre], - instruction, - vec![InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &account_keys.npk(), - &account_keys.vpk(), - ), - npk: account_keys.npk(), - ssk: shared_secret, - identifier: 0, - }], - &program_with_deps, - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - /// A private PDA claimed with a non-default identifier produces a ciphertext that decrypts - /// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`. - #[test] - fn private_pda_claim_with_custom_identifier_encrypts_correct_kind() { - let program = crate::test_methods::pda_claimer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let identifier: u128 = 99; - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier); - let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); - - let (output, _proof) = execute_and_prove( - vec![pre_state], - Program::serialize_instruction(seed).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier, - seed: None, - }], - &program.clone().into(), - ) - .unwrap(); - - assert_eq!( - decrypt_kind(&output, &shared_secret, 0), - PrivateAccountKind::Pda { - program_id: program.id(), - seed, - identifier - }, - ); - } - - /// PDA init: initializes a new PDA under `simple_balance_transfer`'s ownership. - /// The `simple_transfer_proxy` program chains to `simple_balance_transfer` with `pda_seeds` - /// to establish authorization and the private PDA binding. - #[test] - fn private_pda_init() { - let program = crate::test_methods::simple_transfer_proxy(); - let simple_transfer = crate::test_methods::simple_balance_transfer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let shared_secret_pda = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - // PDA (new, private PDA) - let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0); - let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id); - - let auth_id = simple_transfer.id(); - let program_with_deps = - ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into()); - - // is_withdraw=false triggers init path (1 pre-state) - let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, false)).unwrap(); - - let result = execute_and_prove( - vec![pda_pre], - instruction, - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret_pda, - identifier: 0, - seed: None, - }], - &program_with_deps, - ); - - let (output, _proof) = result.expect("PDA init should succeed"); - assert_eq!(output.new_commitments.len(), 1); - } - - /// PDA withdraw: chains to `simple_balance_transfer` to move balance from PDA to recipient. - /// Uses a default PDA (amount=0) because testing with a pre-funded PDA requires a - /// two-tx sequence with membership proofs. - #[test] - fn private_pda_withdraw() { - let program = crate::test_methods::simple_transfer_proxy(); - let simple_transfer = crate::test_methods::simple_balance_transfer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let shared_secret_pda = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - // PDA (new, private PDA) - let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0); - let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id); - - // Recipient (public) - let recipient_id = AccountId::new([88; 32]); - let recipient_pre = AccountWithMetadata::new( - Account { - program_owner: simple_transfer.id(), - balance: 10000, - ..Account::default() - }, - true, - recipient_id, - ); - - let auth_id = simple_transfer.id(); - let program_with_deps = - ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into()); - - // is_withdraw=true, amount=0 (PDA has no balance yet) - let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, true)).unwrap(); - - let result = execute_and_prove( - vec![pda_pre, recipient_pre], - instruction, - vec![ - InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret_pda, - identifier: 0, - seed: None, - }, - InputAccountIdentity::Public, - ], - &program_with_deps, - ); - - let (output, _proof) = result.expect("PDA withdraw should succeed"); - assert_eq!(output.new_commitments.len(), 1); - } - - /// Shared regular private account: receives funds via `authenticated_transfer` directly, - /// no custom program needed. This demonstrates the non-PDA shared account flow where - /// keys are derived from GMS via `derive_keys_for_shared_account`. The shared account - /// uses the standard unauthorized private account path and works with auth-transfer's - /// transfer path like any other private account. - #[test] - fn shared_account_receives_via_simple_transfer() { - let program = crate::test_methods::simple_balance_transfer(); - let shared_keys = test_private_account_keys_1(); - let shared_npk = shared_keys.npk(); - let shared_identifier: u128 = 42; - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&shared_keys.vpk(), &[0_u8; 32], 0).0; - - // Sender: public account with balance, owned by auth-transfer - let sender_id = AccountId::new([99; 32]); - let sender = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 1000, - ..Account::default() - }, - true, - sender_id, - ); - - // Recipient: shared private account (new, unauthorized) - let shared_account_id = AccountId::from((&shared_npk, shared_identifier)); - let recipient = AccountWithMetadata::new(Account::default(), false, shared_account_id); - - let balance_to_move: u128 = 100; - let instruction = Program::serialize_instruction(balance_to_move).unwrap(); - - let result = execute_and_prove( - vec![sender, recipient], - instruction, - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &shared_npk, - &shared_keys.vpk(), - ), - npk: shared_npk, - ssk: shared_secret, - identifier: shared_identifier, - }, - ], - &program.into(), - ); - - let (output, _proof) = result.expect("shared account receive should succeed"); - // Sender is public (no commitment), recipient is private (1 commitment) - assert_eq!(output.new_commitments.len(), 1); - } - - /// `PrivateAuthorizedInit` with a non-default identifier produces a ciphertext that decrypts - /// to `PrivateAccountKind::Regular` carrying the correct identifier. - #[test] - fn private_authorized_init_encrypts_regular_kind_with_identifier() { - let program = crate::test_methods::claimer(); - let keys = test_private_account_keys_1(); - let identifier: u128 = 99; - let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier); - let pre = AccountWithMetadata::new(Account::default(), true, account_id); - - let (output, _) = execute_and_prove( - vec![pre], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()), - ssk, - nsk: keys.nsk, - identifier, - }], - &program.into(), - ) - .unwrap(); - - assert_eq!( - decrypt_kind(&output, &ssk, 0), - PrivateAccountKind::Regular(identifier) - ); - } - - /// `PrivateUnauthorized` with a non-default identifier produces a ciphertext that decrypts - /// to `PrivateAccountKind::Regular` carrying the correct identifier. - #[test] - fn private_unauthorized_init_encrypts_regular_kind_with_identifier() { - let program = crate::test_methods::claimer(); - let keys = test_private_account_keys_1(); - let identifier: u128 = 99; - let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - let recipient_id = AccountId::for_regular_private_account(&keys.npk(), identifier); - let recipient = AccountWithMetadata::new(Account::default(), false, recipient_id); - - let (output, _) = execute_and_prove( - vec![recipient], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()), - npk: keys.npk(), - ssk, - identifier, - }], - &program.into(), - ) - .unwrap(); - - assert_eq!( - decrypt_kind(&output, &ssk, 0), - PrivateAccountKind::Regular(identifier) - ); - } - - /// `PrivateAuthorizedUpdate` with a non-default identifier produces a ciphertext that decrypts - /// to `PrivateAccountKind::Regular` carrying the correct identifier. - #[test] - fn private_authorized_update_encrypts_regular_kind_with_identifier() { - let program = crate::test_methods::noop(); - let keys = test_private_account_keys_1(); - let identifier: u128 = 99; - let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier); - let account = Account { - program_owner: program.id(), - balance: 1, - ..Account::default() - }; - let commitment = Commitment::new(&account_id, &account); - let mut commitment_set = CommitmentSet::with_capacity(1); - commitment_set.extend(std::slice::from_ref(&commitment)); - - let sender = AccountWithMetadata::new(account, true, account_id); - - let (output, _) = execute_and_prove( - vec![sender], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()), - ssk, - nsk: keys.nsk, - membership_proof: commitment_set.get_proof_for(&commitment).unwrap(), - identifier, - }], - &program.into(), - ) - .unwrap(); - - assert_eq!( - decrypt_kind(&output, &ssk, 0), - PrivateAccountKind::Regular(identifier) - ); - } - - /// `PrivatePdaUpdate` with a non-default identifier produces a ciphertext that decrypts - /// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`. - #[test] - fn private_pda_update_encrypts_pda_kind_with_identifier() { - let program = crate::test_methods::pda_spend_proxy(); - let simple_transfer = crate::test_methods::simple_balance_transfer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let identifier: u128 = 99; - let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let simple_transfer_id = simple_transfer.id(); - let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier); - let pda_account = Account { - program_owner: simple_transfer_id, - balance: 1, - ..Account::default() - }; - let pda_commitment = Commitment::new(&pda_id, &pda_account); - let mut commitment_set = CommitmentSet::with_capacity(1); - commitment_set.extend(std::slice::from_ref(&pda_commitment)); - - let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id); - let recipient_pre = - AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32])); - - let program_with_deps = ProgramWithDependencies::new( - program.clone(), - [(simple_transfer_id, simple_transfer)].into(), - ); - - let (output, _) = execute_and_prove( - vec![pda_pre, recipient_pre], - Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(), - vec![ - InputAccountIdentity::PrivatePdaUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - ssk, - nsk: keys.nsk, - membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(), - identifier, - seed: None, - }, - InputAccountIdentity::Public, - ], - &program_with_deps, - ) - .unwrap(); - - assert_eq!( - decrypt_kind(&output, &ssk, 0), - PrivateAccountKind::Pda { - program_id: program.id(), - seed, - identifier - }, - ); - } - - #[test] - fn private_pda_init_identifier_mismatch_fails() { - let program = crate::test_methods::pda_claimer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5); - let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); - - let result = execute_and_prove( - vec![pre_state], - Program::serialize_instruction(seed).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier: 99, - seed: None, - }], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn private_pda_update_identifier_mismatch_fails() { - let program = crate::test_methods::pda_spend_proxy(); - let simple_transfer = crate::test_methods::simple_balance_transfer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let simple_transfer_id = simple_transfer.id(); - let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5); - let pda_account = Account { - program_owner: simple_transfer_id, - balance: 1, - ..Account::default() - }; - let pda_commitment = Commitment::new(&pda_id, &pda_account); - let mut commitment_set = CommitmentSet::with_capacity(1); - commitment_set.extend(std::slice::from_ref(&pda_commitment)); - - let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id); - let recipient_pre = - AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32])); - - let program_with_deps = - ProgramWithDependencies::new(program, [(simple_transfer_id, simple_transfer)].into()); - - let result = execute_and_prove( - vec![pda_pre, recipient_pre], - Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(), - vec![ - InputAccountIdentity::PrivatePdaUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - ssk, - nsk: keys.nsk, - membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(), - identifier: 99, - seed: None, - }, - InputAccountIdentity::Public, - ], - &program_with_deps, - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } -} diff --git a/lee/state_machine/src/privacy_preserving_transaction/circuit/mod.rs b/lee/state_machine/src/privacy_preserving_transaction/circuit/mod.rs new file mode 100644 index 00000000..da8b2be4 --- /dev/null +++ b/lee/state_machine/src/privacy_preserving_transaction/circuit/mod.rs @@ -0,0 +1,177 @@ +use std::collections::{HashMap, VecDeque}; + +use borsh::{BorshDeserialize, BorshSerialize}; +use lee_core::{ + InputAccountIdentity, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, + account::AccountWithMetadata, + program::{ChainedCall, InstructionData, ProgramId, ProgramOutput}, +}; +use risc0_zkvm::{ExecutorEnv, InnerReceipt, ProverOpts, Receipt, default_prover}; + +use crate::{ + PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID, + error::{InvalidProgramBehaviorError, LeeError}, + program::Program, + state::MAX_NUMBER_CHAINED_CALLS, +}; + +/// Proof of the privacy preserving execution circuit. +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub struct Proof(pub(crate) Vec); + +impl Proof { + #[must_use] + pub fn into_inner(self) -> Vec { + self.0 + } + + #[must_use] + pub const fn from_inner(inner: Vec) -> Self { + Self(inner) + } + + pub(crate) fn is_valid_for(&self, circuit_output: &PrivacyPreservingCircuitOutput) -> bool { + let Ok(inner) = borsh::from_slice::(&self.0) else { + return false; + }; + let receipt = Receipt::new(inner, circuit_output.to_bytes()); + receipt.verify(PRIVACY_PRESERVING_CIRCUIT_ID).is_ok() + } +} + +#[derive(Clone)] +pub struct ProgramWithDependencies { + pub program: Program, + // TODO: avoid having a copy of the bytecode of each dependency. + pub dependencies: HashMap, +} + +impl ProgramWithDependencies { + #[must_use] + pub const fn new(program: Program, dependencies: HashMap) -> Self { + Self { + program, + dependencies, + } + } +} + +impl From for ProgramWithDependencies { + fn from(program: Program) -> Self { + Self::new(program, HashMap::new()) + } +} + +/// Generates a proof of the execution of a LEE program inside the privacy preserving execution +/// circuit. +pub fn execute_and_prove( + pre_states: Vec, + instruction_data: InstructionData, + account_identities: Vec, + program_with_dependencies: &ProgramWithDependencies, +) -> Result<(PrivacyPreservingCircuitOutput, Proof), LeeError> { + let ProgramWithDependencies { + program: initial_program, + dependencies, + } = program_with_dependencies; + let mut env_builder = ExecutorEnv::builder(); + let mut program_outputs = Vec::new(); + + let initial_call = ChainedCall { + program_id: initial_program.id(), + instruction_data, + pre_states, + pda_seeds: vec![], + }; + + let mut chained_calls = VecDeque::from_iter([(initial_call, initial_program, None)]); + let mut chain_calls_counter = 0; + while let Some((chained_call, program, caller_program_id)) = chained_calls.pop_front() { + if chain_calls_counter >= MAX_NUMBER_CHAINED_CALLS { + return Err(LeeError::MaxChainedCallsDepthExceeded); + } + + let inner_receipt = execute_and_prove_program( + program, + caller_program_id, + &chained_call.pre_states, + &chained_call.instruction_data, + )?; + + let program_output: ProgramOutput = inner_receipt + .journal + .decode() + .map_err(|e| LeeError::ProgramOutputDeserializationError(e.to_string()))?; + + // TODO: remove clone + program_outputs.push(program_output.clone()); + + // Prove circuit. + env_builder.add_assumption(inner_receipt); + + for new_call in program_output.chained_calls.into_iter().rev() { + let next_program = dependencies.get(&new_call.program_id).ok_or( + InvalidProgramBehaviorError::UndeclaredProgramDependency { + program_id: new_call.program_id, + }, + )?; + chained_calls.push_front((new_call, next_program, Some(chained_call.program_id))); + } + + chain_calls_counter = chain_calls_counter + .checked_add(1) + .expect("we check the max depth at the beginning of the loop"); + } + + let circuit_input = PrivacyPreservingCircuitInput { + program_outputs, + account_identities, + program_id: program_with_dependencies.program.id(), + }; + + env_builder.write(&circuit_input).unwrap(); + let env = env_builder.build().unwrap(); + let prover = default_prover(); + let opts = ProverOpts::succinct(); + let prove_info = prover + .prove_with_opts(env, PRIVACY_PRESERVING_CIRCUIT_ELF, &opts) + .map_err(|e| LeeError::CircuitProvingError(e.to_string()))?; + + let proof = Proof(borsh::to_vec(&prove_info.receipt.inner)?); + + let circuit_output: PrivacyPreservingCircuitOutput = prove_info + .receipt + .journal + .decode() + .map_err(|e| LeeError::CircuitOutputDeserializationError(e.to_string()))?; + + Ok((circuit_output, proof)) +} + +fn execute_and_prove_program( + program: &Program, + caller_program_id: Option, + pre_states: &[AccountWithMetadata], + instruction_data: &InstructionData, +) -> Result { + // Write inputs to the program + let mut env_builder = ExecutorEnv::builder(); + Program::write_inputs( + program.id(), + caller_program_id, + pre_states, + instruction_data, + &mut env_builder, + )?; + let env = env_builder.build().unwrap(); + + // Prove the program + let prover = default_prover(); + Ok(prover + .prove(env, program.elf()) + .map_err(|e| LeeError::ProgramProveFailed(e.to_string()))? + .receipt) +} + +#[cfg(test)] +mod tests; diff --git a/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs b/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs new file mode 100644 index 00000000..970bfbd4 --- /dev/null +++ b/lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs @@ -0,0 +1,731 @@ +#![expect(clippy::shadow_unrelated, reason = "We don't care about it in tests")] + +use lee_core::{ + Commitment, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, + Nullifier, PrivacyPreservingCircuitOutput, SharedSecretKey, + account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data}, + program::{PdaSeed, PrivateAccountKind}, +}; + +use super::*; +use crate::{ + error::LeeError, + privacy_preserving_transaction::circuit::execute_and_prove, + program::Program, + state::{ + CommitmentSet, + tests::{test_private_account_keys_1, test_private_account_keys_2}, + }, +}; + +fn decrypt_kind( + output: &PrivacyPreservingCircuitOutput, + ssk: &SharedSecretKey, + idx: usize, +) -> PrivateAccountKind { + let (kind, _) = EncryptionScheme::decrypt( + &output.encrypted_private_post_states[idx].ciphertext, + ssk, + &output.new_commitments[idx], + u32::try_from(idx).expect("idx fits in u32"), + ) + .unwrap(); + kind +} + +#[test] +fn proof_inner_roundtrip() { + // `Proof::from_inner(b).into_inner()` must return exactly `b`. Catches + // mutations of `into_inner` returning `vec![]`, `vec![0]`, or `vec![1]`, + // and of `from_inner` discarding its argument. + let bytes = vec![0xDE_u8, 0xAD, 0xBE, 0xEF]; + assert_eq!(Proof::from_inner(bytes.clone()).into_inner(), bytes); + assert!(Proof::from_inner(vec![]).into_inner().is_empty()); + assert_eq!(Proof::from_inner(vec![0xFF]).into_inner(), vec![0xFF_u8]); +} + +#[test] +fn prove_privacy_preserving_execution_circuit_public_and_private_pre_accounts() { + let recipient_keys = test_private_account_keys_1(); + let program = crate::test_methods::simple_balance_transfer(); + let sender = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); + let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id); + + let balance_to_move: u128 = 37; + + let expected_sender_post = Account { + program_owner: program.id(), + balance: 100 - balance_to_move, + nonce: Nonce::default(), + data: Data::default(), + }; + + let expected_recipient_post = Account { + program_owner: program.id(), + balance: balance_to_move, + nonce: Nonce::private_account_nonce_init(&recipient_account_id), + data: Data::default(), + }; + + let expected_sender_pre = sender.clone(); + + let shared_secret = + SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0; + + let (output, proof) = execute_and_prove( + vec![sender, recipient], + Program::serialize_instruction(balance_to_move).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: shared_secret, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + ], + &crate::test_methods::simple_balance_transfer().into(), + ) + .unwrap(); + + assert!(proof.is_valid_for(&output)); + + let [sender_pre] = output.public_pre_states.try_into().unwrap(); + let [sender_post] = output.public_post_states.try_into().unwrap(); + assert_eq!(sender_pre, expected_sender_pre); + assert_eq!(sender_post, expected_sender_post); + assert_eq!(output.new_commitments.len(), 1); + assert_eq!(output.new_nullifiers.len(), 1); + assert_eq!(output.encrypted_private_post_states.len(), 1); + + let (_identifier, recipient_post) = EncryptionScheme::decrypt( + &output.encrypted_private_post_states[0].ciphertext, + &shared_secret, + &output.new_commitments[0], + 0, + ) + .unwrap(); + assert_eq!(recipient_post, expected_recipient_post); +} + +#[test] +fn prove_privacy_preserving_execution_circuit_fully_private() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + + let sender_nonce = Nonce(0xdead_beef); + let sender_pre = AccountWithMetadata::new( + Account { + balance: 100, + nonce: sender_nonce, + program_owner: program.id(), + data: Data::default(), + }, + true, + AccountId::for_regular_private_account(&sender_keys.npk(), 0), + ); + let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); + let commitment_sender = Commitment::new(&sender_account_id, &sender_pre.account); + + let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); + let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id); + let balance_to_move: u128 = 37; + + let mut commitment_set = CommitmentSet::with_capacity(2); + commitment_set.extend(std::slice::from_ref(&commitment_sender)); + let expected_new_nullifiers = vec![ + ( + Nullifier::for_account_update(&commitment_sender, &sender_keys.nsk), + commitment_set.digest(), + ), + ( + Nullifier::for_account_initialization(&recipient_account_id), + DUMMY_COMMITMENT_HASH, + ), + ]; + + let program = crate::test_methods::simple_balance_transfer(); + + let expected_private_account_1 = Account { + program_owner: program.id(), + balance: 100 - balance_to_move, + nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk), + ..Default::default() + }; + let expected_private_account_2 = Account { + program_owner: program.id(), + balance: balance_to_move, + nonce: Nonce::private_account_nonce_init(&recipient_account_id), + ..Default::default() + }; + let expected_new_commitments = vec![ + Commitment::new(&sender_account_id, &expected_private_account_1), + Commitment::new(&recipient_account_id, &expected_private_account_2), + ]; + + let shared_secret_1 = + SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0; + + let shared_secret_2 = + SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1).0; + + let (output, proof) = execute_and_prove( + vec![sender_pre, recipient], + Program::serialize_instruction(balance_to_move).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: shared_secret_1, + nsk: sender_keys.nsk, + membership_proof: commitment_set + .get_proof_for(&commitment_sender) + .expect("sender's commitment must be in the set"), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: shared_secret_2, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + ], + &program.into(), + ) + .unwrap(); + + assert!(proof.is_valid_for(&output)); + assert!(output.public_pre_states.is_empty()); + assert!(output.public_post_states.is_empty()); + assert_eq!(output.new_commitments, expected_new_commitments); + assert_eq!(output.new_nullifiers, expected_new_nullifiers); + assert_eq!(output.encrypted_private_post_states.len(), 2); + + let (_identifier, sender_post) = EncryptionScheme::decrypt( + &output.encrypted_private_post_states[0].ciphertext, + &shared_secret_1, + &expected_new_commitments[0], + 0, + ) + .unwrap(); + assert_eq!(sender_post, expected_private_account_1); + + let (_identifier, recipient_post) = EncryptionScheme::decrypt( + &output.encrypted_private_post_states[1].ciphertext, + &shared_secret_2, + &expected_new_commitments[1], + 1, + ) + .unwrap(); + assert_eq!(recipient_post, expected_private_account_2); +} + +#[test] +fn circuit_fails_when_chained_validity_windows_have_empty_intersection() { + let account_keys = test_private_account_keys_1(); + let pre = AccountWithMetadata::new( + Account::default(), + false, + AccountId::for_regular_private_account(&account_keys.npk(), 0), + ); + + let validity_window_chain_caller = crate::test_methods::validity_window_chain_caller(); + let validity_window = crate::test_methods::validity_window(); + + let instruction = Program::serialize_instruction(( + Some(1_u64), + Some(4_u64), + validity_window.id(), + Some(4_u64), + Some(7_u64), + )) + .unwrap(); + + let shared_secret = + SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0).0; + + let program_with_deps = ProgramWithDependencies::new( + validity_window_chain_caller, + [(validity_window.id(), validity_window)].into(), + ); + + let result = execute_and_prove( + vec![pre], + instruction, + vec![InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &account_keys.npk(), + &account_keys.vpk(), + ), + npk: account_keys.npk(), + ssk: shared_secret, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &program_with_deps, + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +/// A private PDA claimed with a non-default identifier produces a ciphertext that decrypts +/// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`. +#[test] +fn private_pda_claim_with_custom_identifier_encrypts_correct_kind() { + let program = crate::test_methods::pda_claimer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let identifier: u128 = 99; + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let (output, _proof) = execute_and_prove( + vec![pre_state], + Program::serialize_instruction(seed).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program.clone().into(), + ) + .unwrap(); + + assert_eq!( + decrypt_kind(&output, &shared_secret, 0), + PrivateAccountKind::Pda { + program_id: program.id(), + seed, + identifier + }, + ); +} + +/// PDA init: initializes a new PDA under `simple_balance_transfer`'s ownership. +/// The `simple_transfer_proxy` program chains to `simple_balance_transfer` with `pda_seeds` +/// to establish authorization and the private PDA binding. +#[test] +fn private_pda_init() { + let program = crate::test_methods::simple_transfer_proxy(); + let simple_transfer = crate::test_methods::simple_balance_transfer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let shared_secret_pda = + SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + // PDA (new, private PDA) + let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0); + let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id); + + let auth_id = simple_transfer.id(); + let program_with_deps = + ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into()); + + // is_withdraw=false triggers init path (1 pre-state) + let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, false)).unwrap(); + + let result = execute_and_prove( + vec![pda_pre], + instruction, + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret_pda, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program_with_deps, + ); + + let (output, _proof) = result.expect("PDA init should succeed"); + assert_eq!(output.new_commitments.len(), 1); +} + +/// PDA withdraw: chains to `simple_balance_transfer` to move balance from PDA to recipient. +/// Uses a default PDA (amount=0) because testing with a pre-funded PDA requires a +/// two-tx sequence with membership proofs. +#[test] +fn private_pda_withdraw() { + let program = crate::test_methods::simple_transfer_proxy(); + let simple_transfer = crate::test_methods::simple_balance_transfer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let shared_secret_pda = + SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + // PDA (new, private PDA) + let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0); + let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id); + + // Recipient (public) + let recipient_id = AccountId::new([88; 32]); + let recipient_pre = AccountWithMetadata::new( + Account { + program_owner: simple_transfer.id(), + balance: 10000, + ..Account::default() + }, + true, + recipient_id, + ); + + let auth_id = simple_transfer.id(); + let program_with_deps = + ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into()); + + // is_withdraw=true, amount=0 (PDA has no balance yet) + let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, true)).unwrap(); + + let result = execute_and_prove( + vec![pda_pre, recipient_pre], + instruction, + vec![ + InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret_pda, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }, + InputAccountIdentity::Public, + ], + &program_with_deps, + ); + + let (output, _proof) = result.expect("PDA withdraw should succeed"); + assert_eq!(output.new_commitments.len(), 1); +} + +/// Shared regular private account: receives funds via `authenticated_transfer` directly, +/// no custom program needed. This demonstrates the non-PDA shared account flow where +/// keys are derived from GMS via `derive_keys_for_shared_account`. The shared account +/// uses the standard unauthorized private account path and works with auth-transfer's +/// transfer path like any other private account. +#[test] +fn shared_account_receives_via_simple_transfer() { + let program = crate::test_methods::simple_balance_transfer(); + let shared_keys = test_private_account_keys_1(); + let shared_npk = shared_keys.npk(); + let shared_identifier: u128 = 42; + let shared_secret = + SharedSecretKey::encapsulate_deterministic(&shared_keys.vpk(), &[0_u8; 32], 0).0; + + // Sender: public account with balance, owned by auth-transfer + let sender_id = AccountId::new([99; 32]); + let sender = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 1000, + ..Account::default() + }, + true, + sender_id, + ); + + // Recipient: shared private account (new, unauthorized) + let shared_account_id = AccountId::from((&shared_npk, shared_identifier)); + let recipient = AccountWithMetadata::new(Account::default(), false, shared_account_id); + + let balance_to_move: u128 = 100; + let instruction = Program::serialize_instruction(balance_to_move).unwrap(); + + let result = execute_and_prove( + vec![sender, recipient], + instruction, + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&shared_npk, &shared_keys.vpk()), + npk: shared_npk, + ssk: shared_secret, + identifier: shared_identifier, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + ], + &program.into(), + ); + + let (output, _proof) = result.expect("shared account receive should succeed"); + // Sender is public (no commitment), recipient is private (1 commitment) + assert_eq!(output.new_commitments.len(), 1); +} + +/// `PrivateAuthorizedInit` with a non-default identifier produces a ciphertext that decrypts +/// to `PrivateAccountKind::Regular` carrying the correct identifier. +#[test] +fn private_authorized_init_encrypts_regular_kind_with_identifier() { + let program = crate::test_methods::claimer(); + let keys = test_private_account_keys_1(); + let identifier: u128 = 99; + let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier); + let pre = AccountWithMetadata::new(Account::default(), true, account_id); + + let (output, _) = execute_and_prove( + vec![pre], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()), + ssk, + nsk: keys.nsk, + identifier, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &program.into(), + ) + .unwrap(); + + assert_eq!( + decrypt_kind(&output, &ssk, 0), + PrivateAccountKind::Regular(identifier) + ); +} + +/// `PrivateUnauthorized` with a non-default identifier produces a ciphertext that decrypts +/// to `PrivateAccountKind::Regular` carrying the correct identifier. +#[test] +fn private_unauthorized_init_encrypts_regular_kind_with_identifier() { + let program = crate::test_methods::claimer(); + let keys = test_private_account_keys_1(); + let identifier: u128 = 99; + let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + let recipient_id = AccountId::for_regular_private_account(&keys.npk(), identifier); + let recipient = AccountWithMetadata::new(Account::default(), false, recipient_id); + + let (output, _) = execute_and_prove( + vec![recipient], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()), + npk: keys.npk(), + ssk, + identifier, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &program.into(), + ) + .unwrap(); + + assert_eq!( + decrypt_kind(&output, &ssk, 0), + PrivateAccountKind::Regular(identifier) + ); +} + +/// `PrivateAuthorizedUpdate` with a non-default identifier produces a ciphertext that decrypts +/// to `PrivateAccountKind::Regular` carrying the correct identifier. +#[test] +fn private_authorized_update_encrypts_regular_kind_with_identifier() { + let program = crate::test_methods::noop(); + let keys = test_private_account_keys_1(); + let identifier: u128 = 99; + let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier); + let account = Account { + program_owner: program.id(), + balance: 1, + ..Account::default() + }; + let commitment = Commitment::new(&account_id, &account); + let mut commitment_set = CommitmentSet::with_capacity(1); + commitment_set.extend(std::slice::from_ref(&commitment)); + + let sender = AccountWithMetadata::new(account, true, account_id); + + let (output, _) = execute_and_prove( + vec![sender], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()), + ssk, + nsk: keys.nsk, + membership_proof: commitment_set.get_proof_for(&commitment).unwrap(), + identifier, + }], + &program.into(), + ) + .unwrap(); + + assert_eq!( + decrypt_kind(&output, &ssk, 0), + PrivateAccountKind::Regular(identifier) + ); +} + +/// `PrivatePdaUpdate` with a non-default identifier produces a ciphertext that decrypts +/// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`. +#[test] +fn private_pda_update_encrypts_pda_kind_with_identifier() { + let program = crate::test_methods::pda_spend_proxy(); + let simple_transfer = crate::test_methods::simple_balance_transfer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let identifier: u128 = 99; + let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let simple_transfer_id = simple_transfer.id(); + let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier); + let pda_account = Account { + program_owner: simple_transfer_id, + balance: 1, + ..Account::default() + }; + let pda_commitment = Commitment::new(&pda_id, &pda_account); + let mut commitment_set = CommitmentSet::with_capacity(1); + commitment_set.extend(std::slice::from_ref(&pda_commitment)); + + let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id); + let recipient_pre = AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32])); + + let program_with_deps = ProgramWithDependencies::new( + program.clone(), + [(simple_transfer_id, simple_transfer)].into(), + ); + + let (output, _) = execute_and_prove( + vec![pda_pre, recipient_pre], + Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(), + vec![ + InputAccountIdentity::PrivatePdaUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + ssk, + nsk: keys.nsk, + membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(), + identifier, + seed: None, + }, + InputAccountIdentity::Public, + ], + &program_with_deps, + ) + .unwrap(); + + assert_eq!( + decrypt_kind(&output, &ssk, 0), + PrivateAccountKind::Pda { + program_id: program.id(), + seed, + identifier + }, + ); +} + +#[test] +fn private_pda_init_identifier_mismatch_fails() { + let program = crate::test_methods::pda_claimer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let result = execute_and_prove( + vec![pre_state], + Program::serialize_instruction(seed).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier: 99, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn private_pda_update_identifier_mismatch_fails() { + let program = crate::test_methods::pda_spend_proxy(); + let simple_transfer = crate::test_methods::simple_balance_transfer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let simple_transfer_id = simple_transfer.id(); + let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5); + let pda_account = Account { + program_owner: simple_transfer_id, + balance: 1, + ..Account::default() + }; + let pda_commitment = Commitment::new(&pda_id, &pda_account); + let mut commitment_set = CommitmentSet::with_capacity(1); + commitment_set.extend(std::slice::from_ref(&pda_commitment)); + + let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id); + let recipient_pre = AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32])); + + let program_with_deps = + ProgramWithDependencies::new(program, [(simple_transfer_id, simple_transfer)].into()); + + let result = execute_and_prove( + vec![pda_pre, recipient_pre], + Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(), + vec![ + InputAccountIdentity::PrivatePdaUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + ssk, + nsk: keys.nsk, + membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(), + identifier: 99, + seed: None, + }, + InputAccountIdentity::Public, + ], + &program_with_deps, + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} diff --git a/lee/state_machine/src/program.rs b/lee/state_machine/src/program/mod.rs similarity index 72% rename from lee/state_machine/src/program.rs rename to lee/state_machine/src/program/mod.rs index 65d60a42..d481c1fa 100644 --- a/lee/state_machine/src/program.rs +++ b/lee/state_machine/src/program/mod.rs @@ -111,42 +111,4 @@ impl Program { } #[cfg(test)] -mod tests { - use lee_core::account::{Account, AccountId, AccountWithMetadata}; - - use crate::program::Program; - - #[test] - fn program_execution() { - let program = crate::test_methods::simple_balance_transfer(); - let balance_to_move: u128 = 11_223_344_556_677; - let instruction_data = Program::serialize_instruction(balance_to_move).unwrap(); - let sender = AccountWithMetadata::new( - Account { - balance: 77_665_544_332_211, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - let recipient = - AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32])); - - let expected_sender_post = Account { - balance: 77_665_544_332_211 - balance_to_move, - ..Account::default() - }; - let expected_recipient_post = Account { - balance: balance_to_move, - ..Account::default() - }; - let program_output = program - .execute(None, &[sender, recipient], &instruction_data) - .unwrap(); - - let [sender_post, recipient_post] = program_output.post_states.try_into().unwrap(); - - assert_eq!(sender_post.account(), &expected_sender_post); - assert_eq!(recipient_post.account(), &expected_recipient_post); - } -} +mod tests; diff --git a/lee/state_machine/src/program/tests.rs b/lee/state_machine/src/program/tests.rs new file mode 100644 index 00000000..330bd0d6 --- /dev/null +++ b/lee/state_machine/src/program/tests.rs @@ -0,0 +1,36 @@ +use lee_core::account::{Account, AccountId, AccountWithMetadata}; + +use crate::program::Program; + +#[test] +fn program_execution() { + let program = crate::test_methods::simple_balance_transfer(); + let balance_to_move: u128 = 11_223_344_556_677; + let instruction_data = Program::serialize_instruction(balance_to_move).unwrap(); + let sender = AccountWithMetadata::new( + Account { + balance: 77_665_544_332_211, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + let recipient = AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32])); + + let expected_sender_post = Account { + balance: 77_665_544_332_211 - balance_to_move, + ..Account::default() + }; + let expected_recipient_post = Account { + balance: balance_to_move, + ..Account::default() + }; + let program_output = program + .execute(None, &[sender, recipient], &instruction_data) + .unwrap(); + + let [sender_post, recipient_post] = program_output.post_states.try_into().unwrap(); + + assert_eq!(sender_post.account(), &expected_sender_post); + assert_eq!(recipient_post.account(), &expected_recipient_post); +} diff --git a/lee/state_machine/src/state.rs b/lee/state_machine/src/state.rs deleted file mode 100644 index c399cea1..00000000 --- a/lee/state_machine/src/state.rs +++ /dev/null @@ -1,4414 +0,0 @@ -use std::collections::{BTreeSet, HashMap, HashSet}; - -use borsh::{BorshDeserialize, BorshSerialize}; -use lee_core::{ - BlockId, Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier, - Timestamp, - account::{Account, AccountId}, - program::ProgramId, -}; - -use crate::{ - error::LeeError, - merkle_tree::MerkleTree, - privacy_preserving_transaction::PrivacyPreservingTransaction, - program::Program, - program_deployment_transaction::ProgramDeploymentTransaction, - public_transaction::PublicTransaction, - validated_state_diff::{StateDiff, ValidatedStateDiff}, -}; - -pub const MAX_NUMBER_CHAINED_CALLS: usize = 10; - -#[derive(Clone, BorshSerialize, BorshDeserialize)] -#[cfg_attr(test, derive(Debug, PartialEq, Eq))] -pub struct CommitmentSet { - merkle_tree: MerkleTree, - commitments: HashMap, - root_history: HashSet, -} - -impl CommitmentSet { - pub(crate) fn digest(&self) -> CommitmentSetDigest { - self.merkle_tree.root() - } - - /// Queries the `CommitmentSet` for a membership proof of commitment. - pub fn get_proof_for(&self, commitment: &Commitment) -> Option { - let index = *self.commitments.get(commitment)?; - - self.merkle_tree - .get_authentication_path_for(index) - .map(|path| (index, path)) - } - - /// Inserts a list of commitments to the `CommitmentSet`. - pub(crate) fn extend(&mut self, commitments: &[Commitment]) { - for commitment in commitments.iter().cloned() { - let index = self.merkle_tree.insert(commitment.to_byte_array()); - self.commitments.insert(commitment, index); - } - self.root_history.insert(self.digest()); - } - - fn contains(&self, commitment: &Commitment) -> bool { - self.commitments.contains_key(commitment) - } - - /// Initializes an empty `CommitmentSet` with a given capacity. - /// If the capacity is not a `power_of_two`, then capacity is taken - /// to be the next `power_of_two`. - pub(crate) fn with_capacity(capacity: usize) -> Self { - Self { - merkle_tree: MerkleTree::with_capacity(capacity), - commitments: HashMap::new(), - root_history: HashSet::new(), - } - } -} - -#[cfg_attr(test, derive(Debug, PartialEq, Eq))] -#[derive(Clone)] -struct NullifierSet(BTreeSet); - -impl NullifierSet { - const fn new() -> Self { - Self(BTreeSet::new()) - } - - fn extend(&mut self, new_nullifiers: &[Nullifier]) { - self.0.extend(new_nullifiers); - } - - fn contains(&self, nullifier: &Nullifier) -> bool { - self.0.contains(nullifier) - } -} - -impl BorshSerialize for NullifierSet { - fn serialize(&self, writer: &mut W) -> std::io::Result<()> { - self.0.iter().collect::>().serialize(writer) - } -} - -impl BorshDeserialize for NullifierSet { - fn deserialize_reader(reader: &mut R) -> std::io::Result { - let vec = Vec::::deserialize_reader(reader)?; - - let mut set = BTreeSet::new(); - for n in vec { - if !set.insert(n) { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "duplicate nullifier in NullifierSet", - )); - } - } - - Ok(Self(set)) - } -} - -#[derive(Clone, BorshSerialize, BorshDeserialize)] -#[cfg_attr(test, derive(Debug, PartialEq, Eq))] -pub struct V03State { - public_state: HashMap, - private_state: (CommitmentSet, NullifierSet), - programs: HashMap, -} - -impl Default for V03State { - fn default() -> Self { - let mut commitment_set = CommitmentSet::with_capacity(32); - commitment_set.extend(&[DUMMY_COMMITMENT]); - let nullifier_set = NullifierSet::new(); - let private_state = (commitment_set, nullifier_set); - - Self { - public_state: HashMap::default(), - private_state, - programs: HashMap::default(), - } - } -} - -impl V03State { - #[must_use] - pub fn new() -> Self { - Self::default() - } - - /// Initializes state with given public account balances leaving other account fields at their - /// default values. - #[must_use] - pub fn with_public_account_balances( - mut self, - balances: impl IntoIterator, - ) -> Self { - let public_accounts = balances.into_iter().map(|(account_id, balance)| { - ( - account_id, - Account { - balance, - ..Account::default() - }, - ) - }); - self.public_state.extend(public_accounts); - self - } - - /// Initializes state with given public accounts. - #[must_use] - pub fn with_public_accounts( - mut self, - public_accounts: impl IntoIterator, - ) -> Self { - self.public_state.extend(public_accounts); - self - } - - /// Initializes state with given private accounts. - #[must_use] - pub fn with_private_accounts( - mut self, - private_accounts: impl IntoIterator, - ) -> Self { - let (commitments, nullifiers): (Vec, Vec) = - private_accounts.into_iter().unzip(); - self.private_state.0.extend(&commitments); - self.private_state.1.extend(&nullifiers); - self - } - - /// Initializes state with given builtin programs. - #[must_use] - pub fn with_programs(mut self, programs: impl IntoIterator) -> Self { - for program in programs { - self.insert_program(program); - } - self - } - - pub(crate) fn insert_program(&mut self, program: Program) { - self.programs.insert(program.id(), program); - } - - pub fn apply_state_diff(&mut self, diff: ValidatedStateDiff) { - let StateDiff { - signer_account_ids, - public_diff, - new_commitments, - new_nullifiers, - program, - } = diff.into_state_diff(); - #[expect( - clippy::iter_over_hash_type, - reason = "Iteration order doesn't matter here" - )] - for (account_id, account) in public_diff { - *self.get_account_by_id_mut(account_id) = account; - } - for account_id in signer_account_ids { - self.get_account_by_id_mut(account_id) - .nonce - .public_account_nonce_increment(); - } - self.private_state.0.extend(&new_commitments); - self.private_state.1.extend(&new_nullifiers); - if let Some(program) = program { - self.insert_program(program); - } - } - - pub fn transition_from_public_transaction( - &mut self, - tx: &PublicTransaction, - block_id: BlockId, - timestamp: Timestamp, - ) -> Result<(), LeeError> { - let diff = ValidatedStateDiff::from_public_transaction(tx, self, block_id, timestamp)?; - self.apply_state_diff(diff); - Ok(()) - } - - pub fn transition_from_privacy_preserving_transaction( - &mut self, - tx: &PrivacyPreservingTransaction, - block_id: BlockId, - timestamp: Timestamp, - ) -> Result<(), LeeError> { - let diff = - ValidatedStateDiff::from_privacy_preserving_transaction(tx, self, block_id, timestamp)?; - self.apply_state_diff(diff); - Ok(()) - } - - pub fn transition_from_program_deployment_transaction( - &mut self, - tx: &ProgramDeploymentTransaction, - ) -> Result<(), LeeError> { - let diff = ValidatedStateDiff::from_program_deployment_transaction(tx, self)?; - self.apply_state_diff(diff); - Ok(()) - } - - fn get_account_by_id_mut(&mut self, account_id: AccountId) -> &mut Account { - self.public_state.entry(account_id).or_default() - } - - #[must_use] - pub fn get_account_by_id(&self, account_id: AccountId) -> Account { - self.public_state - .get(&account_id) - .cloned() - .unwrap_or_else(Account::default) - } - - #[must_use] - pub fn get_proof_for_commitment(&self, commitment: &Commitment) -> Option { - self.private_state.0.get_proof_for(commitment) - } - - pub(crate) const fn programs(&self) -> &HashMap { - &self.programs - } - - #[must_use] - pub fn commitment_set_digest(&self) -> CommitmentSetDigest { - self.private_state.0.digest() - } - - pub(crate) fn check_commitments_are_new( - &self, - new_commitments: &[Commitment], - ) -> Result<(), LeeError> { - for commitment in new_commitments { - if self.private_state.0.contains(commitment) { - return Err(LeeError::InvalidInput("Commitment already seen".to_owned())); - } - } - Ok(()) - } - - pub(crate) fn check_nullifiers_are_valid( - &self, - new_nullifiers: &[(Nullifier, CommitmentSetDigest)], - ) -> Result<(), LeeError> { - for (nullifier, digest) in new_nullifiers { - if self.private_state.1.contains(nullifier) { - return Err(LeeError::InvalidInput("Nullifier already seen".to_owned())); - } - if !self.private_state.0.root_history.contains(digest) { - return Err(LeeError::InvalidInput( - "Unrecognized commitment set digest".to_owned(), - )); - } - } - Ok(()) - } -} - -#[cfg(any(test, feature = "test-utils"))] -impl V03State { - pub fn force_insert_account(&mut self, account_id: AccountId, account: Account) { - self.public_state.insert(account_id, account); - } -} - -#[cfg(test)] -pub mod tests { - #![expect( - clippy::arithmetic_side_effects, - clippy::shadow_unrelated, - reason = "We don't care about it in tests" - )] - - use std::collections::HashMap; - - use lee_core::{ - BlockId, Commitment, EncryptedAccountData, InputAccountIdentity, Nullifier, - NullifierPublicKey, NullifierSecretKey, SharedSecretKey, Timestamp, - account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data}, - encryption::{EphemeralPublicKey, ViewingPublicKey}, - program::{ - BlockValidityWindow, ExecutionValidationError, MAX_NUMBER_CHAINED_CALLS, PdaSeed, - ProgramId, TimestampValidityWindow, WrappedBalanceSum, - }, - }; - - use crate::{ - PublicKey, PublicTransaction, V03State, - error::{InvalidProgramBehaviorError, LeeError}, - execute_and_prove, - privacy_preserving_transaction::{ - PrivacyPreservingTransaction, - circuit::{self, ProgramWithDependencies}, - message::Message, - witness_set::WitnessSet, - }, - program::Program, - public_transaction, - signature::PrivateKey, - }; - - impl V03State { - /// Include test programs in the builtin programs map. - #[must_use] - pub fn with_test_programs(mut self) -> Self { - self.insert_program(crate::test_methods::simple_balance_transfer()); - self.insert_program(crate::test_methods::nonce_changer()); - self.insert_program(crate::test_methods::extra_output()); - self.insert_program(crate::test_methods::missing_output()); - self.insert_program(crate::test_methods::program_owner_changer()); - self.insert_program(crate::test_methods::data_changer()); - self.insert_program(crate::test_methods::minter()); - self.insert_program(crate::test_methods::burner()); - self.insert_program(crate::test_methods::auth_asserting_noop()); - self.insert_program(crate::test_methods::private_pda_delegator()); - self.insert_program(crate::test_methods::pda_claimer()); - self.insert_program(crate::test_methods::two_pda_claimer()); - self.insert_program(crate::test_methods::noop()); - self.insert_program(crate::test_methods::chain_caller()); - self.insert_program(crate::test_methods::modified_transfer_program()); - self.insert_program(crate::test_methods::malicious_authorization_changer()); - self.insert_program(crate::test_methods::validity_window()); - self.insert_program(crate::test_methods::flash_swap_initiator()); - self.insert_program(crate::test_methods::flash_swap_callback()); - self.insert_program(crate::test_methods::malicious_self_program_id()); - self.insert_program(crate::test_methods::malicious_caller_program_id()); - self.insert_program(crate::test_methods::pda_spend_proxy()); - self.insert_program(crate::test_methods::claimer()); - self.insert_program(crate::test_methods::changer_claimer()); - self.insert_program(crate::test_methods::validity_window_chain_caller()); - self.insert_program(crate::test_methods::simple_transfer_proxy()); - self.insert_program(crate::test_methods::malicious_injector()); - self.insert_program(crate::test_methods::malicious_launderer()); - self.insert_program(crate::test_methods::modified_transfer_program()); - self - } - - #[must_use] - pub fn with_non_default_accounts_but_default_program_owners(mut self) -> Self { - let account_with_default_values_except_balance = Account { - balance: 100, - ..Account::default() - }; - let account_with_default_values_except_nonce = Account { - nonce: Nonce(37), - ..Account::default() - }; - let account_with_default_values_except_data = Account { - data: vec![0xca, 0xfe].try_into().unwrap(), - ..Account::default() - }; - self.force_insert_account( - AccountId::new([255; 32]), - account_with_default_values_except_balance, - ); - self.force_insert_account( - AccountId::new([254; 32]), - account_with_default_values_except_nonce, - ); - self.force_insert_account( - AccountId::new([253; 32]), - account_with_default_values_except_data, - ); - self - } - - #[must_use] - pub fn with_account_owned_by_burner_program(mut self) -> Self { - let account = Account { - program_owner: crate::test_methods::burner().id(), - balance: 100, - ..Default::default() - }; - self.force_insert_account(AccountId::new([252; 32]), account); - self - } - - #[must_use] - pub fn with_private_account(mut self, keys: &TestPrivateKeys, account: &Account) -> Self { - let account_id = AccountId::for_regular_private_account(&keys.npk(), 0); - let commitment = Commitment::new(&account_id, account); - self.private_state.0.extend(&[commitment]); - self - } - } - - pub struct TestPublicKeys { - pub signing_key: PrivateKey, - } - - impl TestPublicKeys { - pub fn account_id(&self) -> AccountId { - AccountId::from(&PublicKey::new_from_private_key(&self.signing_key)) - } - } - - pub struct TestPrivateKeys { - pub nsk: NullifierSecretKey, - pub d: [u8; 32], - pub z: [u8; 32], - } - - impl TestPrivateKeys { - pub fn npk(&self) -> NullifierPublicKey { - NullifierPublicKey::from(&self.nsk) - } - - pub fn vpk(&self) -> ViewingPublicKey { - ViewingPublicKey::from_seed(&self.d, &self.z) - } - } - - // ── Flash Swap types (mirrors of guest types for host-side serialisation) ── - - #[derive(serde::Serialize, serde::Deserialize)] - struct CallbackInstruction { - return_funds: bool, - token_program_id: ProgramId, - amount: u128, - } - - #[derive(serde::Serialize, serde::Deserialize)] - enum FlashSwapInstruction { - Initiate { - token_program_id: ProgramId, - callback_program_id: ProgramId, - amount_out: u128, - callback_instruction_data: Vec, - }, - InvariantCheck { - min_vault_balance: u128, - }, - } - - fn public_state_from_balances( - initial_data: &[(AccountId, u128)], - ) -> HashMap { - initial_data - .iter() - .copied() - .map(|(account_id, balance)| { - ( - account_id, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance, - ..Account::default() - }, - ) - }) - .collect() - } - - fn transfer_transaction( - from: AccountId, - from_key: &PrivateKey, - from_nonce: u128, - to: AccountId, - to_key: &PrivateKey, - to_nonce: u128, - balance: u128, - ) -> PublicTransaction { - let account_ids = vec![from, to]; - let nonces = vec![Nonce(from_nonce), Nonce(to_nonce)]; - let program_id = crate::test_methods::simple_balance_transfer().id(); - let message = - public_transaction::Message::try_new(program_id, account_ids, nonces, balance).unwrap(); - let witness_set = - public_transaction::WitnessSet::for_message(&message, &[from_key, to_key]); - PublicTransaction::new(message, witness_set) - } - - fn build_flash_swap_tx( - initiator: &Program, - vault_id: AccountId, - receiver_id: AccountId, - instruction: FlashSwapInstruction, - ) -> PublicTransaction { - let message = public_transaction::Message::try_new( - initiator.id(), - vec![vault_id, receiver_id], - vec![], // no signers — vault is PDA-authorised - instruction, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - PublicTransaction::new(message, witness_set) - } - - #[test] - fn new_works() { - let key1 = PrivateKey::try_new([1; 32]).unwrap(); - let key2 = PrivateKey::try_new([2; 32]).unwrap(); - let addr1 = AccountId::from(&PublicKey::new_from_private_key(&key1)); - let addr2 = AccountId::from(&PublicKey::new_from_private_key(&key2)); - let expected_public_state = { - let mut this = HashMap::new(); - this.insert( - addr1, - Account { - balance: 100, - ..Account::default() - }, - ); - this.insert( - addr2, - Account { - balance: 151, - ..Account::default() - }, - ); - this - }; - let expected_builtin_programs = HashMap::new(); - - let state = - V03State::new().with_public_account_balances([(addr1, 100_u128), (addr2, 151_u128)]); - - assert_eq!(state.public_state, expected_public_state); - assert_eq!(state.programs, expected_builtin_programs); - } - - #[test] - fn new_includes_nullifiers_for_private_accounts() { - let keys1 = test_private_account_keys_1(); - let keys2 = test_private_account_keys_2(); - - let account = Account { - balance: 100, - ..Account::default() - }; - - let account_id1 = AccountId::for_regular_private_account(&keys1.npk(), 0); - let account_id2 = AccountId::for_regular_private_account(&keys2.npk(), 0); - - let init_commitment1 = Commitment::new(&account_id1, &account); - let init_commitment2 = Commitment::new(&account_id2, &account); - let init_nullifier1 = Nullifier::for_account_initialization(&account_id1); - let init_nullifier2 = Nullifier::for_account_initialization(&account_id2); - - let initial_private_accounts = vec![ - (init_commitment1, init_nullifier1), - (init_commitment2, init_nullifier2), - ]; - - let state = V03State::new().with_private_accounts(initial_private_accounts); - - assert!(state.private_state.1.contains(&init_nullifier1)); - assert!(state.private_state.1.contains(&init_nullifier2)); - } - - #[test] - fn insert_program() { - let mut state = V03State::new(); - let program_to_insert = crate::test_methods::simple_balance_transfer(); - let program_id = program_to_insert.id(); - assert!(!state.programs.contains_key(&program_id)); - - state.insert_program(program_to_insert); - - assert!(state.programs.contains_key(&program_id)); - } - - #[test] - fn get_account_by_account_id_non_default_account() { - let key = PrivateKey::try_new([1; 32]).unwrap(); - let account_id = AccountId::from(&PublicKey::new_from_private_key(&key)); - let initial_data = [( - account_id, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - ..Account::default() - }, - )]; - let state = V03State::new().with_public_accounts(initial_data); - let expected_account = &state.public_state[&account_id]; - - let account = state.get_account_by_id(account_id); - - assert_eq!(&account, expected_account); - } - - #[test] - fn get_account_by_account_id_default_account() { - let addr2 = AccountId::new([0; 32]); - let state = V03State::new(); - let expected_account = Account::default(); - - let account = state.get_account_by_id(addr2); - - assert_eq!(account, expected_account); - } - - #[test] - fn builtin_programs_getter() { - let state = V03State::new(); - - let builtin_programs = state.programs(); - - assert_eq!(builtin_programs, &state.programs); - } - - #[test] - fn transition_from_authenticated_transfer_program_invocation_default_account_destination() { - let key = PrivateKey::try_new([1; 32]).unwrap(); - let account_id = AccountId::from(&PublicKey::new_from_private_key(&key)); - let initial_data = [( - account_id, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - ..Account::default() - }, - )]; - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs(); - let from = account_id; - let to_key = PrivateKey::try_new([2; 32]).unwrap(); - let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); - assert_eq!(state.get_account_by_id(to), Account::default()); - let balance_to_move = 5; - - let tx = transfer_transaction(from, &key, 0, to, &to_key, 0, balance_to_move); - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - assert_eq!(state.get_account_by_id(from).balance, 95); - assert_eq!(state.get_account_by_id(to).balance, 5); - assert_eq!(state.get_account_by_id(from).nonce, Nonce(1)); - assert_eq!(state.get_account_by_id(to).nonce, Nonce(1)); - } - - #[test] - fn transition_from_authenticated_transfer_program_invocation_insuficient_balance() { - let key = PrivateKey::try_new([1; 32]).unwrap(); - let account_id = AccountId::from(&PublicKey::new_from_private_key(&key)); - let mut state = V03State::new() - .with_public_account_balances([(account_id, 100)]) - .with_test_programs(); - let from = account_id; - let from_key = key; - let to_key = PrivateKey::try_new([2; 32]).unwrap(); - let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); - let balance_to_move = 101; - assert!(state.get_account_by_id(from).balance < balance_to_move); - - let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move); - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!(result, Err(LeeError::ProgramExecutionFailed(_)))); - assert_eq!(state.get_account_by_id(from).balance, 100); - assert_eq!(state.get_account_by_id(to).balance, 0); - assert_eq!(state.get_account_by_id(from).nonce, Nonce(0)); - assert_eq!(state.get_account_by_id(to).nonce, Nonce(0)); - } - - #[test] - fn transition_from_authenticated_transfer_program_invocation_non_default_account_destination() { - let key1 = PrivateKey::try_new([1; 32]).unwrap(); - let key2 = PrivateKey::try_new([2; 32]).unwrap(); - let account_id1 = AccountId::from(&PublicKey::new_from_private_key(&key1)); - let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2)); - let initial_data = [ - ( - account_id1, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - ..Account::default() - }, - ), - ( - account_id2, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 200, - ..Account::default() - }, - ), - ]; - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs(); - let from = account_id2; - let from_key = key2; - let to = account_id1; - let to_key = key1; - assert_ne!(state.get_account_by_id(to), Account::default()); - let balance_to_move = 8; - - let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move); - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - assert_eq!(state.get_account_by_id(from).balance, 192); - assert_eq!(state.get_account_by_id(to).balance, 108); - assert_eq!(state.get_account_by_id(from).nonce, Nonce(1)); - assert_eq!(state.get_account_by_id(to).nonce, Nonce(1)); - } - - #[test] - fn transition_from_sequence_of_authenticated_transfer_program_invocations() { - let key1 = PrivateKey::try_new([8; 32]).unwrap(); - let account_id1 = AccountId::from(&PublicKey::new_from_private_key(&key1)); - let key2 = PrivateKey::try_new([2; 32]).unwrap(); - let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2)); - let initial_data = [( - account_id1, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - ..Account::default() - }, - )]; - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs(); - let key3 = PrivateKey::try_new([3; 32]).unwrap(); - let account_id3 = AccountId::from(&PublicKey::new_from_private_key(&key3)); - let balance_to_move = 5; - - let tx = transfer_transaction( - account_id1, - &key1, - 0, - account_id2, - &key2, - 0, - balance_to_move, - ); - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - let balance_to_move = 3; - let tx = transfer_transaction( - account_id2, - &key2, - 1, - account_id3, - &key3, - 0, - balance_to_move, - ); - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - assert_eq!(state.get_account_by_id(account_id1).balance, 95); - assert_eq!(state.get_account_by_id(account_id2).balance, 2); - assert_eq!(state.get_account_by_id(account_id3).balance, 3); - assert_eq!(state.get_account_by_id(account_id1).nonce, Nonce(1)); - assert_eq!(state.get_account_by_id(account_id2).nonce, Nonce(2)); - assert_eq!(state.get_account_by_id(account_id3).nonce, Nonce(1)); - } - - #[test] - fn program_should_fail_if_modifies_nonces() { - let account_id = AccountId::new([1; 32]); - let mut state = V03State::new() - .with_public_account_balances([(account_id, 100)]) - .with_test_programs(); - let account_ids = vec![account_id]; - let program_id = crate::test_methods::nonce_changer().id(); - let message = - public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::ModifiedNonce { account_id: err_account_id } - ) - )) if err_account_id == account_id - )); - } - - #[test] - fn program_should_fail_if_output_accounts_exceed_inputs() { - let mut state = V03State::new() - .with_public_account_balances([(AccountId::new([1; 32]), 0)]) - .with_test_programs(); - let account_ids = vec![AccountId::new([1; 32])]; - let program_id = crate::test_methods::extra_output().id(); - let message = - public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::MismatchedPreStatePostStateLength { - pre_state_length, - post_state_length - } - ) - )) if pre_state_length == 1 && post_state_length == 2 - )); - } - - #[test] - fn program_should_fail_with_missing_output_accounts() { - let mut state = V03State::new() - .with_public_account_balances([(AccountId::new([1; 32]), 100)]) - .with_test_programs(); - let account_ids = vec![AccountId::new([1; 32]), AccountId::new([2; 32])]; - let program_id = crate::test_methods::missing_output().id(); - let message = - public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::MismatchedPreStatePostStateLength { - pre_state_length, - post_state_length - } - ) - )) if pre_state_length == 2 && post_state_length == 1 - )); - } - - #[test] - fn program_should_fail_if_modifies_program_owner_with_only_non_default_program_owner() { - let initial_data = [( - AccountId::new([1; 32]), - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - ..Account::default() - }, - )]; - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs(); - let account_id = AccountId::new([1; 32]); - let account = state.get_account_by_id(account_id); - // Assert the target account only differs from the default account in the program owner - // field - assert_ne!(account.program_owner, Account::default().program_owner); - assert_eq!(account.balance, Account::default().balance); - assert_eq!(account.nonce, Account::default().nonce); - assert_eq!(account.data, Account::default().data); - let program_id = crate::test_methods::program_owner_changer().id(); - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } - ))) if err_account_id == account_id - )); - } - - #[test] - fn program_should_fail_if_modifies_program_owner_with_only_non_default_balance() { - let initial_data = HashMap::new(); - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs() - .with_non_default_accounts_but_default_program_owners(); - let account_id = AccountId::new([255; 32]); - let account = state.get_account_by_id(account_id); - // Assert the target account only differs from the default account in balance field - assert_eq!(account.program_owner, Account::default().program_owner); - assert_ne!(account.balance, Account::default().balance); - assert_eq!(account.nonce, Account::default().nonce); - assert_eq!(account.data, Account::default().data); - let program_id = crate::test_methods::program_owner_changer().id(); - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } - ))) if err_account_id == account_id - )); - } - - #[test] - fn program_should_fail_if_modifies_program_owner_with_only_non_default_nonce() { - let initial_data = HashMap::new(); - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs() - .with_non_default_accounts_but_default_program_owners(); - let account_id = AccountId::new([254; 32]); - let account = state.get_account_by_id(account_id); - // Assert the target account only differs from the default account in nonce field - assert_eq!(account.program_owner, Account::default().program_owner); - assert_eq!(account.balance, Account::default().balance); - assert_ne!(account.nonce, Account::default().nonce); - assert_eq!(account.data, Account::default().data); - let program_id = crate::test_methods::program_owner_changer().id(); - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } - ))) if err_account_id == account_id - )); - } - - #[test] - fn program_should_fail_if_modifies_program_owner_with_only_non_default_data() { - let initial_data = HashMap::new(); - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs() - .with_non_default_accounts_but_default_program_owners(); - let account_id = AccountId::new([253; 32]); - let account = state.get_account_by_id(account_id); - // Assert the target account only differs from the default account in data field - assert_eq!(account.program_owner, Account::default().program_owner); - assert_eq!(account.balance, Account::default().balance); - assert_eq!(account.nonce, Account::default().nonce); - assert_ne!(account.data, Account::default().data); - let program_id = crate::test_methods::program_owner_changer().id(); - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } - ))) if err_account_id == account_id - )); - } - - #[test] - fn program_should_fail_if_transfers_balance_from_non_owned_account() { - let sender_account_id = AccountId::new([1; 32]); - let receiver_account_id = AccountId::new([2; 32]); - let mut state = V03State::new() - .with_public_account_balances([(sender_account_id, 100)]) - .with_test_programs(); - let balance_to_move: u128 = 1; - let program_id = crate::test_methods::simple_balance_transfer().id(); - assert_ne!( - state.get_account_by_id(sender_account_id).program_owner, - program_id - ); - let message = public_transaction::Message::try_new( - program_id, - vec![sender_account_id, receiver_account_id], - vec![], - balance_to_move, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::UnauthorizedBalanceDecrease { account_id: err_account_id, owner_program_id, executing_program_id } - ))) if err_account_id == sender_account_id && owner_program_id != program_id && executing_program_id == program_id - )); - } - - #[test] - fn program_should_fail_if_modifies_data_of_non_owned_account() { - let initial_data = HashMap::new(); - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs() - .with_non_default_accounts_but_default_program_owners(); - let account_id = AccountId::new([255; 32]); - let program_id = crate::test_methods::data_changer().id(); - - assert_ne!(state.get_account_by_id(account_id), Account::default()); - assert_ne!( - state.get_account_by_id(account_id).program_owner, - program_id - ); - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], vec![0]) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::UnauthorizedDataModification { account_id: err_account_id, executing_program_id } - ))) if err_account_id == account_id && executing_program_id == program_id - )); - } - - #[test] - fn program_should_fail_if_does_not_preserve_total_balance_by_minting() { - let initial_data = HashMap::new(); - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs(); - let account_id = AccountId::new([1; 32]); - let program_id = crate::test_methods::minter().id(); - - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 2, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states } - ))) if total_balance_pre_states == 0.into() && total_balance_post_states == 1.into() - )); - } - - #[test] - fn program_should_fail_if_does_not_preserve_total_balance_by_burning() { - let initial_data = HashMap::new(); - let mut state = V03State::new() - .with_public_accounts(initial_data) - .with_test_programs() - .with_account_owned_by_burner_program(); - let program_id = crate::test_methods::burner().id(); - let account_id = AccountId::new([252; 32]); - assert_eq!( - state.get_account_by_id(account_id).program_owner, - program_id - ); - let balance_to_burn: u128 = 1; - assert!(state.get_account_by_id(account_id).balance > balance_to_burn); - - let message = public_transaction::Message::try_new( - program_id, - vec![account_id], - vec![], - balance_to_burn, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - let result = state.transition_from_public_transaction(&tx, 2, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states } - ))) if total_balance_pre_states == 100.into() && total_balance_post_states == 99.into() - )); - } - - fn test_public_account_keys_1() -> TestPublicKeys { - TestPublicKeys { - signing_key: PrivateKey::try_new([37; 32]).unwrap(), - } - } - - fn test_public_account_keys_2() -> TestPublicKeys { - TestPublicKeys { - signing_key: PrivateKey::try_new([38; 32]).unwrap(), - } - } - - pub fn test_private_account_keys_1() -> TestPrivateKeys { - TestPrivateKeys { - nsk: [13; 32], - d: [31; 32], - z: [32; 32], - } - } - - pub fn test_private_account_keys_2() -> TestPrivateKeys { - TestPrivateKeys { - nsk: [38; 32], - d: [83; 32], - z: [84; 32], - } - } - - fn shielded_balance_transfer_for_tests( - sender_keys: &TestPublicKeys, - recipient_keys: &TestPrivateKeys, - balance_to_move: u128, - state: &V03State, - ) -> PrivacyPreservingTransaction { - let sender = AccountWithMetadata::new( - state.get_account_by_id(sender_keys.account_id()), - true, - sender_keys.account_id(), - ); - - let sender_nonce = sender.account.nonce; - - let recipient = - AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0)); - - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0); - - let (output, proof) = circuit::execute_and_prove( - vec![sender, recipient], - Program::serialize_instruction(balance_to_move).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivateUnauthorized { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: shared_secret, - identifier: 0, - }, - ], - &crate::test_methods::simple_balance_transfer().into(), - ) - .unwrap(); - - let message = Message::try_from_circuit_output( - vec![sender_keys.account_id()], - vec![sender_nonce], - output, - ) - .unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[&sender_keys.signing_key]); - PrivacyPreservingTransaction::new(message, witness_set) - } - - fn private_balance_transfer_for_tests( - sender_keys: &TestPrivateKeys, - sender_private_account: &Account, - recipient_keys: &TestPrivateKeys, - balance_to_move: u128, - state: &V03State, - ) -> PrivacyPreservingTransaction { - let program = crate::test_methods::simple_balance_transfer(); - let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); - let sender_commitment = Commitment::new(&sender_account_id, sender_private_account); - let sender_pre = AccountWithMetadata::new( - sender_private_account.clone(), - true, - (&sender_keys.npk(), 0), - ); - let recipient_pre = - AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0)); - - let (shared_secret_1, epk_1) = - SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0); - - let (shared_secret_2, epk_2) = - SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1); - - let (output, proof) = circuit::execute_and_prove( - vec![sender_pre, recipient_pre], - Program::serialize_instruction(balance_to_move).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: epk_1, - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: shared_secret_1, - nsk: sender_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&sender_commitment) - .expect("sender's commitment must be in state"), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: epk_2, - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: shared_secret_2, - identifier: 0, - }, - ], - &program.into(), - ) - .unwrap(); - - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - - PrivacyPreservingTransaction::new(message, witness_set) - } - - fn deshielded_balance_transfer_for_tests( - sender_keys: &TestPrivateKeys, - sender_private_account: &Account, - recipient_account_id: &AccountId, - balance_to_move: u128, - state: &V03State, - ) -> PrivacyPreservingTransaction { - let program = crate::test_methods::simple_balance_transfer(); - let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); - let sender_commitment = Commitment::new(&sender_account_id, sender_private_account); - let sender_pre = AccountWithMetadata::new( - sender_private_account.clone(), - true, - (&sender_keys.npk(), 0), - ); - let recipient_pre = AccountWithMetadata::new( - state.get_account_by_id(*recipient_account_id), - false, - *recipient_account_id, - ); - - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0); - - let (output, proof) = circuit::execute_and_prove( - vec![sender_pre, recipient_pre], - Program::serialize_instruction(balance_to_move).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: shared_secret, - nsk: sender_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&sender_commitment) - .expect("sender's commitment must be in state"), - identifier: 0, - }, - InputAccountIdentity::Public, - ], - &program.into(), - ) - .unwrap(); - - let message = - Message::try_from_circuit_output(vec![*recipient_account_id], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - - PrivacyPreservingTransaction::new(message, witness_set) - } - - #[test] - fn transition_from_privacy_preserving_transaction_shielded() { - let sender_keys = test_public_account_keys_1(); - let recipient_keys = test_private_account_keys_1(); - - let mut state = V03State::new().with_public_accounts([( - sender_keys.account_id(), - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 200, - ..Account::default() - }, - )]); - - let balance_to_move = 37; - - let tx = shielded_balance_transfer_for_tests( - &sender_keys, - &recipient_keys, - balance_to_move, - &state, - ); - - let expected_sender_post = { - let mut this = state.get_account_by_id(sender_keys.account_id()); - this.balance -= balance_to_move; - this.nonce.public_account_nonce_increment(); - this - }; - - let [expected_new_commitment] = tx.message().new_commitments.clone().try_into().unwrap(); - assert!(!state.private_state.0.contains(&expected_new_commitment)); - - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .unwrap(); - - let sender_post = state.get_account_by_id(sender_keys.account_id()); - assert_eq!(sender_post, expected_sender_post); - assert!(state.private_state.0.contains(&expected_new_commitment)); - - assert_eq!( - state.get_account_by_id(sender_keys.account_id()).balance, - 200 - balance_to_move - ); - } - - #[test] - fn transition_from_privacy_preserving_transaction_private() { - let sender_keys = test_private_account_keys_1(); - let sender_nonce = Nonce(0xdead_beef); - - let sender_private_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - nonce: sender_nonce, - data: Data::default(), - }; - let recipient_keys = test_private_account_keys_2(); - - let mut state = V03State::new().with_private_account(&sender_keys, &sender_private_account); - - let balance_to_move = 37; - - let tx = private_balance_transfer_for_tests( - &sender_keys, - &sender_private_account, - &recipient_keys, - balance_to_move, - &state, - ); - - let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); - let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); - let expected_new_commitment_1 = Commitment::new( - &sender_account_id, - &Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk), - balance: sender_private_account.balance - balance_to_move, - data: Data::default(), - }, - ); - - let sender_pre_commitment = Commitment::new(&sender_account_id, &sender_private_account); - let expected_new_nullifier = - Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk); - - let expected_new_commitment_2 = Commitment::new( - &recipient_account_id, - &Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - nonce: Nonce::private_account_nonce_init(&recipient_account_id), - balance: balance_to_move, - ..Account::default() - }, - ); - - let previous_public_state = state.public_state.clone(); - assert!(state.private_state.0.contains(&sender_pre_commitment)); - assert!(!state.private_state.0.contains(&expected_new_commitment_1)); - assert!(!state.private_state.0.contains(&expected_new_commitment_2)); - assert!(!state.private_state.1.contains(&expected_new_nullifier)); - - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .unwrap(); - - assert_eq!(state.public_state, previous_public_state); - assert!(state.private_state.0.contains(&sender_pre_commitment)); - assert!(state.private_state.0.contains(&expected_new_commitment_1)); - assert!(state.private_state.0.contains(&expected_new_commitment_2)); - assert!(state.private_state.1.contains(&expected_new_nullifier)); - } - - fn valid_private_transfer_tx_and_state() -> (V03State, PrivacyPreservingTransaction) { - let sender_keys = test_private_account_keys_1(); - let sender_private_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - nonce: Nonce(0xdead_beef), - ..Account::default() - }; - let recipient_keys = test_private_account_keys_2(); - let state = V03State::new().with_private_account(&sender_keys, &sender_private_account); - let tx = private_balance_transfer_for_tests( - &sender_keys, - &sender_private_account, - &recipient_keys, - 37, - &state, - ); - (state, tx) - } - - /// After a valid fully-private tx is proven, tampering with a note's epk should - /// make the shielding proof invalid. - #[test] - fn privacy_tampered_epk_is_rejected() { - use crate::validated_state_diff::ValidatedStateDiff; - - let (state, mut tx) = valid_private_transfer_tx_and_state(); - - // Baseline: the untampered tx verifies - assert!( - ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0).is_ok(), - "the unmodified private transfer must verify" - ); - - // Flip a byte of the first note's epk - tx.message.encrypted_private_post_states[0].epk.0[0] ^= 0xFF; - - assert!( - matches!( - ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0), - Err(LeeError::InvalidPrivacyPreservingProof) - ), - "a tampered epk must be rejected by proof verification" - ); - } - - /// After a valid fully-private tx is proven, tampering with a note's view tag should - /// make the shielding proof invalid. - #[test] - fn privacy_tampered_view_tag_is_rejected() { - use crate::validated_state_diff::ValidatedStateDiff; - - let (state, mut tx) = valid_private_transfer_tx_and_state(); - - // Baseline: the untampered tx verifies. - assert!( - ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0).is_ok(), - "the unmodified private transfer must verify" - ); - - // Flip the first note's view_tag - tx.message.encrypted_private_post_states[0].view_tag ^= 0xFF; - - assert!( - matches!( - ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0), - Err(LeeError::InvalidPrivacyPreservingProof) - ), - "a tampered view_tag must be rejected by proof verification" - ); - } - - #[test] - fn transition_from_privacy_preserving_transaction_deshielded() { - let sender_keys = test_private_account_keys_1(); - let sender_nonce = Nonce(0xdead_beef); - - let sender_private_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - nonce: sender_nonce, - data: Data::default(), - }; - let recipient_keys = test_public_account_keys_1(); - let recipient_initial_balance = 400; - let mut state = V03State::new() - .with_public_accounts([( - recipient_keys.account_id(), - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: recipient_initial_balance, - ..Account::default() - }, - )]) - .with_private_account(&sender_keys, &sender_private_account); - - let balance_to_move = 37; - - let expected_recipient_post = { - let mut this = state.get_account_by_id(recipient_keys.account_id()); - this.balance += balance_to_move; - this - }; - - let tx = deshielded_balance_transfer_for_tests( - &sender_keys, - &sender_private_account, - &recipient_keys.account_id(), - balance_to_move, - &state, - ); - - let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); - let expected_new_commitment = Commitment::new( - &sender_account_id, - &Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk), - balance: sender_private_account.balance - balance_to_move, - data: Data::default(), - }, - ); - - let sender_pre_commitment = Commitment::new(&sender_account_id, &sender_private_account); - let expected_new_nullifier = - Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk); - - assert!(state.private_state.0.contains(&sender_pre_commitment)); - assert!(!state.private_state.0.contains(&expected_new_commitment)); - assert!(!state.private_state.1.contains(&expected_new_nullifier)); - - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .unwrap(); - - let recipient_post = state.get_account_by_id(recipient_keys.account_id()); - assert_eq!(recipient_post, expected_recipient_post); - assert!(state.private_state.0.contains(&sender_pre_commitment)); - assert!(state.private_state.0.contains(&expected_new_commitment)); - assert!(state.private_state.1.contains(&expected_new_nullifier)); - assert_eq!( - state.get_account_by_id(recipient_keys.account_id()).balance, - recipient_initial_balance + balance_to_move - ); - } - - #[test] - fn burner_program_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::burner(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(10_u128).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn minter_program_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::minter(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(10_u128).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn nonce_changer_program_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::nonce_changer(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn data_changer_program_should_fail_for_non_owned_account_in_privacy_preserving_circuit() { - let program = crate::test_methods::data_changer(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: [0, 1, 2, 3, 4, 5, 6, 7], - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(vec![0]).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn data_changer_program_should_fail_for_too_large_data_in_privacy_preserving_circuit() { - let program = crate::test_methods::data_changer(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let large_data: Vec = - vec![ - 0; - usize::try_from(lee_core::account::data::DATA_MAX_LENGTH.as_u64()) - .expect("DATA_MAX_LENGTH fits in usize") - + 1 - ]; - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(large_data).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::ProgramProveFailed(_)))); - } - - #[test] - fn extra_output_program_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::extra_output(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn missing_output_program_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::missing_output(); - let public_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - let public_account_2 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([1; 32]), - ); - - let result = execute_and_prove( - vec![public_account_1, public_account_2], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::Public, InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn program_owner_changer_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::program_owner_changer(); - let public_account = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn transfer_from_non_owned_account_should_fail_in_privacy_preserving_circuit() { - let program = crate::test_methods::simple_balance_transfer(); - let public_account_1 = AccountWithMetadata::new( - Account { - program_owner: [0, 1, 2, 3, 4, 5, 6, 7], - balance: 100, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - let public_account_2 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([1; 32]), - ); - - let result = execute_and_prove( - vec![public_account_1, public_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![InputAccountIdentity::Public, InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_fails_if_visibility_masks_have_incorrect_lenght() { - let program = crate::test_methods::simple_balance_transfer(); - let public_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - let public_account_2 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 0, - ..Account::default() - }, - true, - AccountId::new([1; 32]), - ); - - // Single account_identity entry for a circuit execution with two pre_state accounts. - let result = execute_and_prove( - vec![public_account_1, public_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_fails_if_invalid_auth_keys_are_provided() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - let private_account_2 = - AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0)); - - // Setting the recipient nsk to authorize the sender. - // This should be set to the sender private account in a normal circumstance. - // `PrivateAuthorizedUpdate` derives npk from nsk and asserts equality with - // `pre_state.account_id`, so a mismatched nsk fails that check. - let result = execute_and_prove( - vec![private_account_1, private_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic( - &sender_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - nsk: recipient_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: SharedSecretKey::encapsulate_deterministic( - &recipient_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_should_fail_if_new_private_account_with_non_default_balance_is_provided() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - let private_account_2 = AccountWithMetadata::new( - Account { - // Non default balance - balance: 1, - ..Account::default() - }, - false, - (&recipient_keys.npk(), 0), - ); - - let result = execute_and_prove( - vec![private_account_1, private_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic( - &sender_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: SharedSecretKey::encapsulate_deterministic( - &recipient_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_should_fail_if_new_private_account_with_non_default_program_owner_is_provided() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - let private_account_2 = AccountWithMetadata::new( - Account { - // Non default program_owner - program_owner: [0, 1, 2, 3, 4, 5, 6, 7], - ..Account::default() - }, - false, - (&recipient_keys.npk(), 0), - ); - - let result = execute_and_prove( - vec![private_account_1, private_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic( - &sender_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: SharedSecretKey::encapsulate_deterministic( - &recipient_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_should_fail_if_new_private_account_with_non_default_data_is_provided() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - let private_account_2 = AccountWithMetadata::new( - Account { - // Non default data - data: b"hola mundo".to_vec().try_into().unwrap(), - ..Account::default() - }, - false, - (&recipient_keys.npk(), 0), - ); - - let result = execute_and_prove( - vec![private_account_1, private_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic( - &sender_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: SharedSecretKey::encapsulate_deterministic( - &recipient_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_should_fail_if_new_private_account_with_non_default_nonce_is_provided() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - let private_account_2 = AccountWithMetadata::new( - Account { - // Non default nonce - nonce: Nonce(0xdead_beef), - ..Account::default() - }, - false, - (&recipient_keys.npk(), 0), - ); - - let result = execute_and_prove( - vec![private_account_1, private_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic( - &sender_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: SharedSecretKey::encapsulate_deterministic( - &recipient_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn circuit_should_fail_if_new_private_account_is_provided_with_default_values_but_marked_as_authorized() - { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let recipient_keys = test_private_account_keys_2(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - let private_account_2 = AccountWithMetadata::new( - Account::default(), - // This should be set to false in normal circumstances - true, - (&recipient_keys.npk(), 0), - ); - - let result = execute_and_prove( - vec![private_account_1, private_account_2], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic( - &sender_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateUnauthorized { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - npk: recipient_keys.npk(), - ssk: SharedSecretKey::encapsulate_deterministic( - &recipient_keys.vpk(), - &[0_u8; 32], - 0, - ) - .0, - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - /// A private PDA account that no program claims via `Claim::Pda` and no caller authorizes via - /// `ChainedCall.pda_seeds` has no binding between its supplied npk and its `account_id`, - /// so the circuit must reject. Here `simple_balance_transfer` emits no claim for the - /// second account, leaving position 1 unbound. - #[test] - fn private_pda_without_binding_fails() { - let program = crate::test_methods::simple_balance_transfer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - let public_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - AccountId::new([0; 32]), - ); - let private_pda_account = - AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32])); - - let result = execute_and_prove( - vec![public_account_1, private_pda_account], - Program::serialize_instruction(10_u128).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier: u128::MAX, - seed: None, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - /// Happy path: a program claims a new private PDA via `Claim::Pda(seed)`. The circuit - /// reads the npk for that `pre_state` from `private_account_keys` at the `pre_state`'s - /// position, derives `AccountId` via `AccountId::for_private_pda(program_id, seed, npk)`, and - /// asserts it equals the `pre_state`'s `account_id`. The equality both validates the claim - /// and binds the supplied npk to the `account_id`. - #[test] - fn private_pda_claim_succeeds() { - let program = crate::test_methods::pda_claimer(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([42; 32]); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, u128::MAX); - let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); - - let result = execute_and_prove( - vec![pre_state], - Program::serialize_instruction(seed).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier: u128::MAX, - seed: None, - }], - &program.into(), - ); - - let (output, _proof) = result.expect("private PDA claim should succeed"); - assert_eq!(output.new_nullifiers.len(), 1); - assert_eq!(output.new_commitments.len(), 1); - assert_eq!(output.encrypted_private_post_states.len(), 1); - assert!(output.public_pre_states.is_empty()); - assert!(output.public_post_states.is_empty()); - } - - /// An npk is supplied that does not match the `pre_state`'s `account_id` under - /// `AccountId::for_private_pda(program, claim_seed, npk)`. The claim equality check rejects. - #[test] - fn private_pda_npk_mismatch_fails() { - // `keys_a` produces the `pre_state`'s `account_id` (the registered pair), `keys_b` is - // the mismatched pair supplied in `private_account_keys` for that pre_state. - let program = crate::test_methods::pda_claimer(); - let keys_a = test_private_account_keys_1(); - let keys_b = test_private_account_keys_2(); - let npk_a = keys_a.npk(); - let npk_b = keys_b.npk(); - let seed = PdaSeed::new([42; 32]); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys_b.vpk(), &[0_u8; 32], 0).0; - - // `account_id` is derived from `npk_a`, but `npk_b` is supplied for this pre_state. - // `AccountId::for_private_pda(program, seed, npk_b) != account_id`, so the claim check in - // the circuit must reject. - let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk_a, u128::MAX); - let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); - - let result = execute_and_prove( - vec![pre_state], - Program::serialize_instruction(seed).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk_b, &keys_b.vpk()), - npk: npk_b, - ssk: shared_secret, - identifier: u128::MAX, - seed: None, - }], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - /// Happy path for the caller-seeds authorization of a private PDA. The delegator claims a - /// private PDA via `Claim::Pda(seed)`, then chains to a callee (`noop`) delegating the same - /// seed via `ChainedCall.pda_seeds`. In the callee's step, the `pre_state`'s authorization - /// is established via the private derivation - /// `AccountId::for_private_pda(delegator, seed, npk) == pre.account_id`. - #[test] - fn caller_pda_seeds_authorize_private_pda_for_callee() { - let delegator = crate::test_methods::private_pda_delegator(); - let callee = crate::test_methods::auth_asserting_noop(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let seed = PdaSeed::new([77; 32]); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let account_id = AccountId::for_private_pda(&delegator.id(), &seed, &npk, u128::MAX); - let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); - - let callee_id = callee.id(); - let program_with_deps = - ProgramWithDependencies::new(delegator, [(callee_id, callee)].into()); - - let result = execute_and_prove( - vec![pre_state], - Program::serialize_instruction((seed, seed, callee_id)).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier: u128::MAX, - seed: None, - }], - &program_with_deps, - ); - - let (output, _proof) = - result.expect("caller-seeds authorization of private PDA should succeed"); - assert_eq!(output.new_commitments.len(), 1); - assert_eq!(output.new_nullifiers.len(), 1); - } - - /// The delegator chains with a different seed than the one it claimed with. In the callee - /// step, neither public nor private caller-seeds authorization matches; `pre.is_authorized` - /// was set to `true` by the delegator but no proven source supports it, so the consistency - /// assertion rejects. - #[test] - fn caller_pda_seeds_with_wrong_seed_rejects_private_pda_for_callee() { - let delegator = crate::test_methods::private_pda_delegator(); - let callee = crate::test_methods::auth_asserting_noop(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let claim_seed = PdaSeed::new([77; 32]); - let wrong_delegated_seed = PdaSeed::new([88; 32]); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - - let account_id = AccountId::for_private_pda(&delegator.id(), &claim_seed, &npk, u128::MAX); - let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); - - let callee_id = callee.id(); - let program_with_deps = - ProgramWithDependencies::new(delegator, [(callee_id, callee)].into()); - - let result = execute_and_prove( - vec![pre_state], - Program::serialize_instruction((claim_seed, wrong_delegated_seed, callee_id)).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier: u128::MAX, - seed: None, - }], - &program_with_deps, - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - /// Exploit-scenario pin. A single `(program_id, seed)` pair can derive a family of - /// `AccountId`s, one public PDA and one private PDA per distinct npk. Without the tx-wide - /// family-binding check, a program could claim `PDA_alice` (`alice_npk`) and - /// `PDA_bob` (`bob_npk`) under the same seed in one transaction, and once reuse - /// is supported a later chained call could delegate both to a callee via - /// `pda_seeds: [S]` and mix balances across them. The binding check rejects the setup - /// here: after the first claim records `(program, seed) → PDA_alice`, the second claim - /// tries to record `(program, seed) → PDA_bob` and panics. - #[test] - fn two_private_pda_claims_under_same_seed_are_rejected() { - let program = crate::test_methods::two_pda_claimer(); - let keys_a = test_private_account_keys_1(); - let keys_b = test_private_account_keys_2(); - let seed = PdaSeed::new([55; 32]); - let shared_a = SharedSecretKey::encapsulate_deterministic(&keys_a.vpk(), &[0_u8; 32], 0).0; - let shared_b = SharedSecretKey::encapsulate_deterministic(&keys_b.vpk(), &[0_u8; 32], 0).0; - - let account_a = AccountId::for_private_pda(&program.id(), &seed, &keys_a.npk(), u128::MAX); - let account_b = AccountId::for_private_pda(&program.id(), &seed, &keys_b.npk(), u128::MAX); - - let pre_a = AccountWithMetadata::new(Account::default(), false, account_a); - let pre_b = AccountWithMetadata::new(Account::default(), false, account_b); - - let result = execute_and_prove( - vec![pre_a, pre_b], - Program::serialize_instruction(seed).unwrap(), - vec![ - InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&keys_a.npk(), &keys_a.vpk()), - npk: keys_a.npk(), - ssk: shared_a, - identifier: u128::MAX, - seed: None, - }, - InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&keys_b.npk(), &keys_b.vpk()), - npk: keys_b.npk(), - ssk: shared_b, - identifier: u128::MAX, - seed: None, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - /// A private PDA that is reused at top level without an external seed in the identity still - /// fails binding. The noop program emits no `Claim::Pda` and there is no caller - /// `ChainedCall.pda_seeds`, so position 0 is never bound and the assertion fires. - /// Supplying `seed: Some((seed, owner_program_id))` in the `PrivatePdaUpdate` identity is - /// the correct path for top-level reuse; this test pins the failure when no seed is provided. - #[test] - fn private_pda_top_level_reuse_rejected_by_binding_check() { - let program = crate::test_methods::noop(); - let keys = test_private_account_keys_1(); - let npk = keys.npk(); - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; - let seed = PdaSeed::new([99; 32]); - - // Simulate a previously-claimed private PDA: program_owner != DEFAULT, is_authorized = - // true, account_id derived via the private formula. - let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, u128::MAX); - let owned_pre_state = AccountWithMetadata::new( - Account { - program_owner: program.id(), - ..Account::default() - }, - true, - account_id, - ); - - let result = execute_and_prove( - vec![owned_pre_state], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::PrivatePdaInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), - npk, - ssk: shared_secret, - identifier: u128::MAX, - seed: None, - }], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn private_accounts_can_only_be_initialized_once() { - let sender_keys = test_private_account_keys_1(); - let sender_nonce = Nonce(0xdead_beef); - - let sender_private_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - nonce: sender_nonce, - data: Data::default(), - }; - let recipient_keys = test_private_account_keys_2(); - - let mut state = V03State::new().with_private_account(&sender_keys, &sender_private_account); - - let balance_to_move = 37; - let balance_to_move_2 = 30; - - let tx = private_balance_transfer_for_tests( - &sender_keys, - &sender_private_account, - &recipient_keys, - balance_to_move, - &state, - ); - - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .unwrap(); - - let sender_private_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - nonce: sender_nonce, - data: Data::default(), - }; - - let tx = private_balance_transfer_for_tests( - &sender_keys, - &sender_private_account, - &recipient_keys, - balance_to_move_2, - &state, - ); - - let result = state.transition_from_privacy_preserving_transaction(&tx, 1, 0); - - assert!(matches!(result, Err(LeeError::InvalidInput(_)))); - let LeeError::InvalidInput(error_message) = result.err().unwrap() else { - panic!("Incorrect message error"); - }; - let expected_error_message = "Nullifier already seen".to_owned(); - assert_eq!(error_message, expected_error_message); - } - - #[test] - fn circuit_should_fail_if_there_are_repeated_ids() { - let program = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_private_account_keys_1(); - let private_account_1 = AccountWithMetadata::new( - Account { - program_owner: program.id(), - balance: 100, - ..Account::default() - }, - true, - (&sender_keys.npk(), 0), - ); - - let shared_secret = - SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0; - let result = execute_and_prove( - vec![private_account_1.clone(), private_account_1], - Program::serialize_instruction(100_u128).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: shared_secret, - nsk: sender_keys.nsk, - membership_proof: (1, vec![]), - identifier: 0, - }, - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: shared_secret, - nsk: sender_keys.nsk, - membership_proof: (1, vec![]), - identifier: 0, - }, - ], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn claiming_mechanism() { - let program = crate::test_methods::simple_balance_transfer(); - let from_key = PrivateKey::try_new([1; 32]).unwrap(); - let from = AccountId::from(&PublicKey::new_from_private_key(&from_key)); - let initial_balance = 100; - let initial_data = [(from, initial_balance)]; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let to_key = PrivateKey::try_new([2; 32]).unwrap(); - let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); - let amount: u128 = 37; - - // Check the recipient is an uninitialized account - assert_eq!(state.get_account_by_id(to), Account::default()); - - let expected_recipient_post = Account { - program_owner: program.id(), - balance: amount, - nonce: Nonce(1), - ..Account::default() - }; - - let message = public_transaction::Message::try_new( - program.id(), - vec![from, to], - vec![Nonce(0), Nonce(0)], - amount, - ) - .unwrap(); - let witness_set = - public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]); - let tx = PublicTransaction::new(message, witness_set); - - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - let recipient_post = state.get_account_by_id(to); - - assert_eq!(recipient_post, expected_recipient_post); - } - - #[test] - fn unauthorized_public_account_claiming_fails() { - let program = crate::test_methods::simple_balance_transfer(); - let account_key = PrivateKey::try_new([9; 32]).unwrap(); - let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key)); - let mut state = V03State::new().with_test_programs(); - - assert_eq!(state.get_account_by_id(account_id), Account::default()); - - let message = - public_transaction::Message::try_new(program.id(), vec![account_id], vec![], 0_u128) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 2, 0); - - assert!(matches!(result, Err(LeeError::InvalidProgramBehavior(_)))); - assert_eq!(state.get_account_by_id(account_id), Account::default()); - } - - #[test] - fn authorized_public_account_claiming_succeeds() { - let program = crate::test_methods::simple_balance_transfer(); - let account_key = PrivateKey::try_new([10; 32]).unwrap(); - let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key)); - let mut state = V03State::new().with_test_programs(); - - assert_eq!(state.get_account_by_id(account_id), Account::default()); - - let message = public_transaction::Message::try_new( - program.id(), - vec![account_id], - vec![Nonce(0)], - 0_u128, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[&account_key]); - let tx = PublicTransaction::new(message, witness_set); - - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - assert_eq!( - state.get_account_by_id(account_id), - Account { - program_owner: program.id(), - nonce: Nonce(1), - ..Account::default() - } - ); - } - - #[test] - fn public_chained_call() { - let program = crate::test_methods::chain_caller(); - let key = PrivateKey::try_new([1; 32]).unwrap(); - let from = AccountId::from(&PublicKey::new_from_private_key(&key)); - let to = AccountId::new([2; 32]); - let initial_balance = 1000; - let initial_data = [(from, initial_balance), (to, 0)]; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let from_key = key; - let amount: u128 = 37; - let instruction: (u128, ProgramId, u32, Option) = ( - amount, - crate::test_methods::simple_balance_transfer().id(), - 2, - None, - ); - - let expected_to_post = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: amount * 2, // The `chain_caller` chains the program twice - ..Account::default() - }; - - let message = public_transaction::Message::try_new( - program.id(), - vec![to, from], // The chain_caller program permutes the account order in the chain - // call - vec![Nonce(0)], - instruction, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]); - let tx = PublicTransaction::new(message, witness_set); - - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - let from_post = state.get_account_by_id(from); - let to_post = state.get_account_by_id(to); - // The `chain_caller` program calls the program twice - assert_eq!(from_post.balance, initial_balance - 2 * amount); - assert_eq!(to_post, expected_to_post); - } - - #[test] - fn execution_fails_if_chained_calls_exceeds_depth() { - let program = crate::test_methods::chain_caller(); - let key = PrivateKey::try_new([1; 32]).unwrap(); - let from = AccountId::from(&PublicKey::new_from_private_key(&key)); - let to = AccountId::new([2; 32]); - let initial_balance = 100; - let initial_data = [(from, initial_balance), (to, 0)]; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let from_key = key; - let amount: u128 = 0; - let instruction: (u128, ProgramId, u32, Option) = ( - amount, - crate::test_methods::simple_balance_transfer().id(), - u32::try_from(MAX_NUMBER_CHAINED_CALLS).expect("MAX_NUMBER_CHAINED_CALLS fits in u32") - + 1, - None, - ); - - let message = public_transaction::Message::try_new( - program.id(), - vec![to, from], // The chain_caller program permutes the account order in the chain - // call - vec![Nonce(0)], - instruction, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - assert!(matches!( - result, - Err(LeeError::MaxChainedCallsDepthExceeded) - )); - } - - #[test] - fn execution_that_requires_authentication_of_a_program_derived_account_id_succeeds() { - let chain_caller = crate::test_methods::chain_caller(); - let pda_seed = PdaSeed::new([37; 32]); - let from = AccountId::for_public_pda(&chain_caller.id(), &pda_seed); - let to = AccountId::new([2; 32]); - let initial_balance = 1000; - let initial_data = [(from, initial_balance), (to, 0)]; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let amount: u128 = 58; - let instruction: (u128, ProgramId, u32, Option) = ( - amount, - crate::test_methods::simple_balance_transfer().id(), - 1, - Some(pda_seed), - ); - - let expected_to_post = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: amount, // The `chain_caller` chains the program twice - ..Account::default() - }; - let message = public_transaction::Message::try_new( - chain_caller.id(), - vec![to, from], // The chain_caller program permutes the account order in the chain - // call - vec![], - instruction, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - let from_post = state.get_account_by_id(from); - let to_post = state.get_account_by_id(to); - assert_eq!(from_post.balance, initial_balance - amount); - assert_eq!(to_post, expected_to_post); - } - - #[test] - fn claiming_mechanism_within_chain_call() { - // This test calls the authenticated transfer program through the chain_caller program. - // The transfer is made from an initialized sender to an uninitialized recipient. And - // it is expected that the recipient account is claimed by the authenticated transfer - // program and not the chained_caller program. - let chain_caller = crate::test_methods::chain_caller(); - let simple_transfer = crate::test_methods::simple_balance_transfer(); - let from_key = PrivateKey::try_new([1; 32]).unwrap(); - let from = AccountId::from(&PublicKey::new_from_private_key(&from_key)); - let initial_balance = 100; - let initial_data = [(from, initial_balance)]; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let to_key = PrivateKey::try_new([2; 32]).unwrap(); - let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); - let amount: u128 = 37; - - // Check the recipient is an uninitialized account - assert_eq!(state.get_account_by_id(to), Account::default()); - - let expected_to_post = Account { - // The expected program owner is the authenticated transfer program - program_owner: simple_transfer.id(), - balance: amount, - nonce: Nonce(1), - ..Account::default() - }; - - // The transaction executes the chain_caller program, which internally calls the - // authenticated_transfer program - let instruction: (u128, ProgramId, u32, Option) = ( - amount, - crate::test_methods::simple_balance_transfer().id(), - 1, - None, - ); - let message = public_transaction::Message::try_new( - chain_caller.id(), - vec![to, from], // The chain_caller program permutes the account order in the chain - // call - vec![Nonce(0), Nonce(0)], - instruction, - ) - .unwrap(); - let witness_set = - public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]); - let tx = PublicTransaction::new(message, witness_set); - - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - let from_post = state.get_account_by_id(from); - let to_post = state.get_account_by_id(to); - assert_eq!(from_post.balance, initial_balance - amount); - assert_eq!(to_post, expected_to_post); - } - - #[test] - fn unauthorized_public_account_claiming_fails_when_executed_privately() { - let program = crate::test_methods::simple_balance_transfer(); - let account_id = AccountId::new([11; 32]); - let public_account = AccountWithMetadata::new(Account::default(), false, account_id); - - let result = execute_and_prove( - vec![public_account], - Program::serialize_instruction(0_u128).unwrap(), - vec![InputAccountIdentity::Public], - &program.into(), - ); - - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn authorized_public_account_claiming_succeeds_when_executed_privately() { - let program = crate::test_methods::simple_balance_transfer(); - let program_id = program.id(); - let sender_keys = test_private_account_keys_1(); - let sender_private_account = Account { - program_owner: program_id, - balance: 100, - ..Account::default() - }; - let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); - let sender_commitment = Commitment::new(&sender_account_id, &sender_private_account); - let sender_init_nullifier = Nullifier::for_account_initialization(&sender_account_id); - let mut state = V03State::new() - .with_private_accounts([(sender_commitment.clone(), sender_init_nullifier)]); - let sender_pre = - AccountWithMetadata::new(sender_private_account, true, (&sender_keys.npk(), 0)); - let recipient_private_key = PrivateKey::try_new([2; 32]).unwrap(); - let recipient_account_id = - AccountId::from(&PublicKey::new_from_private_key(&recipient_private_key)); - let recipient_pre = - AccountWithMetadata::new(Account::default(), true, recipient_account_id); - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0); - - let balance = 37; - - let (output, proof) = execute_and_prove( - vec![sender_pre, recipient_pre], - Program::serialize_instruction(balance).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: shared_secret, - nsk: sender_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&sender_commitment) - .expect("sender's commitment must be in state"), - identifier: 0, - }, - InputAccountIdentity::Public, - ], - &program.into(), - ) - .unwrap(); - - let message = - Message::try_from_circuit_output(vec![recipient_account_id], vec![Nonce(0)], output) - .unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_private_key]); - let tx = PrivacyPreservingTransaction::new(message, witness_set); - - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .unwrap(); - - let nullifier = Nullifier::for_account_update(&sender_commitment, &sender_keys.nsk); - assert!(state.private_state.1.contains(&nullifier)); - - assert_eq!( - state.get_account_by_id(recipient_account_id), - Account { - program_owner: program_id, - balance, - nonce: Nonce(1), - ..Account::default() - } - ); - } - - #[test_case::test_case(1; "single call")] - #[test_case::test_case(2; "two calls")] - fn private_chained_call(number_of_calls: u32) { - // Arrange - let chain_caller = crate::test_methods::chain_caller(); - let simple_transfers = crate::test_methods::simple_balance_transfer(); - let from_keys = test_private_account_keys_1(); - let to_keys = test_private_account_keys_2(); - let initial_balance = 100; - let from_account = AccountWithMetadata::new( - Account { - program_owner: simple_transfers.id(), - balance: initial_balance, - ..Account::default() - }, - true, - (&from_keys.npk(), 0), - ); - let to_account = AccountWithMetadata::new( - Account { - program_owner: simple_transfers.id(), - ..Account::default() - }, - true, - (&to_keys.npk(), 0), - ); - - let from_account_id = AccountId::for_regular_private_account(&from_keys.npk(), 0); - let to_account_id = AccountId::for_regular_private_account(&to_keys.npk(), 0); - let from_commitment = Commitment::new(&from_account_id, &from_account.account); - let to_commitment = Commitment::new(&to_account_id, &to_account.account); - let from_init_nullifier = Nullifier::for_account_initialization(&from_account_id); - let to_init_nullifier = Nullifier::for_account_initialization(&to_account_id); - let mut state = V03State::new() - .with_private_accounts([ - (from_commitment.clone(), from_init_nullifier), - (to_commitment.clone(), to_init_nullifier), - ]) - .with_test_programs(); - let amount: u128 = 37; - let instruction: (u128, ProgramId, u32, Option) = ( - amount, - crate::test_methods::simple_balance_transfer().id(), - number_of_calls, - None, - ); - - let (from_ss, from_epk) = - SharedSecretKey::encapsulate_deterministic(&from_keys.vpk(), &[0_u8; 32], 0); - - let (to_ss, to_epk) = - SharedSecretKey::encapsulate_deterministic(&to_keys.vpk(), &[0_u8; 32], 1); - - let mut dependencies = HashMap::new(); - - dependencies.insert(simple_transfers.id(), simple_transfers); - let program_with_deps = ProgramWithDependencies::new(chain_caller, dependencies); - - let from_new_nonce = Nonce::default().private_account_nonce_increment(&from_keys.nsk); - let to_new_nonce = Nonce::default().private_account_nonce_increment(&to_keys.nsk); - - let from_expected_post = Account { - balance: initial_balance - u128::from(number_of_calls) * amount, - nonce: from_new_nonce, - ..from_account.account.clone() - }; - let from_expected_commitment = Commitment::new(&from_account_id, &from_expected_post); - - let to_expected_post = Account { - balance: u128::from(number_of_calls) * amount, - nonce: to_new_nonce, - ..to_account.account.clone() - }; - let to_expected_commitment = Commitment::new(&to_account_id, &to_expected_post); - - // Act - let (output, proof) = execute_and_prove( - vec![to_account, from_account], - Program::serialize_instruction(instruction).unwrap(), - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: to_epk, - view_tag: EncryptedAccountData::compute_view_tag( - &to_keys.npk(), - &to_keys.vpk(), - ), - ssk: to_ss, - nsk: from_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&from_commitment) - .expect("from's commitment must be in state"), - identifier: 0, - }, - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: from_epk, - view_tag: EncryptedAccountData::compute_view_tag( - &from_keys.npk(), - &from_keys.vpk(), - ), - ssk: from_ss, - nsk: to_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&to_commitment) - .expect("to's commitment must be in state"), - identifier: 0, - }, - ], - &program_with_deps, - ) - .unwrap(); - - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - let witness_set = WitnessSet::for_message(&message, proof, &[]); - let transaction = PrivacyPreservingTransaction::new(message, witness_set); - - state - .transition_from_privacy_preserving_transaction(&transaction, 1, 0) - .unwrap(); - - // Assert - assert!( - state - .get_proof_for_commitment(&from_expected_commitment) - .is_some() - ); - assert!( - state - .get_proof_for_commitment(&to_expected_commitment) - .is_some() - ); - } - - #[test] - fn claiming_mechanism_cannot_claim_initialied_accounts() { - let claimer = crate::test_methods::claimer(); - let mut state = V03State::new().with_test_programs(); - let account_id = AccountId::new([2; 32]); - - // Insert an account with non-default program owner - state.force_insert_account( - account_id, - Account { - program_owner: [1, 2, 3, 4, 5, 6, 7, 8], - ..Account::default() - }, - ); - - let message = - public_transaction::Message::try_new(claimer.id(), vec![account_id], vec![], ()) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id: err_account_id } - )) if err_account_id == account_id - )); - } - - /// This test ensures that even if a malicious program tries to perform overflow of balances - /// it will not be able to break the balance validation. - #[test] - fn malicious_program_cannot_break_balance_validation_if_not_in_genesis() { - let sender_key = PrivateKey::try_new([37; 32]).unwrap(); - let sender_id = AccountId::from(&PublicKey::new_from_private_key(&sender_key)); - let sender_init_balance: u128 = 10; - - let recipient_key = PrivateKey::try_new([42; 32]).unwrap(); - let recipient_id = AccountId::from(&PublicKey::new_from_private_key(&recipient_key)); - let recipient_init_balance: u128 = 10; - - let modified_transfer_id = crate::test_methods::modified_transfer_program().id(); - - let mut state = V03State::new() - .with_public_accounts([ - ( - sender_id, - Account { - program_owner: modified_transfer_id, - balance: sender_init_balance, - ..Account::default() - }, - ), - ( - recipient_id, - Account { - program_owner: modified_transfer_id, - balance: recipient_init_balance, - ..Account::default() - }, - ), - ]) - .with_test_programs(); - - let balance_to_move: u128 = 4; - - let sender = AccountWithMetadata::new(state.get_account_by_id(sender_id), true, sender_id); - - let sender_nonce = sender.account.nonce; - - let _recipient = - AccountWithMetadata::new(state.get_account_by_id(recipient_id), false, sender_id); - - let message = public_transaction::Message::try_new( - modified_transfer_id, - vec![sender_id, recipient_id], - vec![sender_nonce], - balance_to_move, - ) - .unwrap(); - - let witness_set = public_transaction::WitnessSet::for_message(&message, &[&sender_key]); - let tx = PublicTransaction::new(message, witness_set); - let res = state.transition_from_public_transaction(&tx, 2, 0); - let expected_total_balance_pre_states = WrappedBalanceSum::from_balances( - [sender_init_balance, recipient_init_balance].into_iter(), - ) - .unwrap(); - let expected_total_balance_post_states = WrappedBalanceSum::from_balances( - [sender_init_balance, recipient_init_balance, u128::MAX, 1].into_iter(), - ) - .unwrap(); - assert!(matches!( - res, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::ExecutionValidationFailed( - ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states } - ) - )) if total_balance_pre_states == expected_total_balance_pre_states && total_balance_post_states == expected_total_balance_post_states - )); - - let sender_post = state.get_account_by_id(sender_id); - let recipient_post = state.get_account_by_id(recipient_id); - - let expected_sender_post = { - let mut this = state.get_account_by_id(sender_id); - this.balance = sender_init_balance; - this.nonce = Nonce(0); - this - }; - - let expected_recipient_post = { - let mut this = state.get_account_by_id(sender_id); - this.balance = recipient_init_balance; - this.nonce = Nonce(0); - this - }; - - assert_eq!(expected_sender_post, sender_post); - assert_eq!(expected_recipient_post, recipient_post); - } - - #[test] - fn private_authorized_uninitialized_account() { - let mut state = V03State::new().with_test_programs(); - - // Set up keys for the authorized private account - let private_keys = test_private_account_keys_1(); - - // Create an authorized private account with default values (new account being initialized) - let authorized_account = - AccountWithMetadata::new(Account::default(), true, (&private_keys.npk(), 0)); - - let program = crate::test_methods::simple_balance_transfer(); - - // Set up parameters for the new account - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0); - - let instruction: u128 = 0; - - // Execute and prove the circuit with the authorized account but no commitment proof - let (output, proof) = execute_and_prove( - vec![authorized_account], - Program::serialize_instruction(instruction).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedInit { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &private_keys.npk(), - &private_keys.vpk(), - ), - ssk: shared_secret, - nsk: private_keys.nsk, - identifier: 0, - }], - &program.into(), - ) - .unwrap(); - - // Create message from circuit output - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - - let tx = PrivacyPreservingTransaction::new(message, witness_set); - let result = state.transition_from_privacy_preserving_transaction(&tx, 1, 0); - assert!(result.is_ok()); - - let account_id = AccountId::for_regular_private_account(&private_keys.npk(), 0); - let nullifier = Nullifier::for_account_initialization(&account_id); - assert!(state.private_state.1.contains(&nullifier)); - } - - #[test] - fn private_unauthorized_uninitialized_account_can_still_be_claimed() { - let mut state = V03State::new().with_test_programs(); - - let private_keys = test_private_account_keys_1(); - // This is intentional: claim authorization was introduced to protect public accounts, - // especially PDAs. Private PDAs are not useful in practice because there is no way to - // operate them without the corresponding private keys, so unauthorized private claiming - // remains allowed. - let unauthorized_account = - AccountWithMetadata::new(Account::default(), false, (&private_keys.npk(), 0)); - - let program = crate::test_methods::claimer(); - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0); - - let (output, proof) = execute_and_prove( - vec![unauthorized_account], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::PrivateUnauthorized { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &private_keys.npk(), - &private_keys.vpk(), - ), - npk: private_keys.npk(), - ssk: shared_secret, - identifier: 0, - }], - &program.into(), - ) - .unwrap(); - - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - let tx = PrivacyPreservingTransaction::new(message, witness_set); - - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .unwrap(); - - let account_id = AccountId::for_regular_private_account(&private_keys.npk(), 0); - let nullifier = Nullifier::for_account_initialization(&account_id); - assert!(state.private_state.1.contains(&nullifier)); - } - - #[test] - fn private_account_claimed_then_used_without_init_flag_should_fail() { - let mut state = V03State::new().with_test_programs(); - - // Set up keys for the private account - let private_keys = test_private_account_keys_1(); - - // Step 1: Create a new private account with authorization - let authorized_account = - AccountWithMetadata::new(Account::default(), true, (&private_keys.npk(), 0)); - - let claimer_program = crate::test_methods::claimer(); - - // Set up parameters for claiming the new account - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0); - - let instruction = (); - - // Step 2: Execute claimer program to claim the account with authentication - let (output, proof) = execute_and_prove( - vec![authorized_account.clone()], - Program::serialize_instruction(instruction).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedInit { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &private_keys.npk(), - &private_keys.vpk(), - ), - ssk: shared_secret, - nsk: private_keys.nsk, - identifier: 0, - }], - &claimer_program.into(), - ) - .unwrap(); - - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - let tx = PrivacyPreservingTransaction::new(message, witness_set); - - // Claim should succeed - assert!( - state - .transition_from_privacy_preserving_transaction(&tx, 1, 0) - .is_ok() - ); - - // Verify the account is now initialized (nullifier exists) - let account_id = AccountId::for_regular_private_account(&private_keys.npk(), 0); - let nullifier = Nullifier::for_account_initialization(&account_id); - assert!(state.private_state.1.contains(&nullifier)); - - // Prepare new state of account - let account_metadata = { - let mut acc = authorized_account; - acc.account.program_owner = crate::test_methods::claimer().id(); - acc - }; - - let noop_program = crate::test_methods::noop(); - let shared_secret2 = - SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0).0; - - // Step 3: Try to execute noop program with authentication but without initialization - let res = execute_and_prove( - vec![account_metadata], - Program::serialize_instruction(()).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedInit { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &private_keys.npk(), - &private_keys.vpk(), - ), - ssk: shared_secret2, - nsk: private_keys.nsk, - identifier: 0, - }], - &noop_program.into(), - ); - - assert!(matches!(res, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn public_changer_claimer_no_data_change_no_claim_succeeds() { - let initial_data = []; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let account_id = AccountId::new([1; 32]); - let program_id = crate::test_methods::changer_claimer().id(); - // Don't change data (None) and don't claim (false) - let instruction: (Option>, bool) = (None, false); - - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - // Should succeed - no changes made, no claim needed - assert!(result.is_ok()); - // Account should remain default/unclaimed - assert_eq!(state.get_account_by_id(account_id), Account::default()); - } - - #[test] - fn public_changer_claimer_data_change_no_claim_fails() { - let initial_data = []; - let mut state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let account_id = AccountId::new([1; 32]); - let program_id = crate::test_methods::changer_claimer().id(); - // Change data but don't claim (false) - should fail - let new_data = vec![1, 2, 3, 4, 5]; - let instruction: (Option>, bool) = (Some(new_data), false); - - let message = - public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - - // Should fail - cannot modify data without claiming the account - assert!(matches!( - result, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim { - account_id: err_account_id - } - )) if err_account_id == account_id - )); - } - - #[test] - fn private_changer_claimer_no_data_change_no_claim_succeeds() { - let program = crate::test_methods::changer_claimer(); - let sender_keys = test_private_account_keys_1(); - let private_account = - AccountWithMetadata::new(Account::default(), true, (&sender_keys.npk(), 0)); - // Don't change data (None) and don't claim (false) - let instruction: (Option>, bool) = (None, false); - - let result = execute_and_prove( - vec![private_account], - Program::serialize_instruction(instruction).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }], - &program.into(), - ); - - // Should succeed - no changes made, no claim needed - assert!(result.is_ok()); - } - - #[test] - fn private_changer_claimer_data_change_no_claim_fails() { - let program = crate::test_methods::changer_claimer(); - let sender_keys = test_private_account_keys_1(); - let private_account = - AccountWithMetadata::new(Account::default(), true, (&sender_keys.npk(), 0)); - // Change data but don't claim (false) - should fail - let new_data = vec![1, 2, 3, 4, 5]; - let instruction: (Option>, bool) = (Some(new_data), false); - - let result = execute_and_prove( - vec![private_account], - Program::serialize_instruction(instruction).unwrap(), - vec![InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.npk(), - &sender_keys.vpk(), - ), - ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) - .0, - nsk: sender_keys.nsk, - membership_proof: (0, vec![]), - identifier: 0, - }], - &program.into(), - ); - - // Should fail - cannot modify data without claiming the account - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test] - fn malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() { - // Arrange - let malicious_program = crate::test_methods::malicious_authorization_changer(); - let simple_transfers = crate::test_methods::simple_balance_transfer(); - let sender_keys = test_public_account_keys_1(); - let recipient_keys = test_private_account_keys_1(); - - let sender_account = AccountWithMetadata::new( - Account { - program_owner: simple_transfers.id(), - balance: 100, - ..Default::default() - }, - false, - sender_keys.account_id(), - ); - let recipient_account = - AccountWithMetadata::new(Account::default(), true, (&recipient_keys.npk(), 0)); - - let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); - let recipient_commitment = - Commitment::new(&recipient_account_id, &recipient_account.account); - let recipient_init_nullifier = Nullifier::for_account_initialization(&recipient_account_id); - let state = V03State::new() - .with_public_accounts(public_state_from_balances(&[( - sender_account.account_id, - sender_account.account.balance, - )])) - .with_private_accounts([(recipient_commitment.clone(), recipient_init_nullifier)]) - .with_test_programs(); - - let balance_to_transfer = 10_u128; - let instruction = (balance_to_transfer, simple_transfers.id()); - - let recipient = - SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0; - - let mut dependencies = HashMap::new(); - dependencies.insert(simple_transfers.id(), simple_transfers); - let program_with_deps = ProgramWithDependencies::new(malicious_program, dependencies); - - // Act - execute the malicious program - this should fail during proving - let result = execute_and_prove( - vec![sender_account, recipient_account], - Program::serialize_instruction(instruction).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(Vec::new()), - view_tag: EncryptedAccountData::compute_view_tag( - &recipient_keys.npk(), - &recipient_keys.vpk(), - ), - ssk: recipient, - nsk: recipient_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&recipient_commitment) - .expect("recipient's commitment must be in state"), - identifier: 0, - }, - ], - &program_with_deps, - ); - - // Assert - should fail because the malicious program tries to manipulate is_authorized - assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); - } - - #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] - #[test_case::test_case((Some(1), Some(3)), 2; "inside range")] - #[test_case::test_case((Some(1), Some(3)), 0; "below range")] - #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] - #[test_case::test_case((Some(1), Some(3)), 4; "above range")] - #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] - #[test_case::test_case((Some(1), None), 10; "lower bound only - above")] - #[test_case::test_case((Some(1), None), 0; "lower bound only - below")] - #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] - #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] - #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] - #[test_case::test_case((None, None), 0; "no bounds - always valid")] - #[test_case::test_case((None, None), 100; "no bounds - always valid 2")] - fn validity_window_works_in_public_transactions( - validity_window: (Option, Option), - block_id: BlockId, - ) { - let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap(); - let validity_window_program = crate::test_methods::validity_window(); - let account_keys = test_public_account_keys_1(); - let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id()); - let mut state = V03State::new().with_test_programs(); - let tx = { - let account_ids = vec![pre.account_id]; - let nonces = vec![]; - let program_id = validity_window_program.id(); - let instruction = ( - block_validity_window, - TimestampValidityWindow::new_unbounded(), - ); - let message = - public_transaction::Message::try_new(program_id, account_ids, nonces, instruction) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - PublicTransaction::new(message, witness_set) - }; - let result = state.transition_from_public_transaction(&tx, block_id, 0); - let is_inside_validity_window = - match (block_validity_window.start(), block_validity_window.end()) { - (Some(s), Some(e)) => s <= block_id && block_id < e, - (Some(s), None) => s <= block_id, - (None, Some(e)) => block_id < e, - (None, None) => true, - }; - if is_inside_validity_window { - assert!(result.is_ok()); - } else { - assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); - } - } - - #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] - #[test_case::test_case((Some(1), Some(3)), 2; "inside range")] - #[test_case::test_case((Some(1), Some(3)), 0; "below range")] - #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] - #[test_case::test_case((Some(1), Some(3)), 4; "above range")] - #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] - #[test_case::test_case((Some(1), None), 10; "lower bound only - above")] - #[test_case::test_case((Some(1), None), 0; "lower bound only - below")] - #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] - #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] - #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] - #[test_case::test_case((None, None), 0; "no bounds - always valid")] - #[test_case::test_case((None, None), 100; "no bounds - always valid 2")] - fn timestamp_validity_window_works_in_public_transactions( - validity_window: (Option, Option), - timestamp: Timestamp, - ) { - let timestamp_validity_window: TimestampValidityWindow = - validity_window.try_into().unwrap(); - let validity_window_program = crate::test_methods::validity_window(); - let account_keys = test_public_account_keys_1(); - let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id()); - let mut state = V03State::new().with_test_programs(); - let tx = { - let account_ids = vec![pre.account_id]; - let nonces = vec![]; - let program_id = validity_window_program.id(); - let instruction = ( - BlockValidityWindow::new_unbounded(), - timestamp_validity_window, - ); - let message = - public_transaction::Message::try_new(program_id, account_ids, nonces, instruction) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - PublicTransaction::new(message, witness_set) - }; - let result = state.transition_from_public_transaction(&tx, 1, timestamp); - let is_inside_validity_window = match ( - timestamp_validity_window.start(), - timestamp_validity_window.end(), - ) { - (Some(s), Some(e)) => s <= timestamp && timestamp < e, - (Some(s), None) => s <= timestamp, - (None, Some(e)) => timestamp < e, - (None, None) => true, - }; - if is_inside_validity_window { - assert!(result.is_ok()); - } else { - assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); - } - } - - #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] - #[test_case::test_case((Some(1), Some(3)), 2; "inside range")] - #[test_case::test_case((Some(1), Some(3)), 0; "below range")] - #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] - #[test_case::test_case((Some(1), Some(3)), 4; "above range")] - #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] - #[test_case::test_case((Some(1), None), 10; "lower bound only - above")] - #[test_case::test_case((Some(1), None), 0; "lower bound only - below")] - #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] - #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] - #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] - #[test_case::test_case((None, None), 0; "no bounds - always valid")] - #[test_case::test_case((None, None), 100; "no bounds - always valid 2")] - fn validity_window_works_in_privacy_preserving_transactions( - validity_window: (Option, Option), - block_id: BlockId, - ) { - let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap(); - let validity_window_program = crate::test_methods::validity_window(); - let account_keys = test_private_account_keys_1(); - let pre = AccountWithMetadata::new(Account::default(), false, (&account_keys.npk(), 0)); - let mut state = V03State::new().with_test_programs(); - let tx = { - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0); - - let instruction = ( - block_validity_window, - TimestampValidityWindow::new_unbounded(), - ); - let (output, proof) = circuit::execute_and_prove( - vec![pre], - Program::serialize_instruction(instruction).unwrap(), - vec![InputAccountIdentity::PrivateUnauthorized { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &account_keys.npk(), - &account_keys.vpk(), - ), - npk: account_keys.npk(), - ssk: shared_secret, - identifier: 0, - }], - &validity_window_program.into(), - ) - .unwrap(); - - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - PrivacyPreservingTransaction::new(message, witness_set) - }; - let result = state.transition_from_privacy_preserving_transaction(&tx, block_id, 0); - let is_inside_validity_window = - match (block_validity_window.start(), block_validity_window.end()) { - (Some(s), Some(e)) => s <= block_id && block_id < e, - (Some(s), None) => s <= block_id, - (None, Some(e)) => block_id < e, - (None, None) => true, - }; - if is_inside_validity_window { - assert!(result.is_ok()); - } else { - assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); - } - } - - #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] - #[test_case::test_case((Some(1), Some(3)), 2; "inside range")] - #[test_case::test_case((Some(1), Some(3)), 0; "below range")] - #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] - #[test_case::test_case((Some(1), Some(3)), 4; "above range")] - #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] - #[test_case::test_case((Some(1), None), 10; "lower bound only - above")] - #[test_case::test_case((Some(1), None), 0; "lower bound only - below")] - #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] - #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] - #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] - #[test_case::test_case((None, None), 0; "no bounds - always valid")] - #[test_case::test_case((None, None), 100; "no bounds - always valid 2")] - fn timestamp_validity_window_works_in_privacy_preserving_transactions( - validity_window: (Option, Option), - timestamp: Timestamp, - ) { - let timestamp_validity_window: TimestampValidityWindow = - validity_window.try_into().unwrap(); - let validity_window_program = crate::test_methods::validity_window(); - let account_keys = test_private_account_keys_1(); - let pre = AccountWithMetadata::new(Account::default(), false, (&account_keys.npk(), 0)); - let mut state = V03State::new().with_test_programs(); - let tx = { - let (shared_secret, epk) = - SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0); - - let instruction = ( - BlockValidityWindow::new_unbounded(), - timestamp_validity_window, - ); - let (output, proof) = circuit::execute_and_prove( - vec![pre], - Program::serialize_instruction(instruction).unwrap(), - vec![InputAccountIdentity::PrivateUnauthorized { - epk, - view_tag: EncryptedAccountData::compute_view_tag( - &account_keys.npk(), - &account_keys.vpk(), - ), - npk: account_keys.npk(), - ssk: shared_secret, - identifier: 0, - }], - &validity_window_program.into(), - ) - .unwrap(); - - let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); - PrivacyPreservingTransaction::new(message, witness_set) - }; - let result = state.transition_from_privacy_preserving_transaction(&tx, 1, timestamp); - let is_inside_validity_window = match ( - timestamp_validity_window.start(), - timestamp_validity_window.end(), - ) { - (Some(s), Some(e)) => s <= timestamp && timestamp < e, - (Some(s), None) => s <= timestamp, - (None, Some(e)) => timestamp < e, - (None, None) => true, - }; - if is_inside_validity_window { - assert!(result.is_ok()); - } else { - assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); - } - } - - #[test] - fn state_serialization_roundtrip() { - let account_id_1 = AccountId::new([1; 32]); - let account_id_2 = AccountId::new([2; 32]); - let initial_data = [(account_id_1, 100_u128), (account_id_2, 151_u128)]; - let state = V03State::new() - .with_public_accounts(public_state_from_balances(&initial_data)) - .with_test_programs(); - let bytes = borsh::to_vec(&state).unwrap(); - let state_from_bytes: V03State = borsh::from_slice(&bytes).unwrap(); - assert_eq!(state, state_from_bytes); - } - - #[test] - fn flash_swap_successful() { - let initiator = crate::test_methods::flash_swap_initiator(); - let callback = crate::test_methods::flash_swap_callback(); - let token = crate::test_methods::simple_balance_transfer(); - - let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); - let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32])); - - let initial_balance: u128 = 1000; - let amount_out: u128 = 100; - - let vault_account = Account { - program_owner: token.id(), - balance: initial_balance, - ..Account::default() - }; - let receiver_account = Account { - program_owner: token.id(), - balance: 0, - ..Account::default() - }; - - let mut state = V03State::new().with_test_programs(); - state.force_insert_account(vault_id, vault_account); - state.force_insert_account(receiver_id, receiver_account); - - // Callback instruction: return funds - let cb_instruction = CallbackInstruction { - return_funds: true, - token_program_id: token.id(), - amount: amount_out, - }; - let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); - - let instruction = FlashSwapInstruction::Initiate { - token_program_id: token.id(), - callback_program_id: callback.id(), - amount_out, - callback_instruction_data: cb_data, - }; - - let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); - let result = state.transition_from_public_transaction(&tx, 1, 0); - assert!(result.is_ok(), "flash swap should succeed: {result:?}"); - - // Vault balance restored, receiver back to 0 - assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance); - assert_eq!(state.get_account_by_id(receiver_id).balance, 0); - } - - #[test] - fn flash_swap_callback_keeps_funds_rollback() { - let initiator = crate::test_methods::flash_swap_initiator(); - let callback = crate::test_methods::flash_swap_callback(); - let token = crate::test_methods::simple_balance_transfer(); - - let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); - let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32])); - - let initial_balance: u128 = 1000; - let amount_out: u128 = 100; - - let vault_account = Account { - program_owner: token.id(), - balance: initial_balance, - ..Account::default() - }; - let receiver_account = Account { - program_owner: token.id(), - balance: 0, - ..Account::default() - }; - - let mut state = V03State::new().with_test_programs(); - state.force_insert_account(vault_id, vault_account); - state.force_insert_account(receiver_id, receiver_account); - - // Callback instruction: do NOT return funds - let cb_instruction = CallbackInstruction { - return_funds: false, - token_program_id: token.id(), - amount: amount_out, - }; - let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); - - let instruction = FlashSwapInstruction::Initiate { - token_program_id: token.id(), - callback_program_id: callback.id(), - amount_out, - callback_instruction_data: cb_data, - }; - - let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); - let result = state.transition_from_public_transaction(&tx, 1, 0); - - // Invariant check fails → entire tx rolls back - assert!( - result.is_err(), - "flash swap should fail when callback keeps funds" - ); - - // State unchanged (rollback) - assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance); - assert_eq!(state.get_account_by_id(receiver_id).balance, 0); - } - - #[test] - fn flash_swap_self_call_targets_correct_program() { - // Zero-amount flash swap: the invariant self-call still runs and succeeds - // because vault balance doesn't decrease. - let initiator = crate::test_methods::flash_swap_initiator(); - let callback = crate::test_methods::flash_swap_callback(); - let token = crate::test_methods::simple_balance_transfer(); - - let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); - let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32])); - - let initial_balance: u128 = 1000; - - let vault_account = Account { - program_owner: token.id(), - balance: initial_balance, - ..Account::default() - }; - let receiver_account = Account { - program_owner: token.id(), - balance: 0, - ..Account::default() - }; - - let mut state = V03State::new().with_test_programs(); - state.force_insert_account(vault_id, vault_account); - state.force_insert_account(receiver_id, receiver_account); - - let cb_instruction = CallbackInstruction { - return_funds: true, - token_program_id: token.id(), - amount: 0, - }; - let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); - - let instruction = FlashSwapInstruction::Initiate { - token_program_id: token.id(), - callback_program_id: callback.id(), - amount_out: 0, - callback_instruction_data: cb_data, - }; - - let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); - let result = state.transition_from_public_transaction(&tx, 1, 0); - assert!( - result.is_ok(), - "zero-amount flash swap should succeed: {result:?}" - ); - } - - #[test] - fn flash_swap_standalone_invariant_check_rejected() { - // Calling InvariantCheck directly (not as a chained self-call) should fail - // because caller_program_id will be None. - let initiator = crate::test_methods::flash_swap_initiator(); - let token = crate::test_methods::simple_balance_transfer(); - - let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); - - let vault_account = Account { - program_owner: token.id(), - balance: 1000, - ..Account::default() - }; - - let mut state = V03State::new().with_test_programs(); - state.force_insert_account(vault_id, vault_account); - - let instruction = FlashSwapInstruction::InvariantCheck { - min_vault_balance: 1000, - }; - - let message = public_transaction::Message::try_new( - initiator.id(), - vec![vault_id], - vec![], - instruction, - ) - .unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - assert!( - result.is_err(), - "standalone InvariantCheck should be rejected (caller_program_id is None)" - ); - } - - #[test] - fn malicious_self_program_id_rejected_in_public_execution() { - let program = crate::test_methods::malicious_self_program_id(); - let acc_id = AccountId::new([99; 32]); - let account = Account::default(); - - let mut state = V03State::new().with_test_programs(); - state.force_insert_account(acc_id, account); - - let message = - public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - assert!( - result.is_err(), - "program with wrong self_program_id in output should be rejected" - ); - } - - #[test] - fn malicious_caller_program_id_rejected_in_public_execution() { - let program = crate::test_methods::malicious_caller_program_id(); - let acc_id = AccountId::new([99; 32]); - let account = Account::default(); - - let mut state = V03State::new().with_test_programs(); - state.force_insert_account(acc_id, account); - - let message = - public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap(); - let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - - let result = state.transition_from_public_transaction(&tx, 1, 0); - assert!( - result.is_err(), - "program with spoofed caller_program_id in output should be rejected" - ); - } - - #[test] - fn two_private_pda_family_members_receive_and_spend() { - let funder_keys = test_public_account_keys_1(); - let alice_keys = test_private_account_keys_1(); - let alice_npk = alice_keys.npk(); - - let proxy = crate::test_methods::pda_spend_proxy(); - let simple_transfer = crate::test_methods::simple_balance_transfer(); - let proxy_id = proxy.id(); - let simple_transfer_id = simple_transfer.id(); - let seed = PdaSeed::new([42; 32]); - let amount: u128 = 100; - - let spend_with_deps = ProgramWithDependencies::new( - proxy, - [(simple_transfer_id, simple_transfer.clone())].into(), - ); - - let funder_id = funder_keys.account_id(); - let alice_pda_0_id = AccountId::for_private_pda(&proxy_id, &seed, &alice_npk, 0); - let alice_pda_1_id = AccountId::for_private_pda(&proxy_id, &seed, &alice_npk, 1); - let recipient_id = test_public_account_keys_2().account_id(); - let recipient_signing_key = test_public_account_keys_2().signing_key; - - let mut state = - V03State::new().with_public_accounts(public_state_from_balances(&[(funder_id, 500)])); - - let alice_pda_0_account = Account { - program_owner: simple_transfer_id, - balance: amount, - nonce: Nonce::private_account_nonce_init(&alice_pda_0_id), - ..Account::default() - }; - let alice_pda_1_account = Account { - program_owner: simple_transfer_id, - balance: amount, - nonce: Nonce::private_account_nonce_init(&alice_pda_1_id), - ..Account::default() - }; - - let (alice_shared_0, alice_epk_0) = - SharedSecretKey::encapsulate_deterministic(&alice_keys.vpk(), &[0_u8; 32], 0); - let (alice_shared_1, alice_epk_1) = - SharedSecretKey::encapsulate_deterministic(&alice_keys.vpk(), &[0_u8; 32], 1); - - // Fund alice_pda_0 via authenticated_transfer directly. - { - let funder_account = state.get_account_by_id(funder_id); - let funder_nonce = funder_account.nonce; - let (output, proof) = execute_and_prove( - vec![ - AccountWithMetadata::new(funder_account, true, funder_id), - AccountWithMetadata::new(Account::default(), false, alice_pda_0_id), - ], - Program::serialize_instruction(amount).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivatePdaInit { - epk: alice_epk_0.clone(), - view_tag: EncryptedAccountData::compute_view_tag( - &alice_npk, - &alice_keys.vpk(), - ), - npk: alice_npk, - ssk: alice_shared_0, - identifier: 0, - seed: Some((seed, proxy_id)), - }, - ], - &simple_transfer.clone().into(), - ) - .unwrap(); - let message = - Message::try_from_circuit_output(vec![funder_id], vec![funder_nonce], output) - .unwrap(); - let witness_set = WitnessSet::for_message(&message, proof, &[&funder_keys.signing_key]); - state - .transition_from_privacy_preserving_transaction( - &PrivacyPreservingTransaction::new(message, witness_set), - 1, - 0, - ) - .unwrap(); - } - - // Fund alice_pda_1 the same way with identifier 1. - { - let funder_account = state.get_account_by_id(funder_id); - let funder_nonce = funder_account.nonce; - let (output, proof) = execute_and_prove( - vec![ - AccountWithMetadata::new(funder_account, true, funder_id), - AccountWithMetadata::new(Account::default(), false, alice_pda_1_id), - ], - Program::serialize_instruction(amount).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivatePdaInit { - epk: alice_epk_1.clone(), - view_tag: EncryptedAccountData::compute_view_tag( - &alice_npk, - &alice_keys.vpk(), - ), - npk: alice_npk, - ssk: alice_shared_1, - identifier: 1, - seed: Some((seed, proxy_id)), - }, - ], - &simple_transfer.into(), - ) - .unwrap(); - let message = - Message::try_from_circuit_output(vec![funder_id], vec![funder_nonce], output) - .unwrap(); - let witness_set = WitnessSet::for_message(&message, proof, &[&funder_keys.signing_key]); - state - .transition_from_privacy_preserving_transaction( - &PrivacyPreservingTransaction::new(message, witness_set), - 2, - 0, - ) - .unwrap(); - } - - let commitment_pda_0 = Commitment::new(&alice_pda_0_id, &alice_pda_0_account); - let commitment_pda_1 = Commitment::new(&alice_pda_1_id, &alice_pda_1_account); - - assert!(state.get_proof_for_commitment(&commitment_pda_0).is_some()); - assert!(state.get_proof_for_commitment(&commitment_pda_1).is_some()); - - // Alice spends alice_pda_0 into the public recipient. - { - let recipient_account = state.get_account_by_id(recipient_id); - let (output, proof) = execute_and_prove( - vec![ - AccountWithMetadata::new(alice_pda_0_account, true, alice_pda_0_id), - AccountWithMetadata::new(recipient_account, true, recipient_id), - ], - Program::serialize_instruction((seed, amount, simple_transfer_id)).unwrap(), - vec![ - InputAccountIdentity::PrivatePdaUpdate { - epk: alice_epk_0, - view_tag: EncryptedAccountData::compute_view_tag( - &alice_npk, - &alice_keys.vpk(), - ), - ssk: alice_shared_0, - nsk: alice_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&commitment_pda_0) - .expect("pda_0 must be in state"), - identifier: 0, - seed: None, - }, - InputAccountIdentity::Public, - ], - &spend_with_deps, - ) - .unwrap(); - let message = - Message::try_from_circuit_output(vec![recipient_id], vec![Nonce(0)], output) - .unwrap(); - let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_signing_key]); - state - .transition_from_privacy_preserving_transaction( - &PrivacyPreservingTransaction::new(message, witness_set), - 3, - 0, - ) - .unwrap(); - } - - // Alice spends alice_pda_1 into the same public recipient. - { - let recipient_account = state.get_account_by_id(recipient_id); - let (output, proof) = execute_and_prove( - vec![ - AccountWithMetadata::new(alice_pda_1_account.clone(), true, alice_pda_1_id), - AccountWithMetadata::new(recipient_account, false, recipient_id), - ], - Program::serialize_instruction((seed, amount, simple_transfer_id)).unwrap(), - vec![ - InputAccountIdentity::PrivatePdaUpdate { - epk: alice_epk_1, - view_tag: EncryptedAccountData::compute_view_tag( - &alice_npk, - &alice_keys.vpk(), - ), - ssk: alice_shared_1, - nsk: alice_keys.nsk, - membership_proof: state - .get_proof_for_commitment(&commitment_pda_1) - .expect("pda_1 must be in state"), - identifier: 1, - seed: None, - }, - InputAccountIdentity::Public, - ], - &spend_with_deps, - ) - .unwrap(); - let message = - Message::try_from_circuit_output(vec![recipient_id], vec![], output).unwrap(); - let witness_set = WitnessSet::for_message(&message, proof, &[]); - state - .transition_from_privacy_preserving_transaction( - &PrivacyPreservingTransaction::new(message, witness_set), - 4, - 0, - ) - .unwrap(); - } - - assert_eq!(state.get_account_by_id(recipient_id).balance, 2 * amount); - - // Re-fund alice_pda_1 top-level via simple_transfer using PrivatePdaUpdate with an - // external seed. - let alice_pda_1_account_after_spend = Account { - program_owner: simple_transfer_id, - balance: 0, - nonce: alice_pda_1_account - .nonce - .private_account_nonce_increment(&alice_keys.nsk), - ..Account::default() - }; - let commitment_pda_1_after_spend = - Commitment::new(&alice_pda_1_id, &alice_pda_1_account_after_spend); - let alice_shared_1_refund = SharedSecretKey([12; 32]); - { - let recipient_account = state.get_account_by_id(recipient_id); - let recipient_nonce = recipient_account.nonce; - let (output, proof) = execute_and_prove( - vec![ - AccountWithMetadata::new(recipient_account, true, recipient_id), - AccountWithMetadata::new( - alice_pda_1_account_after_spend, - false, - alice_pda_1_id, - ), - ], - Program::serialize_instruction(amount).unwrap(), - vec![ - InputAccountIdentity::Public, - InputAccountIdentity::PrivatePdaUpdate { - epk: EphemeralPublicKey(vec![12_u8; 1088]), - view_tag: EncryptedAccountData::compute_view_tag( - &alice_npk, - &alice_keys.vpk(), - ), - nsk: alice_keys.nsk, - ssk: alice_shared_1_refund, - membership_proof: state - .get_proof_for_commitment(&commitment_pda_1_after_spend) - .expect("pda_1 after spend must be in state"), - identifier: 1, - seed: Some((seed, proxy_id)), - }, - ], - &crate::test_methods::simple_balance_transfer().into(), - ) - .unwrap(); - let message = - Message::try_from_circuit_output(vec![recipient_id], vec![recipient_nonce], output) - .unwrap(); - let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_signing_key]); - state - .transition_from_privacy_preserving_transaction( - &PrivacyPreservingTransaction::new(message, witness_set), - 5, - 0, - ) - .unwrap(); - } - - assert_eq!(state.get_account_by_id(recipient_id).balance, amount); - } -} diff --git a/lee/state_machine/src/state/mod.rs b/lee/state_machine/src/state/mod.rs new file mode 100644 index 00000000..fcf71a21 --- /dev/null +++ b/lee/state_machine/src/state/mod.rs @@ -0,0 +1,320 @@ +use std::collections::{BTreeSet, HashMap, HashSet}; + +use borsh::{BorshDeserialize, BorshSerialize}; +use lee_core::{ + BlockId, Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier, + Timestamp, + account::{Account, AccountId}, + program::ProgramId, +}; + +use crate::{ + error::LeeError, + merkle_tree::MerkleTree, + privacy_preserving_transaction::PrivacyPreservingTransaction, + program::Program, + program_deployment_transaction::ProgramDeploymentTransaction, + public_transaction::PublicTransaction, + validated_state_diff::{StateDiff, ValidatedStateDiff}, +}; + +pub const MAX_NUMBER_CHAINED_CALLS: usize = 10; + +#[derive(Clone, BorshSerialize, BorshDeserialize)] +#[cfg_attr(test, derive(Debug, PartialEq, Eq))] +pub struct CommitmentSet { + merkle_tree: MerkleTree, + commitments: HashMap, + root_history: HashSet, +} + +impl CommitmentSet { + pub(crate) fn digest(&self) -> CommitmentSetDigest { + self.merkle_tree.root() + } + + /// Queries the `CommitmentSet` for a membership proof of commitment. + pub fn get_proof_for(&self, commitment: &Commitment) -> Option { + let index = *self.commitments.get(commitment)?; + + self.merkle_tree + .get_authentication_path_for(index) + .map(|path| (index, path)) + } + + /// Inserts a list of commitments to the `CommitmentSet`. + pub(crate) fn extend(&mut self, commitments: &[Commitment]) { + for commitment in commitments.iter().cloned() { + let index = self.merkle_tree.insert(commitment.to_byte_array()); + self.commitments.insert(commitment, index); + } + self.root_history.insert(self.digest()); + } + + fn contains(&self, commitment: &Commitment) -> bool { + self.commitments.contains_key(commitment) + } + + /// Initializes an empty `CommitmentSet` with a given capacity. + /// If the capacity is not a `power_of_two`, then capacity is taken + /// to be the next `power_of_two`. + pub(crate) fn with_capacity(capacity: usize) -> Self { + Self { + merkle_tree: MerkleTree::with_capacity(capacity), + commitments: HashMap::new(), + root_history: HashSet::new(), + } + } +} + +#[cfg_attr(test, derive(Debug, PartialEq, Eq))] +#[derive(Clone)] +struct NullifierSet(BTreeSet); + +impl NullifierSet { + const fn new() -> Self { + Self(BTreeSet::new()) + } + + fn extend(&mut self, new_nullifiers: &[Nullifier]) { + self.0.extend(new_nullifiers); + } + + fn contains(&self, nullifier: &Nullifier) -> bool { + self.0.contains(nullifier) + } +} + +impl BorshSerialize for NullifierSet { + fn serialize(&self, writer: &mut W) -> std::io::Result<()> { + self.0.iter().collect::>().serialize(writer) + } +} + +impl BorshDeserialize for NullifierSet { + fn deserialize_reader(reader: &mut R) -> std::io::Result { + let vec = Vec::::deserialize_reader(reader)?; + + let mut set = BTreeSet::new(); + for n in vec { + if !set.insert(n) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "duplicate nullifier in NullifierSet", + )); + } + } + + Ok(Self(set)) + } +} + +#[derive(Clone, BorshSerialize, BorshDeserialize)] +#[cfg_attr(test, derive(Debug, PartialEq, Eq))] +pub struct V03State { + public_state: HashMap, + private_state: (CommitmentSet, NullifierSet), + programs: HashMap, +} + +impl Default for V03State { + fn default() -> Self { + let mut commitment_set = CommitmentSet::with_capacity(32); + commitment_set.extend(&[DUMMY_COMMITMENT]); + let nullifier_set = NullifierSet::new(); + let private_state = (commitment_set, nullifier_set); + + Self { + public_state: HashMap::default(), + private_state, + programs: HashMap::default(), + } + } +} + +impl V03State { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Initializes state with given public account balances leaving other account fields at their + /// default values. + #[must_use] + pub fn with_public_account_balances( + mut self, + balances: impl IntoIterator, + ) -> Self { + let public_accounts = balances.into_iter().map(|(account_id, balance)| { + ( + account_id, + Account { + balance, + ..Account::default() + }, + ) + }); + self.public_state.extend(public_accounts); + self + } + + /// Initializes state with given public accounts. + #[must_use] + pub fn with_public_accounts( + mut self, + public_accounts: impl IntoIterator, + ) -> Self { + self.public_state.extend(public_accounts); + self + } + + /// Initializes state with given private accounts. + #[must_use] + pub fn with_private_accounts( + mut self, + private_accounts: impl IntoIterator, + ) -> Self { + let (commitments, nullifiers): (Vec, Vec) = + private_accounts.into_iter().unzip(); + self.private_state.0.extend(&commitments); + self.private_state.1.extend(&nullifiers); + self + } + + /// Initializes state with given builtin programs. + #[must_use] + pub fn with_programs(mut self, programs: impl IntoIterator) -> Self { + for program in programs { + self.insert_program(program); + } + self + } + + pub(crate) fn insert_program(&mut self, program: Program) { + self.programs.insert(program.id(), program); + } + + pub fn apply_state_diff(&mut self, diff: ValidatedStateDiff) { + let StateDiff { + signer_account_ids, + public_diff, + new_commitments, + new_nullifiers, + program, + } = diff.into_state_diff(); + #[expect( + clippy::iter_over_hash_type, + reason = "Iteration order doesn't matter here" + )] + for (account_id, account) in public_diff { + *self.get_account_by_id_mut(account_id) = account; + } + for account_id in signer_account_ids { + self.get_account_by_id_mut(account_id) + .nonce + .public_account_nonce_increment(); + } + self.private_state.0.extend(&new_commitments); + self.private_state.1.extend(&new_nullifiers); + if let Some(program) = program { + self.insert_program(program); + } + } + + pub fn transition_from_public_transaction( + &mut self, + tx: &PublicTransaction, + block_id: BlockId, + timestamp: Timestamp, + ) -> Result<(), LeeError> { + let diff = ValidatedStateDiff::from_public_transaction(tx, self, block_id, timestamp)?; + self.apply_state_diff(diff); + Ok(()) + } + + pub fn transition_from_privacy_preserving_transaction( + &mut self, + tx: &PrivacyPreservingTransaction, + block_id: BlockId, + timestamp: Timestamp, + ) -> Result<(), LeeError> { + let diff = + ValidatedStateDiff::from_privacy_preserving_transaction(tx, self, block_id, timestamp)?; + self.apply_state_diff(diff); + Ok(()) + } + + pub fn transition_from_program_deployment_transaction( + &mut self, + tx: &ProgramDeploymentTransaction, + ) -> Result<(), LeeError> { + let diff = ValidatedStateDiff::from_program_deployment_transaction(tx, self)?; + self.apply_state_diff(diff); + Ok(()) + } + + fn get_account_by_id_mut(&mut self, account_id: AccountId) -> &mut Account { + self.public_state.entry(account_id).or_default() + } + + #[must_use] + pub fn get_account_by_id(&self, account_id: AccountId) -> Account { + self.public_state + .get(&account_id) + .cloned() + .unwrap_or_else(Account::default) + } + + #[must_use] + pub fn get_proof_for_commitment(&self, commitment: &Commitment) -> Option { + self.private_state.0.get_proof_for(commitment) + } + + pub(crate) const fn programs(&self) -> &HashMap { + &self.programs + } + + #[must_use] + pub fn commitment_set_digest(&self) -> CommitmentSetDigest { + self.private_state.0.digest() + } + + pub(crate) fn check_commitments_are_new( + &self, + new_commitments: &[Commitment], + ) -> Result<(), LeeError> { + for commitment in new_commitments { + if self.private_state.0.contains(commitment) { + return Err(LeeError::InvalidInput("Commitment already seen".to_owned())); + } + } + Ok(()) + } + + pub(crate) fn check_nullifiers_are_valid( + &self, + new_nullifiers: &[(Nullifier, CommitmentSetDigest)], + ) -> Result<(), LeeError> { + for (nullifier, digest) in new_nullifiers { + if self.private_state.1.contains(nullifier) { + return Err(LeeError::InvalidInput("Nullifier already seen".to_owned())); + } + if !self.private_state.0.root_history.contains(digest) { + return Err(LeeError::InvalidInput( + "Unrecognized commitment set digest".to_owned(), + )); + } + } + Ok(()) + } +} + +#[cfg(any(test, feature = "test-utils"))] +impl V03State { + pub fn force_insert_account(&mut self, account_id: AccountId, account: Account) { + self.public_state.insert(account_id, account); + } +} + +#[cfg(test)] +pub mod tests; diff --git a/lee/state_machine/src/state/tests/authenticated_transfer.rs b/lee/state_machine/src/state/tests/authenticated_transfer.rs new file mode 100644 index 00000000..8d227fc3 --- /dev/null +++ b/lee/state_machine/src/state/tests/authenticated_transfer.rs @@ -0,0 +1,149 @@ +use super::*; + +#[test] +fn transition_from_authenticated_transfer_program_invocation_default_account_destination() { + let key = PrivateKey::try_new([1; 32]).unwrap(); + let account_id = AccountId::from(&PublicKey::new_from_private_key(&key)); + let initial_data = [( + account_id, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + ..Account::default() + }, + )]; + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs(); + let from = account_id; + let to_key = PrivateKey::try_new([2; 32]).unwrap(); + let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); + assert_eq!(state.get_account_by_id(to), Account::default()); + let balance_to_move = 5; + + let tx = transfer_transaction(from, &key, 0, to, &to_key, 0, balance_to_move); + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + assert_eq!(state.get_account_by_id(from).balance, 95); + assert_eq!(state.get_account_by_id(to).balance, 5); + assert_eq!(state.get_account_by_id(from).nonce, Nonce(1)); + assert_eq!(state.get_account_by_id(to).nonce, Nonce(1)); +} + +#[test] +fn transition_from_authenticated_transfer_program_invocation_insuficient_balance() { + let key = PrivateKey::try_new([1; 32]).unwrap(); + let account_id = AccountId::from(&PublicKey::new_from_private_key(&key)); + let mut state = V03State::new() + .with_public_account_balances([(account_id, 100)]) + .with_test_programs(); + let from = account_id; + let from_key = key; + let to_key = PrivateKey::try_new([2; 32]).unwrap(); + let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); + let balance_to_move = 101; + assert!(state.get_account_by_id(from).balance < balance_to_move); + + let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move); + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!(result, Err(LeeError::ProgramExecutionFailed(_)))); + assert_eq!(state.get_account_by_id(from).balance, 100); + assert_eq!(state.get_account_by_id(to).balance, 0); + assert_eq!(state.get_account_by_id(from).nonce, Nonce(0)); + assert_eq!(state.get_account_by_id(to).nonce, Nonce(0)); +} + +#[test] +fn transition_from_authenticated_transfer_program_invocation_non_default_account_destination() { + let key1 = PrivateKey::try_new([1; 32]).unwrap(); + let key2 = PrivateKey::try_new([2; 32]).unwrap(); + let account_id1 = AccountId::from(&PublicKey::new_from_private_key(&key1)); + let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2)); + let initial_data = [ + ( + account_id1, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + ..Account::default() + }, + ), + ( + account_id2, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 200, + ..Account::default() + }, + ), + ]; + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs(); + let from = account_id2; + let from_key = key2; + let to = account_id1; + let to_key = key1; + assert_ne!(state.get_account_by_id(to), Account::default()); + let balance_to_move = 8; + + let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move); + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + assert_eq!(state.get_account_by_id(from).balance, 192); + assert_eq!(state.get_account_by_id(to).balance, 108); + assert_eq!(state.get_account_by_id(from).nonce, Nonce(1)); + assert_eq!(state.get_account_by_id(to).nonce, Nonce(1)); +} + +#[test] +fn transition_from_sequence_of_authenticated_transfer_program_invocations() { + let key1 = PrivateKey::try_new([8; 32]).unwrap(); + let account_id1 = AccountId::from(&PublicKey::new_from_private_key(&key1)); + let key2 = PrivateKey::try_new([2; 32]).unwrap(); + let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2)); + let initial_data = [( + account_id1, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + ..Account::default() + }, + )]; + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs(); + let key3 = PrivateKey::try_new([3; 32]).unwrap(); + let account_id3 = AccountId::from(&PublicKey::new_from_private_key(&key3)); + let balance_to_move = 5; + + let tx = transfer_transaction( + account_id1, + &key1, + 0, + account_id2, + &key2, + 0, + balance_to_move, + ); + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + let balance_to_move = 3; + let tx = transfer_transaction( + account_id2, + &key2, + 1, + account_id3, + &key3, + 0, + balance_to_move, + ); + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + assert_eq!(state.get_account_by_id(account_id1).balance, 95); + assert_eq!(state.get_account_by_id(account_id2).balance, 2); + assert_eq!(state.get_account_by_id(account_id3).balance, 3); + assert_eq!(state.get_account_by_id(account_id1).nonce, Nonce(1)); + assert_eq!(state.get_account_by_id(account_id2).nonce, Nonce(2)); + assert_eq!(state.get_account_by_id(account_id3).nonce, Nonce(1)); +} diff --git a/lee/state_machine/src/state/tests/changer_claimer.rs b/lee/state_machine/src/state/tests/changer_claimer.rs new file mode 100644 index 00000000..f9ddce25 --- /dev/null +++ b/lee/state_machine/src/state/tests/changer_claimer.rs @@ -0,0 +1,118 @@ +use super::*; + +#[test] +fn public_changer_claimer_no_data_change_no_claim_succeeds() { + let initial_data = []; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let account_id = AccountId::new([1; 32]); + let program_id = crate::test_methods::changer_claimer().id(); + // Don't change data (None) and don't claim (false) + let instruction: (Option>, bool) = (None, false); + + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + // Should succeed - no changes made, no claim needed + assert!(result.is_ok()); + // Account should remain default/unclaimed + assert_eq!(state.get_account_by_id(account_id), Account::default()); +} + +#[test] +fn public_changer_claimer_data_change_no_claim_fails() { + let initial_data = []; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let account_id = AccountId::new([1; 32]); + let program_id = crate::test_methods::changer_claimer().id(); + // Change data but don't claim (false) - should fail + let new_data = vec![1, 2, 3, 4, 5]; + let instruction: (Option>, bool) = (Some(new_data), false); + + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + // Should fail - cannot modify data without claiming the account + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim { + account_id: err_account_id + } + )) if err_account_id == account_id + )); +} + +#[test] +fn private_changer_claimer_no_data_change_no_claim_succeeds() { + let program = crate::test_methods::changer_claimer(); + let sender_keys = test_private_account_keys_1(); + let private_account = + AccountWithMetadata::new(Account::default(), true, (&sender_keys.npk(), 0)); + // Don't change data (None) and don't claim (false) + let instruction: (Option>, bool) = (None, false); + + let result = execute_and_prove( + vec![private_account], + Program::serialize_instruction(instruction).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }], + &program.into(), + ); + + // Should succeed - no changes made, no claim needed + assert!(result.is_ok()); +} + +#[test] +fn private_changer_claimer_data_change_no_claim_fails() { + let program = crate::test_methods::changer_claimer(); + let sender_keys = test_private_account_keys_1(); + let private_account = + AccountWithMetadata::new(Account::default(), true, (&sender_keys.npk(), 0)); + // Change data but don't claim (false) - should fail + let new_data = vec![1, 2, 3, 4, 5]; + let instruction: (Option>, bool) = (Some(new_data), false); + + let result = execute_and_prove( + vec![private_account], + Program::serialize_instruction(instruction).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }], + &program.into(), + ); + + // Should fail - cannot modify data without claiming the account + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} diff --git a/lee/state_machine/src/state/tests/circuit.rs b/lee/state_machine/src/state/tests/circuit.rs new file mode 100644 index 00000000..0860e98b --- /dev/null +++ b/lee/state_machine/src/state/tests/circuit.rs @@ -0,0 +1,1249 @@ +use super::*; + +#[test] +fn circuit_fails_if_visibility_masks_have_incorrect_lenght() { + let program = crate::test_methods::simple_balance_transfer(); + let public_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + let public_account_2 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([1; 32]), + ); + + // Single account_identity entry for a circuit execution with two pre_state accounts. + let result = execute_and_prove( + vec![public_account_1, public_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn circuit_fails_if_invalid_auth_keys_are_provided() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + let private_account_2 = + AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0)); + + // Setting the recipient nsk to authorize the sender. + // This should be set to the sender private account in a normal circumstance. + // `PrivateAuthorizedUpdate` derives npk from nsk and asserts equality with + // `pre_state.account_id`, so a mismatched nsk fails that check. + let result = execute_and_prove( + vec![private_account_1, private_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) + .0, + nsk: recipient_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: SharedSecretKey::encapsulate_deterministic( + &recipient_keys.vpk(), + &[0_u8; 32], + 0, + ) + .0, + commitment_root: DUMMY_COMMITMENT_HASH, + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn circuit_should_fail_if_new_private_account_with_non_default_balance_is_provided() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + let private_account_2 = AccountWithMetadata::new( + Account { + // Non default balance + balance: 1, + ..Account::default() + }, + false, + (&recipient_keys.npk(), 0), + ); + + let result = execute_and_prove( + vec![private_account_1, private_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) + .0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: SharedSecretKey::encapsulate_deterministic( + &recipient_keys.vpk(), + &[0_u8; 32], + 0, + ) + .0, + commitment_root: DUMMY_COMMITMENT_HASH, + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn circuit_should_fail_if_new_private_account_with_non_default_program_owner_is_provided() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + let private_account_2 = AccountWithMetadata::new( + Account { + // Non default program_owner + program_owner: [0, 1, 2, 3, 4, 5, 6, 7], + ..Account::default() + }, + false, + (&recipient_keys.npk(), 0), + ); + + let result = execute_and_prove( + vec![private_account_1, private_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) + .0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: SharedSecretKey::encapsulate_deterministic( + &recipient_keys.vpk(), + &[0_u8; 32], + 0, + ) + .0, + commitment_root: DUMMY_COMMITMENT_HASH, + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn circuit_should_fail_if_new_private_account_with_non_default_data_is_provided() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + let private_account_2 = AccountWithMetadata::new( + Account { + // Non default data + data: b"hola mundo".to_vec().try_into().unwrap(), + ..Account::default() + }, + false, + (&recipient_keys.npk(), 0), + ); + + let result = execute_and_prove( + vec![private_account_1, private_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) + .0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: SharedSecretKey::encapsulate_deterministic( + &recipient_keys.vpk(), + &[0_u8; 32], + 0, + ) + .0, + commitment_root: DUMMY_COMMITMENT_HASH, + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn circuit_should_fail_if_new_private_account_with_non_default_nonce_is_provided() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + let private_account_2 = AccountWithMetadata::new( + Account { + // Non default nonce + nonce: Nonce(0xdead_beef), + ..Account::default() + }, + false, + (&recipient_keys.npk(), 0), + ); + + let result = execute_and_prove( + vec![private_account_1, private_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) + .0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: SharedSecretKey::encapsulate_deterministic( + &recipient_keys.vpk(), + &[0_u8; 32], + 0, + ) + .0, + commitment_root: DUMMY_COMMITMENT_HASH, + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn circuit_should_fail_if_new_private_account_is_provided_with_default_values_but_marked_as_authorized() + { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let recipient_keys = test_private_account_keys_2(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + let private_account_2 = AccountWithMetadata::new( + Account::default(), + // This should be set to false in normal circumstances + true, + (&recipient_keys.npk(), 0), + ); + + let result = execute_and_prove( + vec![private_account_1, private_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0) + .0, + nsk: sender_keys.nsk, + membership_proof: (0, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: SharedSecretKey::encapsulate_deterministic( + &recipient_keys.vpk(), + &[0_u8; 32], + 0, + ) + .0, + commitment_root: DUMMY_COMMITMENT_HASH, + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +/// A private PDA account that no program claims via `Claim::Pda` and no caller authorizes via +/// `ChainedCall.pda_seeds` has no binding between its supplied npk and its `account_id`, +/// so the circuit must reject. Here `simple_balance_transfer` emits no claim for the +/// second account, leaving position 1 unbound. +#[test] +fn private_pda_without_binding_fails() { + let program = crate::test_methods::simple_balance_transfer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + let public_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + let private_pda_account = + AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32])); + + let result = execute_and_prove( + vec![public_account_1, private_pda_account], + Program::serialize_instruction(10_u128).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +/// Happy path: a program claims a new private PDA via `Claim::Pda(seed)`. The circuit +/// reads the npk for that `pre_state` from `private_account_keys` at the `pre_state`'s +/// position, derives `AccountId` via `AccountId::for_private_pda(program_id, seed, npk)`, and +/// asserts it equals the `pre_state`'s `account_id`. The equality both validates the claim +/// and binds the supplied npk to the `account_id`. +#[test] +fn private_pda_claim_succeeds() { + let program = crate::test_methods::pda_claimer(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([42; 32]); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, u128::MAX); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let result = execute_and_prove( + vec![pre_state], + Program::serialize_instruction(seed).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program.into(), + ); + + let (output, _proof) = result.expect("private PDA claim should succeed"); + assert_eq!(output.new_nullifiers.len(), 1); + assert_eq!(output.new_commitments.len(), 1); + assert_eq!(output.encrypted_private_post_states.len(), 1); + assert!(output.public_pre_states.is_empty()); + assert!(output.public_post_states.is_empty()); +} + +/// An npk is supplied that does not match the `pre_state`'s `account_id` under +/// `AccountId::for_private_pda(program, claim_seed, npk)`. The claim equality check rejects. +#[test] +fn private_pda_npk_mismatch_fails() { + // `keys_a` produces the `pre_state`'s `account_id` (the registered pair), `keys_b` is + // the mismatched pair supplied in `private_account_keys` for that pre_state. + let program = crate::test_methods::pda_claimer(); + let keys_a = test_private_account_keys_1(); + let keys_b = test_private_account_keys_2(); + let npk_a = keys_a.npk(); + let npk_b = keys_b.npk(); + let seed = PdaSeed::new([42; 32]); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys_b.vpk(), &[0_u8; 32], 0).0; + + // `account_id` is derived from `npk_a`, but `npk_b` is supplied for this pre_state. + // `AccountId::for_private_pda(program, seed, npk_b) != account_id`, so the claim check in + // the circuit must reject. + let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk_a, u128::MAX); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let result = execute_and_prove( + vec![pre_state], + Program::serialize_instruction(seed).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk_b, &keys_b.vpk()), + npk: npk_b, + ssk: shared_secret, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +/// Happy path for the caller-seeds authorization of a private PDA. The delegator claims a +/// private PDA via `Claim::Pda(seed)`, then chains to a callee (`noop`) delegating the same +/// seed via `ChainedCall.pda_seeds`. In the callee's step, the `pre_state`'s authorization +/// is established via the private derivation +/// `AccountId::for_private_pda(delegator, seed, npk) == pre.account_id`. +#[test] +fn caller_pda_seeds_authorize_private_pda_for_callee() { + let delegator = crate::test_methods::private_pda_delegator(); + let callee = crate::test_methods::auth_asserting_noop(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let seed = PdaSeed::new([77; 32]); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let account_id = AccountId::for_private_pda(&delegator.id(), &seed, &npk, u128::MAX); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let callee_id = callee.id(); + let program_with_deps = ProgramWithDependencies::new(delegator, [(callee_id, callee)].into()); + + let result = execute_and_prove( + vec![pre_state], + Program::serialize_instruction((seed, seed, callee_id)).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program_with_deps, + ); + + let (output, _proof) = + result.expect("caller-seeds authorization of private PDA should succeed"); + assert_eq!(output.new_commitments.len(), 1); + assert_eq!(output.new_nullifiers.len(), 1); +} + +/// The delegator chains with a different seed than the one it claimed with. In the callee +/// step, neither public nor private caller-seeds authorization matches; `pre.is_authorized` +/// was set to `true` by the delegator but no proven source supports it, so the consistency +/// assertion rejects. +#[test] +fn caller_pda_seeds_with_wrong_seed_rejects_private_pda_for_callee() { + let delegator = crate::test_methods::private_pda_delegator(); + let callee = crate::test_methods::auth_asserting_noop(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let claim_seed = PdaSeed::new([77; 32]); + let wrong_delegated_seed = PdaSeed::new([88; 32]); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + + let account_id = AccountId::for_private_pda(&delegator.id(), &claim_seed, &npk, u128::MAX); + let pre_state = AccountWithMetadata::new(Account::default(), false, account_id); + + let callee_id = callee.id(); + let program_with_deps = ProgramWithDependencies::new(delegator, [(callee_id, callee)].into()); + + let result = execute_and_prove( + vec![pre_state], + Program::serialize_instruction((claim_seed, wrong_delegated_seed, callee_id)).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program_with_deps, + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +/// Exploit-scenario pin. A single `(program_id, seed)` pair can derive a family of +/// `AccountId`s, one public PDA and one private PDA per distinct npk. Without the tx-wide +/// family-binding check, a program could claim `PDA_alice` (`alice_npk`) and +/// `PDA_bob` (`bob_npk`) under the same seed in one transaction, and once reuse +/// is supported a later chained call could delegate both to a callee via +/// `pda_seeds: [S]` and mix balances across them. The binding check rejects the setup +/// here: after the first claim records `(program, seed) → PDA_alice`, the second claim +/// tries to record `(program, seed) → PDA_bob` and panics. +#[test] +fn two_private_pda_claims_under_same_seed_are_rejected() { + let program = crate::test_methods::two_pda_claimer(); + let keys_a = test_private_account_keys_1(); + let keys_b = test_private_account_keys_2(); + let seed = PdaSeed::new([55; 32]); + let shared_a = SharedSecretKey::encapsulate_deterministic(&keys_a.vpk(), &[0_u8; 32], 0).0; + let shared_b = SharedSecretKey::encapsulate_deterministic(&keys_b.vpk(), &[0_u8; 32], 0).0; + + let account_a = AccountId::for_private_pda(&program.id(), &seed, &keys_a.npk(), u128::MAX); + let account_b = AccountId::for_private_pda(&program.id(), &seed, &keys_b.npk(), u128::MAX); + + let pre_a = AccountWithMetadata::new(Account::default(), false, account_a); + let pre_b = AccountWithMetadata::new(Account::default(), false, account_b); + + let result = execute_and_prove( + vec![pre_a, pre_b], + Program::serialize_instruction(seed).unwrap(), + vec![ + InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&keys_a.npk(), &keys_a.vpk()), + npk: keys_a.npk(), + ssk: shared_a, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }, + InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&keys_b.npk(), &keys_b.vpk()), + npk: keys_b.npk(), + ssk: shared_b, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +/// A private PDA that is reused at top level without an external seed in the identity still +/// fails binding. The noop program emits no `Claim::Pda` and there is no caller +/// `ChainedCall.pda_seeds`, so position 0 is never bound and the assertion fires. +/// Supplying `seed: Some((seed, owner_program_id))` in the `PrivatePdaUpdate` identity is +/// the correct path for top-level reuse; this test pins the failure when no seed is provided. +#[test] +fn private_pda_top_level_reuse_rejected_by_binding_check() { + let program = crate::test_methods::noop(); + let keys = test_private_account_keys_1(); + let npk = keys.npk(); + let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0; + let seed = PdaSeed::new([99; 32]); + + // Simulate a previously-claimed private PDA: program_owner != DEFAULT, is_authorized = + // true, account_id derived via the private formula. + let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, u128::MAX); + let owned_pre_state = AccountWithMetadata::new( + Account { + program_owner: program.id(), + ..Account::default() + }, + true, + account_id, + ); + + let result = execute_and_prove( + vec![owned_pre_state], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::PrivatePdaInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()), + npk, + ssk: shared_secret, + identifier: u128::MAX, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: None, + }], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn private_accounts_can_only_be_initialized_once() { + let sender_keys = test_private_account_keys_1(); + let sender_nonce = Nonce(0xdead_beef); + + let sender_private_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + nonce: sender_nonce, + data: Data::default(), + }; + let recipient_keys = test_private_account_keys_2(); + + let mut state = V03State::new().with_private_account(&sender_keys, &sender_private_account); + + let balance_to_move = 37; + let balance_to_move_2 = 30; + + let tx = private_balance_transfer_for_tests( + &sender_keys, + &sender_private_account, + &recipient_keys, + balance_to_move, + &state, + ); + + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .unwrap(); + + let sender_private_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + nonce: sender_nonce, + data: Data::default(), + }; + + let tx = private_balance_transfer_for_tests( + &sender_keys, + &sender_private_account, + &recipient_keys, + balance_to_move_2, + &state, + ); + + let result = state.transition_from_privacy_preserving_transaction(&tx, 1, 0); + + assert!(matches!(result, Err(LeeError::InvalidInput(_)))); + let LeeError::InvalidInput(error_message) = result.err().unwrap() else { + panic!("Incorrect message error"); + }; + let expected_error_message = "Nullifier already seen".to_owned(); + assert_eq!(error_message, expected_error_message); +} + +#[test] +fn circuit_should_fail_if_there_are_repeated_ids() { + let program = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_private_account_keys_1(); + let private_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + (&sender_keys.npk(), 0), + ); + + let shared_secret = + SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0; + let result = execute_and_prove( + vec![private_account_1.clone(), private_account_1], + Program::serialize_instruction(100_u128).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: shared_secret, + nsk: sender_keys.nsk, + membership_proof: (1, vec![]), + identifier: 0, + }, + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: shared_secret, + nsk: sender_keys.nsk, + membership_proof: (1, vec![]), + identifier: 0, + }, + ], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn private_authorized_uninitialized_account() { + let mut state = V03State::new().with_test_programs(); + + // Set up keys for the authorized private account + let private_keys = test_private_account_keys_1(); + + // Create an authorized private account with default values (new account being initialized) + let authorized_account = + AccountWithMetadata::new(Account::default(), true, (&private_keys.npk(), 0)); + + let program = crate::test_methods::simple_balance_transfer(); + + // Set up parameters for the new account + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0); + + let instruction: u128 = 0; + + // Execute and prove the circuit with the authorized account but no commitment proof + let (output, proof) = execute_and_prove( + vec![authorized_account], + Program::serialize_instruction(instruction).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedInit { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &private_keys.npk(), + &private_keys.vpk(), + ), + ssk: shared_secret, + nsk: private_keys.nsk, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &program.into(), + ) + .unwrap(); + + // Create message from circuit output + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + + let tx = PrivacyPreservingTransaction::new(message, witness_set); + let result = state.transition_from_privacy_preserving_transaction(&tx, 1, 0); + assert!(result.is_ok()); + + let account_id = AccountId::for_regular_private_account(&private_keys.npk(), 0); + let nullifier = Nullifier::for_account_initialization(&account_id); + assert!(state.private_state.1.contains(&nullifier)); +} + +#[test] +fn private_unauthorized_uninitialized_account_can_still_be_claimed() { + let mut state = V03State::new().with_test_programs(); + + let private_keys = test_private_account_keys_1(); + // This is intentional: claim authorization was introduced to protect public accounts, + // especially PDAs. Private PDAs are not useful in practice because there is no way to + // operate them without the corresponding private keys, so unauthorized private claiming + // remains allowed. + let unauthorized_account = + AccountWithMetadata::new(Account::default(), false, (&private_keys.npk(), 0)); + + let program = crate::test_methods::claimer(); + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0); + + let (output, proof) = execute_and_prove( + vec![unauthorized_account], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::PrivateUnauthorized { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &private_keys.npk(), + &private_keys.vpk(), + ), + npk: private_keys.npk(), + ssk: shared_secret, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .unwrap(); + + let account_id = AccountId::for_regular_private_account(&private_keys.npk(), 0); + let nullifier = Nullifier::for_account_initialization(&account_id); + assert!(state.private_state.1.contains(&nullifier)); +} + +#[test] +fn private_account_claimed_then_used_without_init_flag_should_fail() { + let mut state = V03State::new().with_test_programs(); + + // Set up keys for the private account + let private_keys = test_private_account_keys_1(); + + // Step 1: Create a new private account with authorization + let authorized_account = + AccountWithMetadata::new(Account::default(), true, (&private_keys.npk(), 0)); + + let claimer_program = crate::test_methods::claimer(); + + // Set up parameters for claiming the new account + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0); + + let instruction = (); + + // Step 2: Execute claimer program to claim the account with authentication + let (output, proof) = execute_and_prove( + vec![authorized_account.clone()], + Program::serialize_instruction(instruction).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedInit { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &private_keys.npk(), + &private_keys.vpk(), + ), + ssk: shared_secret, + nsk: private_keys.nsk, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &claimer_program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + // Claim should succeed + assert!( + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .is_ok() + ); + + // Verify the account is now initialized (nullifier exists) + let account_id = AccountId::for_regular_private_account(&private_keys.npk(), 0); + let nullifier = Nullifier::for_account_initialization(&account_id); + assert!(state.private_state.1.contains(&nullifier)); + + // Prepare new state of account + let account_metadata = { + let mut acc = authorized_account; + acc.account.program_owner = crate::test_methods::claimer().id(); + acc + }; + + let noop_program = crate::test_methods::noop(); + let shared_secret2 = + SharedSecretKey::encapsulate_deterministic(&private_keys.vpk(), &[0_u8; 32], 0).0; + + // Step 3: Try to execute noop program with authentication but without initialization + let res = execute_and_prove( + vec![account_metadata], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::PrivateAuthorizedInit { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &private_keys.npk(), + &private_keys.vpk(), + ), + ssk: shared_secret2, + nsk: private_keys.nsk, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &noop_program.into(), + ); + + assert!(matches!(res, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn two_private_pda_family_members_receive_and_spend() { + let funder_keys = test_public_account_keys_1(); + let alice_keys = test_private_account_keys_1(); + let alice_npk = alice_keys.npk(); + + let proxy = crate::test_methods::pda_spend_proxy(); + let simple_transfer = crate::test_methods::simple_balance_transfer(); + let proxy_id = proxy.id(); + let simple_transfer_id = simple_transfer.id(); + let seed = PdaSeed::new([42; 32]); + let amount: u128 = 100; + + let spend_with_deps = ProgramWithDependencies::new( + proxy, + [(simple_transfer_id, simple_transfer.clone())].into(), + ); + + let funder_id = funder_keys.account_id(); + let alice_pda_0_id = AccountId::for_private_pda(&proxy_id, &seed, &alice_npk, 0); + let alice_pda_1_id = AccountId::for_private_pda(&proxy_id, &seed, &alice_npk, 1); + let recipient_id = test_public_account_keys_2().account_id(); + let recipient_signing_key = test_public_account_keys_2().signing_key; + + let mut state = + V03State::new().with_public_accounts(public_state_from_balances(&[(funder_id, 500)])); + + let alice_pda_0_account = Account { + program_owner: simple_transfer_id, + balance: amount, + nonce: Nonce::private_account_nonce_init(&alice_pda_0_id), + ..Account::default() + }; + let alice_pda_1_account = Account { + program_owner: simple_transfer_id, + balance: amount, + nonce: Nonce::private_account_nonce_init(&alice_pda_1_id), + ..Account::default() + }; + + let (alice_shared_0, alice_epk_0) = + SharedSecretKey::encapsulate_deterministic(&alice_keys.vpk(), &[0_u8; 32], 0); + let (alice_shared_1, alice_epk_1) = + SharedSecretKey::encapsulate_deterministic(&alice_keys.vpk(), &[0_u8; 32], 1); + + // Fund alice_pda_0 via authenticated_transfer directly. + { + let funder_account = state.get_account_by_id(funder_id); + let funder_nonce = funder_account.nonce; + let (output, proof) = execute_and_prove( + vec![ + AccountWithMetadata::new(funder_account, true, funder_id), + AccountWithMetadata::new(Account::default(), false, alice_pda_0_id), + ], + Program::serialize_instruction(amount).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivatePdaInit { + epk: alice_epk_0.clone(), + view_tag: EncryptedAccountData::compute_view_tag(&alice_npk, &alice_keys.vpk()), + npk: alice_npk, + ssk: alice_shared_0, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: Some((seed, proxy_id)), + }, + ], + &simple_transfer.clone().into(), + ) + .unwrap(); + let message = + Message::try_from_circuit_output(vec![funder_id], vec![funder_nonce], output).unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[&funder_keys.signing_key]); + state + .transition_from_privacy_preserving_transaction( + &PrivacyPreservingTransaction::new(message, witness_set), + 1, + 0, + ) + .unwrap(); + } + + // Fund alice_pda_1 the same way with identifier 1. + { + let funder_account = state.get_account_by_id(funder_id); + let funder_nonce = funder_account.nonce; + let (output, proof) = execute_and_prove( + vec![ + AccountWithMetadata::new(funder_account, true, funder_id), + AccountWithMetadata::new(Account::default(), false, alice_pda_1_id), + ], + Program::serialize_instruction(amount).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivatePdaInit { + epk: alice_epk_1.clone(), + view_tag: EncryptedAccountData::compute_view_tag(&alice_npk, &alice_keys.vpk()), + npk: alice_npk, + ssk: alice_shared_1, + identifier: 1, + commitment_root: DUMMY_COMMITMENT_HASH, + seed: Some((seed, proxy_id)), + }, + ], + &simple_transfer.into(), + ) + .unwrap(); + let message = + Message::try_from_circuit_output(vec![funder_id], vec![funder_nonce], output).unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[&funder_keys.signing_key]); + state + .transition_from_privacy_preserving_transaction( + &PrivacyPreservingTransaction::new(message, witness_set), + 2, + 0, + ) + .unwrap(); + } + + let commitment_pda_0 = Commitment::new(&alice_pda_0_id, &alice_pda_0_account); + let commitment_pda_1 = Commitment::new(&alice_pda_1_id, &alice_pda_1_account); + + assert!(state.get_proof_for_commitment(&commitment_pda_0).is_some()); + assert!(state.get_proof_for_commitment(&commitment_pda_1).is_some()); + + // Alice spends alice_pda_0 into the public recipient. + { + let recipient_account = state.get_account_by_id(recipient_id); + let (output, proof) = execute_and_prove( + vec![ + AccountWithMetadata::new(alice_pda_0_account, true, alice_pda_0_id), + AccountWithMetadata::new(recipient_account, true, recipient_id), + ], + Program::serialize_instruction((seed, amount, simple_transfer_id)).unwrap(), + vec![ + InputAccountIdentity::PrivatePdaUpdate { + epk: alice_epk_0, + view_tag: EncryptedAccountData::compute_view_tag(&alice_npk, &alice_keys.vpk()), + ssk: alice_shared_0, + nsk: alice_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&commitment_pda_0) + .expect("pda_0 must be in state"), + identifier: 0, + seed: None, + }, + InputAccountIdentity::Public, + ], + &spend_with_deps, + ) + .unwrap(); + let message = + Message::try_from_circuit_output(vec![recipient_id], vec![Nonce(0)], output).unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_signing_key]); + state + .transition_from_privacy_preserving_transaction( + &PrivacyPreservingTransaction::new(message, witness_set), + 3, + 0, + ) + .unwrap(); + } + + // Alice spends alice_pda_1 into the same public recipient. + { + let recipient_account = state.get_account_by_id(recipient_id); + let (output, proof) = execute_and_prove( + vec![ + AccountWithMetadata::new(alice_pda_1_account.clone(), true, alice_pda_1_id), + AccountWithMetadata::new(recipient_account, false, recipient_id), + ], + Program::serialize_instruction((seed, amount, simple_transfer_id)).unwrap(), + vec![ + InputAccountIdentity::PrivatePdaUpdate { + epk: alice_epk_1, + view_tag: EncryptedAccountData::compute_view_tag(&alice_npk, &alice_keys.vpk()), + ssk: alice_shared_1, + nsk: alice_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&commitment_pda_1) + .expect("pda_1 must be in state"), + identifier: 1, + seed: None, + }, + InputAccountIdentity::Public, + ], + &spend_with_deps, + ) + .unwrap(); + let message = Message::try_from_circuit_output(vec![recipient_id], vec![], output).unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[]); + state + .transition_from_privacy_preserving_transaction( + &PrivacyPreservingTransaction::new(message, witness_set), + 4, + 0, + ) + .unwrap(); + } + + assert_eq!(state.get_account_by_id(recipient_id).balance, 2 * amount); + + // Re-fund alice_pda_1 top-level via simple_transfer using PrivatePdaUpdate with an + // external seed. + let alice_pda_1_account_after_spend = Account { + program_owner: simple_transfer_id, + balance: 0, + nonce: alice_pda_1_account + .nonce + .private_account_nonce_increment(&alice_keys.nsk), + ..Account::default() + }; + let commitment_pda_1_after_spend = + Commitment::new(&alice_pda_1_id, &alice_pda_1_account_after_spend); + let alice_shared_1_refund = SharedSecretKey([12; 32]); + { + let recipient_account = state.get_account_by_id(recipient_id); + let recipient_nonce = recipient_account.nonce; + let (output, proof) = execute_and_prove( + vec![ + AccountWithMetadata::new(recipient_account, true, recipient_id), + AccountWithMetadata::new(alice_pda_1_account_after_spend, false, alice_pda_1_id), + ], + Program::serialize_instruction(amount).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivatePdaUpdate { + epk: EphemeralPublicKey(vec![12_u8; ML_KEM_768_CIPHERTEXT_LEN]), + view_tag: EncryptedAccountData::compute_view_tag(&alice_npk, &alice_keys.vpk()), + nsk: alice_keys.nsk, + ssk: alice_shared_1_refund, + membership_proof: state + .get_proof_for_commitment(&commitment_pda_1_after_spend) + .expect("pda_1 after spend must be in state"), + identifier: 1, + seed: Some((seed, proxy_id)), + }, + ], + &crate::test_methods::simple_balance_transfer().into(), + ) + .unwrap(); + let message = + Message::try_from_circuit_output(vec![recipient_id], vec![recipient_nonce], output) + .unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_signing_key]); + state + .transition_from_privacy_preserving_transaction( + &PrivacyPreservingTransaction::new(message, witness_set), + 5, + 0, + ) + .unwrap(); + } + + assert_eq!(state.get_account_by_id(recipient_id).balance, amount); +} diff --git a/lee/state_machine/src/state/tests/claiming.rs b/lee/state_machine/src/state/tests/claiming.rs new file mode 100644 index 00000000..aba5c447 --- /dev/null +++ b/lee/state_machine/src/state/tests/claiming.rs @@ -0,0 +1,619 @@ +use super::*; + +#[test] +fn claiming_mechanism() { + let program = crate::test_methods::simple_balance_transfer(); + let from_key = PrivateKey::try_new([1; 32]).unwrap(); + let from = AccountId::from(&PublicKey::new_from_private_key(&from_key)); + let initial_balance = 100; + let initial_data = [(from, initial_balance)]; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let to_key = PrivateKey::try_new([2; 32]).unwrap(); + let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); + let amount: u128 = 37; + + // Check the recipient is an uninitialized account + assert_eq!(state.get_account_by_id(to), Account::default()); + + let expected_recipient_post = Account { + program_owner: program.id(), + balance: amount, + nonce: Nonce(1), + ..Account::default() + }; + + let message = public_transaction::Message::try_new( + program.id(), + vec![from, to], + vec![Nonce(0), Nonce(0)], + amount, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + let recipient_post = state.get_account_by_id(to); + + assert_eq!(recipient_post, expected_recipient_post); +} + +#[test] +fn unauthorized_public_account_claiming_fails() { + let program = crate::test_methods::simple_balance_transfer(); + let account_key = PrivateKey::try_new([9; 32]).unwrap(); + let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key)); + let mut state = V03State::new().with_test_programs(); + + assert_eq!(state.get_account_by_id(account_id), Account::default()); + + let message = + public_transaction::Message::try_new(program.id(), vec![account_id], vec![], 0_u128) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 2, 0); + + assert!(matches!(result, Err(LeeError::InvalidProgramBehavior(_)))); + assert_eq!(state.get_account_by_id(account_id), Account::default()); +} + +#[test] +fn authorized_public_account_claiming_succeeds() { + let program = crate::test_methods::simple_balance_transfer(); + let account_key = PrivateKey::try_new([10; 32]).unwrap(); + let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key)); + let mut state = V03State::new().with_test_programs(); + + assert_eq!(state.get_account_by_id(account_id), Account::default()); + + let message = public_transaction::Message::try_new( + program.id(), + vec![account_id], + vec![Nonce(0)], + 0_u128, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&account_key]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + assert_eq!( + state.get_account_by_id(account_id), + Account { + program_owner: program.id(), + nonce: Nonce(1), + ..Account::default() + } + ); +} + +#[test] +fn public_chained_call() { + let program = crate::test_methods::chain_caller(); + let key = PrivateKey::try_new([1; 32]).unwrap(); + let from = AccountId::from(&PublicKey::new_from_private_key(&key)); + let to = AccountId::new([2; 32]); + let initial_balance = 1000; + let initial_data = [(from, initial_balance), (to, 0)]; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let from_key = key; + let amount: u128 = 37; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + crate::test_methods::simple_balance_transfer().id(), + 2, + None, + ); + + let expected_to_post = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: amount * 2, // The `chain_caller` chains the program twice + ..Account::default() + }; + + let message = public_transaction::Message::try_new( + program.id(), + vec![to, from], // The chain_caller program permutes the account order in the chain + // call + vec![Nonce(0)], + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + let from_post = state.get_account_by_id(from); + let to_post = state.get_account_by_id(to); + // The `chain_caller` program calls the program twice + assert_eq!(from_post.balance, initial_balance - 2 * amount); + assert_eq!(to_post, expected_to_post); +} + +#[test] +fn execution_fails_if_chained_calls_exceeds_depth() { + let program = crate::test_methods::chain_caller(); + let key = PrivateKey::try_new([1; 32]).unwrap(); + let from = AccountId::from(&PublicKey::new_from_private_key(&key)); + let to = AccountId::new([2; 32]); + let initial_balance = 100; + let initial_data = [(from, initial_balance), (to, 0)]; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let from_key = key; + let amount: u128 = 0; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + crate::test_methods::simple_balance_transfer().id(), + u32::try_from(MAX_NUMBER_CHAINED_CALLS).expect("MAX_NUMBER_CHAINED_CALLS fits in u32") + 1, + None, + ); + + let message = public_transaction::Message::try_new( + program.id(), + vec![to, from], // The chain_caller program permutes the account order in the chain + // call + vec![Nonce(0)], + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!(matches!( + result, + Err(LeeError::MaxChainedCallsDepthExceeded) + )); +} + +#[test] +fn execution_that_requires_authentication_of_a_program_derived_account_id_succeeds() { + let chain_caller = crate::test_methods::chain_caller(); + let pda_seed = PdaSeed::new([37; 32]); + let from = AccountId::for_public_pda(&chain_caller.id(), &pda_seed); + let to = AccountId::new([2; 32]); + let initial_balance = 1000; + let initial_data = [(from, initial_balance), (to, 0)]; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let amount: u128 = 58; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + crate::test_methods::simple_balance_transfer().id(), + 1, + Some(pda_seed), + ); + + let expected_to_post = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: amount, // The `chain_caller` chains the program twice + ..Account::default() + }; + let message = public_transaction::Message::try_new( + chain_caller.id(), + vec![to, from], // The chain_caller program permutes the account order in the chain + // call + vec![], + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + let from_post = state.get_account_by_id(from); + let to_post = state.get_account_by_id(to); + assert_eq!(from_post.balance, initial_balance - amount); + assert_eq!(to_post, expected_to_post); +} + +#[test] +fn claiming_mechanism_within_chain_call() { + // This test calls the authenticated transfer program through the chain_caller program. + // The transfer is made from an initialized sender to an uninitialized recipient. And + // it is expected that the recipient account is claimed by the authenticated transfer + // program and not the chained_caller program. + let chain_caller = crate::test_methods::chain_caller(); + let simple_transfer = crate::test_methods::simple_balance_transfer(); + let from_key = PrivateKey::try_new([1; 32]).unwrap(); + let from = AccountId::from(&PublicKey::new_from_private_key(&from_key)); + let initial_balance = 100; + let initial_data = [(from, initial_balance)]; + let mut state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let to_key = PrivateKey::try_new([2; 32]).unwrap(); + let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); + let amount: u128 = 37; + + // Check the recipient is an uninitialized account + assert_eq!(state.get_account_by_id(to), Account::default()); + + let expected_to_post = Account { + // The expected program owner is the authenticated transfer program + program_owner: simple_transfer.id(), + balance: amount, + nonce: Nonce(1), + ..Account::default() + }; + + // The transaction executes the chain_caller program, which internally calls the + // authenticated_transfer program + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + crate::test_methods::simple_balance_transfer().id(), + 1, + None, + ); + let message = public_transaction::Message::try_new( + chain_caller.id(), + vec![to, from], // The chain_caller program permutes the account order in the chain + // call + vec![Nonce(0), Nonce(0)], + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + let from_post = state.get_account_by_id(from); + let to_post = state.get_account_by_id(to); + assert_eq!(from_post.balance, initial_balance - amount); + assert_eq!(to_post, expected_to_post); +} + +#[test] +fn unauthorized_public_account_claiming_fails_when_executed_privately() { + let program = crate::test_methods::simple_balance_transfer(); + let account_id = AccountId::new([11; 32]); + let public_account = AccountWithMetadata::new(Account::default(), false, account_id); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(0_u128).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn authorized_public_account_claiming_succeeds_when_executed_privately() { + let program = crate::test_methods::simple_balance_transfer(); + let program_id = program.id(); + let sender_keys = test_private_account_keys_1(); + let sender_private_account = Account { + program_owner: program_id, + balance: 100, + ..Account::default() + }; + let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); + let sender_commitment = Commitment::new(&sender_account_id, &sender_private_account); + let sender_init_nullifier = Nullifier::for_account_initialization(&sender_account_id); + let mut state = + V03State::new().with_private_accounts([(sender_commitment.clone(), sender_init_nullifier)]); + let sender_pre = + AccountWithMetadata::new(sender_private_account, true, (&sender_keys.npk(), 0)); + let recipient_private_key = PrivateKey::try_new([2; 32]).unwrap(); + let recipient_account_id = + AccountId::from(&PublicKey::new_from_private_key(&recipient_private_key)); + let recipient_pre = AccountWithMetadata::new(Account::default(), true, recipient_account_id); + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0); + + let balance = 37; + + let (output, proof) = execute_and_prove( + vec![sender_pre, recipient_pre], + Program::serialize_instruction(balance).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: shared_secret, + nsk: sender_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&sender_commitment) + .expect("sender's commitment must be in state"), + identifier: 0, + }, + InputAccountIdentity::Public, + ], + &program.into(), + ) + .unwrap(); + + let message = + Message::try_from_circuit_output(vec![recipient_account_id], vec![Nonce(0)], output) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_private_key]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .unwrap(); + + let nullifier = Nullifier::for_account_update(&sender_commitment, &sender_keys.nsk); + assert!(state.private_state.1.contains(&nullifier)); + + assert_eq!( + state.get_account_by_id(recipient_account_id), + Account { + program_owner: program_id, + balance, + nonce: Nonce(1), + ..Account::default() + } + ); +} + +#[test_case::test_case(1; "single call")] +#[test_case::test_case(2; "two calls")] +fn private_chained_call(number_of_calls: u32) { + // Arrange + let chain_caller = crate::test_methods::chain_caller(); + let simple_transfers = crate::test_methods::simple_balance_transfer(); + let from_keys = test_private_account_keys_1(); + let to_keys = test_private_account_keys_2(); + let initial_balance = 100; + let from_account = AccountWithMetadata::new( + Account { + program_owner: simple_transfers.id(), + balance: initial_balance, + ..Account::default() + }, + true, + (&from_keys.npk(), 0), + ); + let to_account = AccountWithMetadata::new( + Account { + program_owner: simple_transfers.id(), + ..Account::default() + }, + true, + (&to_keys.npk(), 0), + ); + + let from_account_id = AccountId::for_regular_private_account(&from_keys.npk(), 0); + let to_account_id = AccountId::for_regular_private_account(&to_keys.npk(), 0); + let from_commitment = Commitment::new(&from_account_id, &from_account.account); + let to_commitment = Commitment::new(&to_account_id, &to_account.account); + let from_init_nullifier = Nullifier::for_account_initialization(&from_account_id); + let to_init_nullifier = Nullifier::for_account_initialization(&to_account_id); + let mut state = V03State::new() + .with_private_accounts([ + (from_commitment.clone(), from_init_nullifier), + (to_commitment.clone(), to_init_nullifier), + ]) + .with_test_programs(); + let amount: u128 = 37; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + crate::test_methods::simple_balance_transfer().id(), + number_of_calls, + None, + ); + + let (from_ss, from_epk) = + SharedSecretKey::encapsulate_deterministic(&from_keys.vpk(), &[0_u8; 32], 0); + + let (to_ss, to_epk) = + SharedSecretKey::encapsulate_deterministic(&to_keys.vpk(), &[0_u8; 32], 1); + + let mut dependencies = HashMap::new(); + + dependencies.insert(simple_transfers.id(), simple_transfers); + let program_with_deps = ProgramWithDependencies::new(chain_caller, dependencies); + + let from_new_nonce = Nonce::default().private_account_nonce_increment(&from_keys.nsk); + let to_new_nonce = Nonce::default().private_account_nonce_increment(&to_keys.nsk); + + let from_expected_post = Account { + balance: initial_balance - u128::from(number_of_calls) * amount, + nonce: from_new_nonce, + ..from_account.account.clone() + }; + let from_expected_commitment = Commitment::new(&from_account_id, &from_expected_post); + + let to_expected_post = Account { + balance: u128::from(number_of_calls) * amount, + nonce: to_new_nonce, + ..to_account.account.clone() + }; + let to_expected_commitment = Commitment::new(&to_account_id, &to_expected_post); + + // Act + let (output, proof) = execute_and_prove( + vec![to_account, from_account], + Program::serialize_instruction(instruction).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: to_epk, + view_tag: EncryptedAccountData::compute_view_tag(&to_keys.npk(), &to_keys.vpk()), + ssk: to_ss, + nsk: from_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&from_commitment) + .expect("from's commitment must be in state"), + identifier: 0, + }, + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: from_epk, + view_tag: EncryptedAccountData::compute_view_tag( + &from_keys.npk(), + &from_keys.vpk(), + ), + ssk: from_ss, + nsk: to_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&to_commitment) + .expect("to's commitment must be in state"), + identifier: 0, + }, + ], + &program_with_deps, + ) + .unwrap(); + + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + let witness_set = WitnessSet::for_message(&message, proof, &[]); + let transaction = PrivacyPreservingTransaction::new(message, witness_set); + + state + .transition_from_privacy_preserving_transaction(&transaction, 1, 0) + .unwrap(); + + // Assert + assert!( + state + .get_proof_for_commitment(&from_expected_commitment) + .is_some() + ); + assert!( + state + .get_proof_for_commitment(&to_expected_commitment) + .is_some() + ); +} + +#[test] +fn claiming_mechanism_cannot_claim_initialied_accounts() { + let claimer = crate::test_methods::claimer(); + let mut state = V03State::new().with_test_programs(); + let account_id = AccountId::new([2; 32]); + + // Insert an account with non-default program owner + state.force_insert_account( + account_id, + Account { + program_owner: [1, 2, 3, 4, 5, 6, 7, 8], + ..Account::default() + }, + ); + + let message = + public_transaction::Message::try_new(claimer.id(), vec![account_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id: err_account_id } + )) if err_account_id == account_id + )); +} + +/// This test ensures that even if a malicious program tries to perform overflow of balances +/// it will not be able to break the balance validation. +#[test] +fn malicious_program_cannot_break_balance_validation_if_not_in_genesis() { + let sender_key = PrivateKey::try_new([37; 32]).unwrap(); + let sender_id = AccountId::from(&PublicKey::new_from_private_key(&sender_key)); + let sender_init_balance: u128 = 10; + + let recipient_key = PrivateKey::try_new([42; 32]).unwrap(); + let recipient_id = AccountId::from(&PublicKey::new_from_private_key(&recipient_key)); + let recipient_init_balance: u128 = 10; + + let modified_transfer_id = crate::test_methods::modified_transfer_program().id(); + + let mut state = V03State::new() + .with_public_accounts([ + ( + sender_id, + Account { + program_owner: modified_transfer_id, + balance: sender_init_balance, + ..Account::default() + }, + ), + ( + recipient_id, + Account { + program_owner: modified_transfer_id, + balance: recipient_init_balance, + ..Account::default() + }, + ), + ]) + .with_test_programs(); + + let balance_to_move: u128 = 4; + + let sender = AccountWithMetadata::new(state.get_account_by_id(sender_id), true, sender_id); + + let sender_nonce = sender.account.nonce; + + let _recipient = + AccountWithMetadata::new(state.get_account_by_id(recipient_id), false, sender_id); + + let message = public_transaction::Message::try_new( + modified_transfer_id, + vec![sender_id, recipient_id], + vec![sender_nonce], + balance_to_move, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&sender_key]); + let tx = PublicTransaction::new(message, witness_set); + let res = state.transition_from_public_transaction(&tx, 2, 0); + let expected_total_balance_pre_states = + WrappedBalanceSum::from_balances([sender_init_balance, recipient_init_balance].into_iter()) + .unwrap(); + let expected_total_balance_post_states = WrappedBalanceSum::from_balances( + [sender_init_balance, recipient_init_balance, u128::MAX, 1].into_iter(), + ) + .unwrap(); + assert!(matches!( + res, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states } + ) + )) if total_balance_pre_states == expected_total_balance_pre_states && total_balance_post_states == expected_total_balance_post_states + )); + + let sender_post = state.get_account_by_id(sender_id); + let recipient_post = state.get_account_by_id(recipient_id); + + let expected_sender_post = { + let mut this = state.get_account_by_id(sender_id); + this.balance = sender_init_balance; + this.nonce = Nonce(0); + this + }; + + let expected_recipient_post = { + let mut this = state.get_account_by_id(sender_id); + this.balance = recipient_init_balance; + this.nonce = Nonce(0); + this + }; + + assert_eq!(expected_sender_post, sender_post); + assert_eq!(expected_recipient_post, recipient_post); +} diff --git a/lee/state_machine/src/state/tests/flash_swap.rs b/lee/state_machine/src/state/tests/flash_swap.rs new file mode 100644 index 00000000..be8f1c10 --- /dev/null +++ b/lee/state_machine/src/state/tests/flash_swap.rs @@ -0,0 +1,235 @@ +use super::*; + +#[test] +fn flash_swap_successful() { + let initiator = crate::test_methods::flash_swap_initiator(); + let callback = crate::test_methods::flash_swap_callback(); + let token = crate::test_methods::simple_balance_transfer(); + + let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); + let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32])); + + let initial_balance: u128 = 1000; + let amount_out: u128 = 100; + + let vault_account = Account { + program_owner: token.id(), + balance: initial_balance, + ..Account::default() + }; + let receiver_account = Account { + program_owner: token.id(), + balance: 0, + ..Account::default() + }; + + let mut state = V03State::new().with_test_programs(); + state.force_insert_account(vault_id, vault_account); + state.force_insert_account(receiver_id, receiver_account); + + // Callback instruction: return funds + let cb_instruction = CallbackInstruction { + return_funds: true, + token_program_id: token.id(), + amount: amount_out, + }; + let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); + + let instruction = FlashSwapInstruction::Initiate { + token_program_id: token.id(), + callback_program_id: callback.id(), + amount_out, + callback_instruction_data: cb_data, + }; + + let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!(result.is_ok(), "flash swap should succeed: {result:?}"); + + // Vault balance restored, receiver back to 0 + assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance); + assert_eq!(state.get_account_by_id(receiver_id).balance, 0); +} + +#[test] +fn flash_swap_callback_keeps_funds_rollback() { + let initiator = crate::test_methods::flash_swap_initiator(); + let callback = crate::test_methods::flash_swap_callback(); + let token = crate::test_methods::simple_balance_transfer(); + + let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); + let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32])); + + let initial_balance: u128 = 1000; + let amount_out: u128 = 100; + + let vault_account = Account { + program_owner: token.id(), + balance: initial_balance, + ..Account::default() + }; + let receiver_account = Account { + program_owner: token.id(), + balance: 0, + ..Account::default() + }; + + let mut state = V03State::new().with_test_programs(); + state.force_insert_account(vault_id, vault_account); + state.force_insert_account(receiver_id, receiver_account); + + // Callback instruction: do NOT return funds + let cb_instruction = CallbackInstruction { + return_funds: false, + token_program_id: token.id(), + amount: amount_out, + }; + let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); + + let instruction = FlashSwapInstruction::Initiate { + token_program_id: token.id(), + callback_program_id: callback.id(), + amount_out, + callback_instruction_data: cb_data, + }; + + let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); + let result = state.transition_from_public_transaction(&tx, 1, 0); + + // Invariant check fails → entire tx rolls back + assert!( + result.is_err(), + "flash swap should fail when callback keeps funds" + ); + + // State unchanged (rollback) + assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance); + assert_eq!(state.get_account_by_id(receiver_id).balance, 0); +} + +#[test] +fn flash_swap_self_call_targets_correct_program() { + // Zero-amount flash swap: the invariant self-call still runs and succeeds + // because vault balance doesn't decrease. + let initiator = crate::test_methods::flash_swap_initiator(); + let callback = crate::test_methods::flash_swap_callback(); + let token = crate::test_methods::simple_balance_transfer(); + + let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); + let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32])); + + let initial_balance: u128 = 1000; + + let vault_account = Account { + program_owner: token.id(), + balance: initial_balance, + ..Account::default() + }; + let receiver_account = Account { + program_owner: token.id(), + balance: 0, + ..Account::default() + }; + + let mut state = V03State::new().with_test_programs(); + state.force_insert_account(vault_id, vault_account); + state.force_insert_account(receiver_id, receiver_account); + + let cb_instruction = CallbackInstruction { + return_funds: true, + token_program_id: token.id(), + amount: 0, + }; + let cb_data = Program::serialize_instruction(cb_instruction).unwrap(); + + let instruction = FlashSwapInstruction::Initiate { + token_program_id: token.id(), + callback_program_id: callback.id(), + amount_out: 0, + callback_instruction_data: cb_data, + }; + + let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction); + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!( + result.is_ok(), + "zero-amount flash swap should succeed: {result:?}" + ); +} + +#[test] +fn flash_swap_standalone_invariant_check_rejected() { + // Calling InvariantCheck directly (not as a chained self-call) should fail + // because caller_program_id will be None. + let initiator = crate::test_methods::flash_swap_initiator(); + let token = crate::test_methods::simple_balance_transfer(); + + let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32])); + + let vault_account = Account { + program_owner: token.id(), + balance: 1000, + ..Account::default() + }; + + let mut state = V03State::new().with_test_programs(); + state.force_insert_account(vault_id, vault_account); + + let instruction = FlashSwapInstruction::InvariantCheck { + min_vault_balance: 1000, + }; + + let message = + public_transaction::Message::try_new(initiator.id(), vec![vault_id], vec![], instruction) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!( + result.is_err(), + "standalone InvariantCheck should be rejected (caller_program_id is None)" + ); +} + +#[test] +fn malicious_self_program_id_rejected_in_public_execution() { + let program = crate::test_methods::malicious_self_program_id(); + let acc_id = AccountId::new([99; 32]); + let account = Account::default(); + + let mut state = V03State::new().with_test_programs(); + state.force_insert_account(acc_id, account); + + let message = + public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!( + result.is_err(), + "program with wrong self_program_id in output should be rejected" + ); +} + +#[test] +fn malicious_caller_program_id_rejected_in_public_execution() { + let program = crate::test_methods::malicious_caller_program_id(); + let acc_id = AccountId::new([99; 32]); + let account = Account::default(); + + let mut state = V03State::new().with_test_programs(); + state.force_insert_account(acc_id, account); + + let message = + public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + assert!( + result.is_err(), + "program with spoofed caller_program_id in output should be rejected" + ); +} diff --git a/lee/state_machine/src/state/tests/genesis.rs b/lee/state_machine/src/state/tests/genesis.rs new file mode 100644 index 00000000..d782f325 --- /dev/null +++ b/lee/state_machine/src/state/tests/genesis.rs @@ -0,0 +1,128 @@ +use super::*; + +#[test] +fn new_works() { + let key1 = PrivateKey::try_new([1; 32]).unwrap(); + let key2 = PrivateKey::try_new([2; 32]).unwrap(); + let addr1 = AccountId::from(&PublicKey::new_from_private_key(&key1)); + let addr2 = AccountId::from(&PublicKey::new_from_private_key(&key2)); + let expected_public_state = { + let mut this = HashMap::new(); + this.insert( + addr1, + Account { + balance: 100, + ..Account::default() + }, + ); + this.insert( + addr2, + Account { + balance: 151, + ..Account::default() + }, + ); + this + }; + let expected_builtin_programs = HashMap::new(); + + let state = + V03State::new().with_public_account_balances([(addr1, 100_u128), (addr2, 151_u128)]); + + assert_eq!(state.public_state, expected_public_state); + assert_eq!(state.programs, expected_builtin_programs); +} + +#[test] +fn new_includes_nullifiers_for_private_accounts() { + let keys1 = test_private_account_keys_1(); + let keys2 = test_private_account_keys_2(); + + let account = Account { + balance: 100, + ..Account::default() + }; + + let account_id1 = AccountId::for_regular_private_account(&keys1.npk(), 0); + let account_id2 = AccountId::for_regular_private_account(&keys2.npk(), 0); + + let init_commitment1 = Commitment::new(&account_id1, &account); + let init_commitment2 = Commitment::new(&account_id2, &account); + let init_nullifier1 = Nullifier::for_account_initialization(&account_id1); + let init_nullifier2 = Nullifier::for_account_initialization(&account_id2); + + let initial_private_accounts = vec![ + (init_commitment1, init_nullifier1), + (init_commitment2, init_nullifier2), + ]; + + let state = V03State::new().with_private_accounts(initial_private_accounts); + + assert!(state.private_state.1.contains(&init_nullifier1)); + assert!(state.private_state.1.contains(&init_nullifier2)); +} + +#[test] +fn insert_program() { + let mut state = V03State::new(); + let program_to_insert = crate::test_methods::simple_balance_transfer(); + let program_id = program_to_insert.id(); + assert!(!state.programs.contains_key(&program_id)); + + state.insert_program(program_to_insert); + + assert!(state.programs.contains_key(&program_id)); +} + +#[test] +fn get_account_by_account_id_non_default_account() { + let key = PrivateKey::try_new([1; 32]).unwrap(); + let account_id = AccountId::from(&PublicKey::new_from_private_key(&key)); + let initial_data = [( + account_id, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + ..Account::default() + }, + )]; + let state = V03State::new().with_public_accounts(initial_data); + let expected_account = &state.public_state[&account_id]; + + let account = state.get_account_by_id(account_id); + + assert_eq!(&account, expected_account); +} + +#[test] +fn get_account_by_account_id_default_account() { + let addr2 = AccountId::new([0; 32]); + let state = V03State::new(); + let expected_account = Account::default(); + + let account = state.get_account_by_id(addr2); + + assert_eq!(account, expected_account); +} + +#[test] +fn builtin_programs_getter() { + let state = V03State::new(); + + let builtin_programs = state.programs(); + + assert_eq!(builtin_programs, &state.programs); +} + +#[test] +fn state_serialization_roundtrip() { + let account_id_1 = AccountId::new([1; 32]); + let account_id_2 = AccountId::new([2; 32]); + let initial_data = [(account_id_1, 100_u128), (account_id_2, 151_u128)]; + let state = V03State::new() + .with_public_accounts(public_state_from_balances(&initial_data)) + .with_test_programs(); + let bytes = borsh::to_vec(&state).unwrap(); + let state_from_bytes: V03State = borsh::from_slice(&bytes).unwrap(); + assert_eq!(state, state_from_bytes); +} diff --git a/lee/state_machine/src/state/tests/mod.rs b/lee/state_machine/src/state/tests/mod.rs new file mode 100644 index 00000000..f6b56be6 --- /dev/null +++ b/lee/state_machine/src/state/tests/mod.rs @@ -0,0 +1,444 @@ +#![expect( + clippy::arithmetic_side_effects, + clippy::shadow_unrelated, + reason = "We don't care about it in tests" +)] + +use std::collections::HashMap; + +use lee_core::{ + BlockId, Commitment, DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, + Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, Timestamp, + account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data}, + encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey}, + program::{ + BlockValidityWindow, ExecutionValidationError, MAX_NUMBER_CHAINED_CALLS, PdaSeed, + ProgramId, TimestampValidityWindow, WrappedBalanceSum, + }, +}; + +use crate::{ + PublicKey, PublicTransaction, V03State, + error::{InvalidProgramBehaviorError, LeeError}, + execute_and_prove, + privacy_preserving_transaction::{ + PrivacyPreservingTransaction, circuit::ProgramWithDependencies, message::Message, + witness_set::WitnessSet, + }, + program::Program, + public_transaction, + signature::PrivateKey, +}; + +mod authenticated_transfer; +mod changer_claimer; +mod circuit; +mod claiming; +mod flash_swap; +mod genesis; +mod privacy_preserving; +mod public_program_rules; +mod validity_window; + +impl V03State { + /// Include test programs in the builtin programs map. + #[must_use] + pub fn with_test_programs(mut self) -> Self { + self.insert_program(crate::test_methods::simple_balance_transfer()); + self.insert_program(crate::test_methods::nonce_changer()); + self.insert_program(crate::test_methods::extra_output()); + self.insert_program(crate::test_methods::missing_output()); + self.insert_program(crate::test_methods::program_owner_changer()); + self.insert_program(crate::test_methods::data_changer()); + self.insert_program(crate::test_methods::minter()); + self.insert_program(crate::test_methods::burner()); + self.insert_program(crate::test_methods::auth_asserting_noop()); + self.insert_program(crate::test_methods::private_pda_delegator()); + self.insert_program(crate::test_methods::pda_claimer()); + self.insert_program(crate::test_methods::two_pda_claimer()); + self.insert_program(crate::test_methods::noop()); + self.insert_program(crate::test_methods::chain_caller()); + self.insert_program(crate::test_methods::modified_transfer_program()); + self.insert_program(crate::test_methods::malicious_authorization_changer()); + self.insert_program(crate::test_methods::validity_window()); + self.insert_program(crate::test_methods::flash_swap_initiator()); + self.insert_program(crate::test_methods::flash_swap_callback()); + self.insert_program(crate::test_methods::malicious_self_program_id()); + self.insert_program(crate::test_methods::malicious_caller_program_id()); + self.insert_program(crate::test_methods::pda_spend_proxy()); + self.insert_program(crate::test_methods::claimer()); + self.insert_program(crate::test_methods::changer_claimer()); + self.insert_program(crate::test_methods::validity_window_chain_caller()); + self.insert_program(crate::test_methods::simple_transfer_proxy()); + self.insert_program(crate::test_methods::malicious_injector()); + self.insert_program(crate::test_methods::malicious_launderer()); + self.insert_program(crate::test_methods::modified_transfer_program()); + self + } + + #[must_use] + pub fn with_non_default_accounts_but_default_program_owners(mut self) -> Self { + let account_with_default_values_except_balance = Account { + balance: 100, + ..Account::default() + }; + let account_with_default_values_except_nonce = Account { + nonce: Nonce(37), + ..Account::default() + }; + let account_with_default_values_except_data = Account { + data: vec![0xca, 0xfe].try_into().unwrap(), + ..Account::default() + }; + self.force_insert_account( + AccountId::new([255; 32]), + account_with_default_values_except_balance, + ); + self.force_insert_account( + AccountId::new([254; 32]), + account_with_default_values_except_nonce, + ); + self.force_insert_account( + AccountId::new([253; 32]), + account_with_default_values_except_data, + ); + self + } + + #[must_use] + pub fn with_account_owned_by_burner_program(mut self) -> Self { + let account = Account { + program_owner: crate::test_methods::burner().id(), + balance: 100, + ..Default::default() + }; + self.force_insert_account(AccountId::new([252; 32]), account); + self + } + + #[must_use] + pub fn with_private_account(mut self, keys: &TestPrivateKeys, account: &Account) -> Self { + let account_id = AccountId::for_regular_private_account(&keys.npk(), 0); + let commitment = Commitment::new(&account_id, account); + self.private_state.0.extend(&[commitment]); + self + } +} + +pub struct TestPublicKeys { + pub signing_key: PrivateKey, +} + +impl TestPublicKeys { + pub fn account_id(&self) -> AccountId { + AccountId::from(&PublicKey::new_from_private_key(&self.signing_key)) + } +} + +pub struct TestPrivateKeys { + pub nsk: NullifierSecretKey, + pub d: [u8; 32], + pub z: [u8; 32], +} + +impl TestPrivateKeys { + pub fn npk(&self) -> NullifierPublicKey { + NullifierPublicKey::from(&self.nsk) + } + + pub fn vpk(&self) -> ViewingPublicKey { + ViewingPublicKey::from_seed(&self.d, &self.z) + } +} + +// ── Flash Swap types (mirrors of guest types for host-side serialisation) ── + +#[derive(serde::Serialize, serde::Deserialize)] +struct CallbackInstruction { + return_funds: bool, + token_program_id: ProgramId, + amount: u128, +} + +#[derive(serde::Serialize, serde::Deserialize)] +enum FlashSwapInstruction { + Initiate { + token_program_id: ProgramId, + callback_program_id: ProgramId, + amount_out: u128, + callback_instruction_data: Vec, + }, + InvariantCheck { + min_vault_balance: u128, + }, +} + +fn public_state_from_balances(initial_data: &[(AccountId, u128)]) -> HashMap { + initial_data + .iter() + .copied() + .map(|(account_id, balance)| { + ( + account_id, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance, + ..Account::default() + }, + ) + }) + .collect() +} + +fn transfer_transaction( + from: AccountId, + from_key: &PrivateKey, + from_nonce: u128, + to: AccountId, + to_key: &PrivateKey, + to_nonce: u128, + balance: u128, +) -> PublicTransaction { + let account_ids = vec![from, to]; + let nonces = vec![Nonce(from_nonce), Nonce(to_nonce)]; + let program_id = crate::test_methods::simple_balance_transfer().id(); + let message = + public_transaction::Message::try_new(program_id, account_ids, nonces, balance).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[from_key, to_key]); + PublicTransaction::new(message, witness_set) +} + +fn build_flash_swap_tx( + initiator: &Program, + vault_id: AccountId, + receiver_id: AccountId, + instruction: FlashSwapInstruction, +) -> PublicTransaction { + let message = public_transaction::Message::try_new( + initiator.id(), + vec![vault_id, receiver_id], + vec![], // no signers — vault is PDA-authorised + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + PublicTransaction::new(message, witness_set) +} + +fn test_public_account_keys_1() -> TestPublicKeys { + TestPublicKeys { + signing_key: PrivateKey::try_new([37; 32]).unwrap(), + } +} + +fn test_public_account_keys_2() -> TestPublicKeys { + TestPublicKeys { + signing_key: PrivateKey::try_new([38; 32]).unwrap(), + } +} + +pub fn test_private_account_keys_1() -> TestPrivateKeys { + TestPrivateKeys { + nsk: [13; 32], + d: [31; 32], + z: [32; 32], + } +} + +pub fn test_private_account_keys_2() -> TestPrivateKeys { + TestPrivateKeys { + nsk: [38; 32], + d: [83; 32], + z: [84; 32], + } +} + +fn shielded_balance_transfer_for_tests( + sender_keys: &TestPublicKeys, + recipient_keys: &TestPrivateKeys, + balance_to_move: u128, + state: &V03State, +) -> PrivacyPreservingTransaction { + let sender = AccountWithMetadata::new( + state.get_account_by_id(sender_keys.account_id()), + true, + sender_keys.account_id(), + ); + + let sender_nonce = sender.account.nonce; + + let recipient = AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0)); + + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0); + + let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove( + vec![sender, recipient], + Program::serialize_instruction(balance_to_move).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivateUnauthorized { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: shared_secret, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + ], + &crate::test_methods::simple_balance_transfer().into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output( + vec![sender_keys.account_id()], + vec![sender_nonce], + output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[&sender_keys.signing_key]); + PrivacyPreservingTransaction::new(message, witness_set) +} + +fn private_balance_transfer_for_tests( + sender_keys: &TestPrivateKeys, + sender_private_account: &Account, + recipient_keys: &TestPrivateKeys, + balance_to_move: u128, + state: &V03State, +) -> PrivacyPreservingTransaction { + let program = crate::test_methods::simple_balance_transfer(); + let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); + let sender_commitment = Commitment::new(&sender_account_id, sender_private_account); + let sender_pre = AccountWithMetadata::new( + sender_private_account.clone(), + true, + (&sender_keys.npk(), 0), + ); + let recipient_pre = + AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0)); + + let (shared_secret_1, epk_1) = + SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0); + + let (shared_secret_2, epk_2) = + SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1); + + let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove( + vec![sender_pre, recipient_pre], + Program::serialize_instruction(balance_to_move).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: epk_1, + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: shared_secret_1, + nsk: sender_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&sender_commitment) + .expect("sender's commitment must be in state"), + identifier: 0, + }, + InputAccountIdentity::PrivateUnauthorized { + epk: epk_2, + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + npk: recipient_keys.npk(), + ssk: shared_secret_2, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }, + ], + &program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + + PrivacyPreservingTransaction::new(message, witness_set) +} + +fn deshielded_balance_transfer_for_tests( + sender_keys: &TestPrivateKeys, + sender_private_account: &Account, + recipient_account_id: &AccountId, + balance_to_move: u128, + state: &V03State, +) -> PrivacyPreservingTransaction { + let program = crate::test_methods::simple_balance_transfer(); + let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); + let sender_commitment = Commitment::new(&sender_account_id, sender_private_account); + let sender_pre = AccountWithMetadata::new( + sender_private_account.clone(), + true, + (&sender_keys.npk(), 0), + ); + let recipient_pre = AccountWithMetadata::new( + state.get_account_by_id(*recipient_account_id), + false, + *recipient_account_id, + ); + + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0); + + let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove( + vec![sender_pre, recipient_pre], + Program::serialize_instruction(balance_to_move).unwrap(), + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.npk(), + &sender_keys.vpk(), + ), + ssk: shared_secret, + nsk: sender_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&sender_commitment) + .expect("sender's commitment must be in state"), + identifier: 0, + }, + InputAccountIdentity::Public, + ], + &program.into(), + ) + .unwrap(); + + let message = + Message::try_from_circuit_output(vec![*recipient_account_id], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + + PrivacyPreservingTransaction::new(message, witness_set) +} + +fn valid_private_transfer_tx_and_state() -> (V03State, PrivacyPreservingTransaction) { + let sender_keys = test_private_account_keys_1(); + let sender_private_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + nonce: Nonce(0xdead_beef), + ..Account::default() + }; + let recipient_keys = test_private_account_keys_2(); + let state = V03State::new().with_private_account(&sender_keys, &sender_private_account); + let tx = private_balance_transfer_for_tests( + &sender_keys, + &sender_private_account, + &recipient_keys, + 37, + &state, + ); + (state, tx) +} diff --git a/lee/state_machine/src/state/tests/privacy_preserving.rs b/lee/state_machine/src/state/tests/privacy_preserving.rs new file mode 100644 index 00000000..0205b5fd --- /dev/null +++ b/lee/state_machine/src/state/tests/privacy_preserving.rs @@ -0,0 +1,539 @@ +use super::*; + +#[test] +fn transition_from_privacy_preserving_transaction_shielded() { + let sender_keys = test_public_account_keys_1(); + let recipient_keys = test_private_account_keys_1(); + + let mut state = V03State::new().with_public_accounts([( + sender_keys.account_id(), + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 200, + ..Account::default() + }, + )]); + + let balance_to_move = 37; + + let tx = + shielded_balance_transfer_for_tests(&sender_keys, &recipient_keys, balance_to_move, &state); + + let expected_sender_post = { + let mut this = state.get_account_by_id(sender_keys.account_id()); + this.balance -= balance_to_move; + this.nonce.public_account_nonce_increment(); + this + }; + + let [expected_new_commitment] = tx.message().new_commitments.clone().try_into().unwrap(); + assert!(!state.private_state.0.contains(&expected_new_commitment)); + + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .unwrap(); + + let sender_post = state.get_account_by_id(sender_keys.account_id()); + assert_eq!(sender_post, expected_sender_post); + assert!(state.private_state.0.contains(&expected_new_commitment)); + + assert_eq!( + state.get_account_by_id(sender_keys.account_id()).balance, + 200 - balance_to_move + ); +} + +#[test] +fn transition_from_privacy_preserving_transaction_private() { + let sender_keys = test_private_account_keys_1(); + let sender_nonce = Nonce(0xdead_beef); + + let sender_private_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + nonce: sender_nonce, + data: Data::default(), + }; + let recipient_keys = test_private_account_keys_2(); + + let mut state = V03State::new().with_private_account(&sender_keys, &sender_private_account); + + let balance_to_move = 37; + + let tx = private_balance_transfer_for_tests( + &sender_keys, + &sender_private_account, + &recipient_keys, + balance_to_move, + &state, + ); + + let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); + let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); + let expected_new_commitment_1 = Commitment::new( + &sender_account_id, + &Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk), + balance: sender_private_account.balance - balance_to_move, + data: Data::default(), + }, + ); + + let sender_pre_commitment = Commitment::new(&sender_account_id, &sender_private_account); + let expected_new_nullifier = + Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk); + + let expected_new_commitment_2 = Commitment::new( + &recipient_account_id, + &Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + nonce: Nonce::private_account_nonce_init(&recipient_account_id), + balance: balance_to_move, + ..Account::default() + }, + ); + + let previous_public_state = state.public_state.clone(); + assert!(state.private_state.0.contains(&sender_pre_commitment)); + assert!(!state.private_state.0.contains(&expected_new_commitment_1)); + assert!(!state.private_state.0.contains(&expected_new_commitment_2)); + assert!(!state.private_state.1.contains(&expected_new_nullifier)); + + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .unwrap(); + + assert_eq!(state.public_state, previous_public_state); + assert!(state.private_state.0.contains(&sender_pre_commitment)); + assert!(state.private_state.0.contains(&expected_new_commitment_1)); + assert!(state.private_state.0.contains(&expected_new_commitment_2)); + assert!(state.private_state.1.contains(&expected_new_nullifier)); +} + +/// After a valid fully-private tx is proven, tampering with a note's epk should +/// make the shielding proof invalid. +#[test] +fn privacy_tampered_epk_is_rejected() { + use crate::validated_state_diff::ValidatedStateDiff; + + let (state, mut tx) = valid_private_transfer_tx_and_state(); + + // Baseline: the untampered tx verifies + assert!( + ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0).is_ok(), + "the unmodified private transfer must verify" + ); + + // Flip a byte of the first note's epk + tx.message.encrypted_private_post_states[0].epk.0[0] ^= 0xFF; + + assert!( + matches!( + ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0), + Err(LeeError::InvalidPrivacyPreservingProof) + ), + "a tampered epk must be rejected by proof verification" + ); +} + +/// After a valid fully-private tx is proven, tampering with a note's view tag should +/// make the shielding proof invalid. +#[test] +fn privacy_tampered_view_tag_is_rejected() { + use crate::validated_state_diff::ValidatedStateDiff; + + let (state, mut tx) = valid_private_transfer_tx_and_state(); + + // Baseline: the untampered tx verifies. + assert!( + ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0).is_ok(), + "the unmodified private transfer must verify" + ); + + // Flip the first note's view_tag + tx.message.encrypted_private_post_states[0].view_tag ^= 0xFF; + + assert!( + matches!( + ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0), + Err(LeeError::InvalidPrivacyPreservingProof) + ), + "a tampered view_tag must be rejected by proof verification" + ); +} + +#[test] +fn transition_from_privacy_preserving_transaction_deshielded() { + let sender_keys = test_private_account_keys_1(); + let sender_nonce = Nonce(0xdead_beef); + + let sender_private_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + nonce: sender_nonce, + data: Data::default(), + }; + let recipient_keys = test_public_account_keys_1(); + let recipient_initial_balance = 400; + let mut state = V03State::new() + .with_public_accounts([( + recipient_keys.account_id(), + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: recipient_initial_balance, + ..Account::default() + }, + )]) + .with_private_account(&sender_keys, &sender_private_account); + + let balance_to_move = 37; + + let expected_recipient_post = { + let mut this = state.get_account_by_id(recipient_keys.account_id()); + this.balance += balance_to_move; + this + }; + + let tx = deshielded_balance_transfer_for_tests( + &sender_keys, + &sender_private_account, + &recipient_keys.account_id(), + balance_to_move, + &state, + ); + + let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0); + let expected_new_commitment = Commitment::new( + &sender_account_id, + &Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk), + balance: sender_private_account.balance - balance_to_move, + data: Data::default(), + }, + ); + + let sender_pre_commitment = Commitment::new(&sender_account_id, &sender_private_account); + let expected_new_nullifier = + Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk); + + assert!(state.private_state.0.contains(&sender_pre_commitment)); + assert!(!state.private_state.0.contains(&expected_new_commitment)); + assert!(!state.private_state.1.contains(&expected_new_nullifier)); + + state + .transition_from_privacy_preserving_transaction(&tx, 1, 0) + .unwrap(); + + let recipient_post = state.get_account_by_id(recipient_keys.account_id()); + assert_eq!(recipient_post, expected_recipient_post); + assert!(state.private_state.0.contains(&sender_pre_commitment)); + assert!(state.private_state.0.contains(&expected_new_commitment)); + assert!(state.private_state.1.contains(&expected_new_nullifier)); + assert_eq!( + state.get_account_by_id(recipient_keys.account_id()).balance, + recipient_initial_balance + balance_to_move + ); +} + +#[test] +fn burner_program_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::burner(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 100, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(10_u128).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn minter_program_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::minter(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(10_u128).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn nonce_changer_program_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::nonce_changer(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn data_changer_program_should_fail_for_non_owned_account_in_privacy_preserving_circuit() { + let program = crate::test_methods::data_changer(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: [0, 1, 2, 3, 4, 5, 6, 7], + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(vec![0]).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn data_changer_program_should_fail_for_too_large_data_in_privacy_preserving_circuit() { + let program = crate::test_methods::data_changer(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let large_data: Vec = + vec![ + 0; + usize::try_from(lee_core::account::data::DATA_MAX_LENGTH.as_u64()) + .expect("DATA_MAX_LENGTH fits in usize") + + 1 + ]; + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(large_data).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::ProgramProveFailed(_)))); +} + +#[test] +fn extra_output_program_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::extra_output(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn missing_output_program_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::missing_output(); + let public_account_1 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + let public_account_2 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([1; 32]), + ); + + let result = execute_and_prove( + vec![public_account_1, public_account_2], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::Public, InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn program_owner_changer_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::program_owner_changer(); + let public_account = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + + let result = execute_and_prove( + vec![public_account], + Program::serialize_instruction(()).unwrap(), + vec![InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn transfer_from_non_owned_account_should_fail_in_privacy_preserving_circuit() { + let program = crate::test_methods::simple_balance_transfer(); + let public_account_1 = AccountWithMetadata::new( + Account { + program_owner: [0, 1, 2, 3, 4, 5, 6, 7], + balance: 100, + ..Account::default() + }, + true, + AccountId::new([0; 32]), + ); + let public_account_2 = AccountWithMetadata::new( + Account { + program_owner: program.id(), + balance: 0, + ..Account::default() + }, + true, + AccountId::new([1; 32]), + ); + + let result = execute_and_prove( + vec![public_account_1, public_account_2], + Program::serialize_instruction(10_u128).unwrap(), + vec![InputAccountIdentity::Public, InputAccountIdentity::Public], + &program.into(), + ); + + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} + +#[test] +fn malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() { + // Arrange + let malicious_program = crate::test_methods::malicious_authorization_changer(); + let simple_transfers = crate::test_methods::simple_balance_transfer(); + let sender_keys = test_public_account_keys_1(); + let recipient_keys = test_private_account_keys_1(); + + let sender_account = AccountWithMetadata::new( + Account { + program_owner: simple_transfers.id(), + balance: 100, + ..Default::default() + }, + false, + sender_keys.account_id(), + ); + let recipient_account = + AccountWithMetadata::new(Account::default(), true, (&recipient_keys.npk(), 0)); + + let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0); + let recipient_commitment = Commitment::new(&recipient_account_id, &recipient_account.account); + let recipient_init_nullifier = Nullifier::for_account_initialization(&recipient_account_id); + let state = V03State::new() + .with_public_accounts(public_state_from_balances(&[( + sender_account.account_id, + sender_account.account.balance, + )])) + .with_private_accounts([(recipient_commitment.clone(), recipient_init_nullifier)]) + .with_test_programs(); + + let balance_to_transfer = 10_u128; + let instruction = (balance_to_transfer, simple_transfers.id()); + + let recipient = + SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0; + + let mut dependencies = HashMap::new(); + dependencies.insert(simple_transfers.id(), simple_transfers); + let program_with_deps = ProgramWithDependencies::new(malicious_program, dependencies); + + // Act - execute the malicious program - this should fail during proving + let result = execute_and_prove( + vec![sender_account, recipient_account], + Program::serialize_instruction(instruction).unwrap(), + vec![ + InputAccountIdentity::Public, + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(Vec::new()), + view_tag: EncryptedAccountData::compute_view_tag( + &recipient_keys.npk(), + &recipient_keys.vpk(), + ), + ssk: recipient, + nsk: recipient_keys.nsk, + membership_proof: state + .get_proof_for_commitment(&recipient_commitment) + .expect("recipient's commitment must be in state"), + identifier: 0, + }, + ], + &program_with_deps, + ); + + // Assert - should fail because the malicious program tries to manipulate is_authorized + assert!(matches!(result, Err(LeeError::CircuitProvingError(_)))); +} diff --git a/lee/state_machine/src/state/tests/public_program_rules.rs b/lee/state_machine/src/state/tests/public_program_rules.rs new file mode 100644 index 00000000..1c67aeff --- /dev/null +++ b/lee/state_machine/src/state/tests/public_program_rules.rs @@ -0,0 +1,325 @@ +use super::*; + +#[test] +fn program_should_fail_if_modifies_nonces() { + let account_id = AccountId::new([1; 32]); + let mut state = V03State::new() + .with_public_account_balances([(account_id, 100)]) + .with_test_programs(); + let account_ids = vec![account_id]; + let program_id = crate::test_methods::nonce_changer().id(); + let message = + public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::ModifiedNonce { account_id: err_account_id } + ) + )) if err_account_id == account_id + )); +} + +#[test] +fn program_should_fail_if_output_accounts_exceed_inputs() { + let mut state = V03State::new() + .with_public_account_balances([(AccountId::new([1; 32]), 0)]) + .with_test_programs(); + let account_ids = vec![AccountId::new([1; 32])]; + let program_id = crate::test_methods::extra_output().id(); + let message = + public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::MismatchedPreStatePostStateLength { + pre_state_length, + post_state_length + } + ) + )) if pre_state_length == 1 && post_state_length == 2 + )); +} + +#[test] +fn program_should_fail_with_missing_output_accounts() { + let mut state = V03State::new() + .with_public_account_balances([(AccountId::new([1; 32]), 100)]) + .with_test_programs(); + let account_ids = vec![AccountId::new([1; 32]), AccountId::new([2; 32])]; + let program_id = crate::test_methods::missing_output().id(); + let message = + public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::MismatchedPreStatePostStateLength { + pre_state_length, + post_state_length + } + ) + )) if pre_state_length == 2 && post_state_length == 1 + )); +} + +#[test] +fn program_should_fail_if_modifies_program_owner_with_only_non_default_program_owner() { + let initial_data = [( + AccountId::new([1; 32]), + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + ..Account::default() + }, + )]; + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs(); + let account_id = AccountId::new([1; 32]); + let account = state.get_account_by_id(account_id); + // Assert the target account only differs from the default account in the program owner + // field + assert_ne!(account.program_owner, Account::default().program_owner); + assert_eq!(account.balance, Account::default().balance); + assert_eq!(account.nonce, Account::default().nonce); + assert_eq!(account.data, Account::default().data); + let program_id = crate::test_methods::program_owner_changer().id(); + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } + ))) if err_account_id == account_id + )); +} + +#[test] +fn program_should_fail_if_modifies_program_owner_with_only_non_default_balance() { + let initial_data = HashMap::new(); + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs() + .with_non_default_accounts_but_default_program_owners(); + let account_id = AccountId::new([255; 32]); + let account = state.get_account_by_id(account_id); + // Assert the target account only differs from the default account in balance field + assert_eq!(account.program_owner, Account::default().program_owner); + assert_ne!(account.balance, Account::default().balance); + assert_eq!(account.nonce, Account::default().nonce); + assert_eq!(account.data, Account::default().data); + let program_id = crate::test_methods::program_owner_changer().id(); + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } + ))) if err_account_id == account_id + )); +} + +#[test] +fn program_should_fail_if_modifies_program_owner_with_only_non_default_nonce() { + let initial_data = HashMap::new(); + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs() + .with_non_default_accounts_but_default_program_owners(); + let account_id = AccountId::new([254; 32]); + let account = state.get_account_by_id(account_id); + // Assert the target account only differs from the default account in nonce field + assert_eq!(account.program_owner, Account::default().program_owner); + assert_eq!(account.balance, Account::default().balance); + assert_ne!(account.nonce, Account::default().nonce); + assert_eq!(account.data, Account::default().data); + let program_id = crate::test_methods::program_owner_changer().id(); + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } + ))) if err_account_id == account_id + )); +} + +#[test] +fn program_should_fail_if_modifies_program_owner_with_only_non_default_data() { + let initial_data = HashMap::new(); + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs() + .with_non_default_accounts_but_default_program_owners(); + let account_id = AccountId::new([253; 32]); + let account = state.get_account_by_id(account_id); + // Assert the target account only differs from the default account in data field + assert_eq!(account.program_owner, Account::default().program_owner); + assert_eq!(account.balance, Account::default().balance); + assert_eq!(account.nonce, Account::default().nonce); + assert_ne!(account.data, Account::default().data); + let program_id = crate::test_methods::program_owner_changer().id(); + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id } + ))) if err_account_id == account_id + )); +} + +#[test] +fn program_should_fail_if_transfers_balance_from_non_owned_account() { + let sender_account_id = AccountId::new([1; 32]); + let receiver_account_id = AccountId::new([2; 32]); + let mut state = V03State::new() + .with_public_account_balances([(sender_account_id, 100)]) + .with_test_programs(); + let balance_to_move: u128 = 1; + let program_id = crate::test_methods::simple_balance_transfer().id(); + assert_ne!( + state.get_account_by_id(sender_account_id).program_owner, + program_id + ); + let message = public_transaction::Message::try_new( + program_id, + vec![sender_account_id, receiver_account_id], + vec![], + balance_to_move, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::UnauthorizedBalanceDecrease { account_id: err_account_id, owner_program_id, executing_program_id } + ))) if err_account_id == sender_account_id && owner_program_id != program_id && executing_program_id == program_id + )); +} + +#[test] +fn program_should_fail_if_modifies_data_of_non_owned_account() { + let initial_data = HashMap::new(); + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs() + .with_non_default_accounts_but_default_program_owners(); + let account_id = AccountId::new([255; 32]); + let program_id = crate::test_methods::data_changer().id(); + + assert_ne!(state.get_account_by_id(account_id), Account::default()); + assert_ne!( + state.get_account_by_id(account_id).program_owner, + program_id + ); + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], vec![0]) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::UnauthorizedDataModification { account_id: err_account_id, executing_program_id } + ))) if err_account_id == account_id && executing_program_id == program_id + )); +} + +#[test] +fn program_should_fail_if_does_not_preserve_total_balance_by_minting() { + let initial_data = HashMap::new(); + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs(); + let account_id = AccountId::new([1; 32]); + let program_id = crate::test_methods::minter().id(); + + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 2, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states } + ))) if total_balance_pre_states == 0.into() && total_balance_post_states == 1.into() + )); +} + +#[test] +fn program_should_fail_if_does_not_preserve_total_balance_by_burning() { + let initial_data = HashMap::new(); + let mut state = V03State::new() + .with_public_accounts(initial_data) + .with_test_programs() + .with_account_owned_by_burner_program(); + let program_id = crate::test_methods::burner().id(); + let account_id = AccountId::new([252; 32]); + assert_eq!( + state.get_account_by_id(account_id).program_owner, + program_id + ); + let balance_to_burn: u128 = 1; + assert!(state.get_account_by_id(account_id).balance > balance_to_burn); + + let message = + public_transaction::Message::try_new(program_id, vec![account_id], vec![], balance_to_burn) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + let result = state.transition_from_public_transaction(&tx, 2, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed( + ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states } + ))) if total_balance_pre_states == 100.into() && total_balance_post_states == 99.into() + )); +} diff --git a/lee/state_machine/src/state/tests/validity_window.rs b/lee/state_machine/src/state/tests/validity_window.rs new file mode 100644 index 00000000..5a46bd5b --- /dev/null +++ b/lee/state_machine/src/state/tests/validity_window.rs @@ -0,0 +1,243 @@ +use super::*; + +#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] +#[test_case::test_case((Some(1), Some(3)), 2; "inside range")] +#[test_case::test_case((Some(1), Some(3)), 0; "below range")] +#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] +#[test_case::test_case((Some(1), Some(3)), 4; "above range")] +#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] +#[test_case::test_case((Some(1), None), 10; "lower bound only - above")] +#[test_case::test_case((Some(1), None), 0; "lower bound only - below")] +#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] +#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] +#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] +#[test_case::test_case((None, None), 0; "no bounds - always valid")] +#[test_case::test_case((None, None), 100; "no bounds - always valid 2")] +fn validity_window_works_in_public_transactions( + validity_window: (Option, Option), + block_id: BlockId, +) { + let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap(); + let validity_window_program = crate::test_methods::validity_window(); + let account_keys = test_public_account_keys_1(); + let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id()); + let mut state = V03State::new().with_test_programs(); + let tx = { + let account_ids = vec![pre.account_id]; + let nonces = vec![]; + let program_id = validity_window_program.id(); + let instruction = ( + block_validity_window, + TimestampValidityWindow::new_unbounded(), + ); + let message = + public_transaction::Message::try_new(program_id, account_ids, nonces, instruction) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + PublicTransaction::new(message, witness_set) + }; + let result = state.transition_from_public_transaction(&tx, block_id, 0); + let is_inside_validity_window = + match (block_validity_window.start(), block_validity_window.end()) { + (Some(s), Some(e)) => s <= block_id && block_id < e, + (Some(s), None) => s <= block_id, + (None, Some(e)) => block_id < e, + (None, None) => true, + }; + if is_inside_validity_window { + assert!(result.is_ok()); + } else { + assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); + } +} + +#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] +#[test_case::test_case((Some(1), Some(3)), 2; "inside range")] +#[test_case::test_case((Some(1), Some(3)), 0; "below range")] +#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] +#[test_case::test_case((Some(1), Some(3)), 4; "above range")] +#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] +#[test_case::test_case((Some(1), None), 10; "lower bound only - above")] +#[test_case::test_case((Some(1), None), 0; "lower bound only - below")] +#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] +#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] +#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] +#[test_case::test_case((None, None), 0; "no bounds - always valid")] +#[test_case::test_case((None, None), 100; "no bounds - always valid 2")] +fn timestamp_validity_window_works_in_public_transactions( + validity_window: (Option, Option), + timestamp: Timestamp, +) { + let timestamp_validity_window: TimestampValidityWindow = validity_window.try_into().unwrap(); + let validity_window_program = crate::test_methods::validity_window(); + let account_keys = test_public_account_keys_1(); + let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id()); + let mut state = V03State::new().with_test_programs(); + let tx = { + let account_ids = vec![pre.account_id]; + let nonces = vec![]; + let program_id = validity_window_program.id(); + let instruction = ( + BlockValidityWindow::new_unbounded(), + timestamp_validity_window, + ); + let message = + public_transaction::Message::try_new(program_id, account_ids, nonces, instruction) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + PublicTransaction::new(message, witness_set) + }; + let result = state.transition_from_public_transaction(&tx, 1, timestamp); + let is_inside_validity_window = match ( + timestamp_validity_window.start(), + timestamp_validity_window.end(), + ) { + (Some(s), Some(e)) => s <= timestamp && timestamp < e, + (Some(s), None) => s <= timestamp, + (None, Some(e)) => timestamp < e, + (None, None) => true, + }; + if is_inside_validity_window { + assert!(result.is_ok()); + } else { + assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); + } +} + +#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] +#[test_case::test_case((Some(1), Some(3)), 2; "inside range")] +#[test_case::test_case((Some(1), Some(3)), 0; "below range")] +#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] +#[test_case::test_case((Some(1), Some(3)), 4; "above range")] +#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] +#[test_case::test_case((Some(1), None), 10; "lower bound only - above")] +#[test_case::test_case((Some(1), None), 0; "lower bound only - below")] +#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] +#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] +#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] +#[test_case::test_case((None, None), 0; "no bounds - always valid")] +#[test_case::test_case((None, None), 100; "no bounds - always valid 2")] +fn validity_window_works_in_privacy_preserving_transactions( + validity_window: (Option, Option), + block_id: BlockId, +) { + let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap(); + let validity_window_program = crate::test_methods::validity_window(); + let account_keys = test_private_account_keys_1(); + let pre = AccountWithMetadata::new(Account::default(), false, (&account_keys.npk(), 0)); + let mut state = V03State::new().with_test_programs(); + let tx = { + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0); + + let instruction = ( + block_validity_window, + TimestampValidityWindow::new_unbounded(), + ); + let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove( + vec![pre], + Program::serialize_instruction(instruction).unwrap(), + vec![InputAccountIdentity::PrivateUnauthorized { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &account_keys.npk(), + &account_keys.vpk(), + ), + npk: account_keys.npk(), + ssk: shared_secret, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &validity_window_program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + PrivacyPreservingTransaction::new(message, witness_set) + }; + let result = state.transition_from_privacy_preserving_transaction(&tx, block_id, 0); + let is_inside_validity_window = + match (block_validity_window.start(), block_validity_window.end()) { + (Some(s), Some(e)) => s <= block_id && block_id < e, + (Some(s), None) => s <= block_id, + (None, Some(e)) => block_id < e, + (None, None) => true, + }; + if is_inside_validity_window { + assert!(result.is_ok()); + } else { + assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); + } +} + +#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")] +#[test_case::test_case((Some(1), Some(3)), 2; "inside range")] +#[test_case::test_case((Some(1), Some(3)), 0; "below range")] +#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")] +#[test_case::test_case((Some(1), Some(3)), 4; "above range")] +#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")] +#[test_case::test_case((Some(1), None), 10; "lower bound only - above")] +#[test_case::test_case((Some(1), None), 0; "lower bound only - below")] +#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")] +#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")] +#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")] +#[test_case::test_case((None, None), 0; "no bounds - always valid")] +#[test_case::test_case((None, None), 100; "no bounds - always valid 2")] +fn timestamp_validity_window_works_in_privacy_preserving_transactions( + validity_window: (Option, Option), + timestamp: Timestamp, +) { + let timestamp_validity_window: TimestampValidityWindow = validity_window.try_into().unwrap(); + let validity_window_program = crate::test_methods::validity_window(); + let account_keys = test_private_account_keys_1(); + let pre = AccountWithMetadata::new(Account::default(), false, (&account_keys.npk(), 0)); + let mut state = V03State::new().with_test_programs(); + let tx = { + let (shared_secret, epk) = + SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0); + + let instruction = ( + BlockValidityWindow::new_unbounded(), + timestamp_validity_window, + ); + let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove( + vec![pre], + Program::serialize_instruction(instruction).unwrap(), + vec![InputAccountIdentity::PrivateUnauthorized { + epk, + view_tag: EncryptedAccountData::compute_view_tag( + &account_keys.npk(), + &account_keys.vpk(), + ), + npk: account_keys.npk(), + ssk: shared_secret, + identifier: 0, + commitment_root: DUMMY_COMMITMENT_HASH, + }], + &validity_window_program.into(), + ) + .unwrap(); + + let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); + PrivacyPreservingTransaction::new(message, witness_set) + }; + let result = state.transition_from_privacy_preserving_transaction(&tx, 1, timestamp); + let is_inside_validity_window = match ( + timestamp_validity_window.start(), + timestamp_validity_window.end(), + ) { + (Some(s), Some(e)) => s <= timestamp && timestamp < e, + (Some(s), None) => s <= timestamp, + (None, Some(e)) => timestamp < e, + (None, None) => true, + }; + if is_inside_validity_window { + assert!(result.is_ok()); + } else { + assert!(matches!(result, Err(LeeError::OutOfValidityWindow))); + } +} diff --git a/lee/state_machine/src/validated_state_diff.rs b/lee/state_machine/src/validated_state_diff.rs deleted file mode 100644 index 44a307af..00000000 --- a/lee/state_machine/src/validated_state_diff.rs +++ /dev/null @@ -1,1055 +0,0 @@ -use std::{ - collections::{HashMap, HashSet, VecDeque}, - hash::Hash, -}; - -use lee_core::{ - BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, Timestamp, - account::{Account, AccountId, AccountWithMetadata}, - program::{ - ChainedCall, Claim, DEFAULT_PROGRAM_ID, ProgramId, compute_public_authorized_pdas, - validate_execution, - }, -}; -use log::debug; - -use crate::{ - V03State, ensure, - error::{InvalidProgramBehaviorError, LeeError}, - privacy_preserving_transaction::{ - PrivacyPreservingTransaction, circuit::Proof, message::Message, - }, - program::Program, - program_deployment_transaction::ProgramDeploymentTransaction, - public_transaction::PublicTransaction, - state::MAX_NUMBER_CHAINED_CALLS, -}; - -pub struct StateDiff { - pub signer_account_ids: Vec, - pub public_diff: HashMap, - pub new_commitments: Vec, - pub new_nullifiers: Vec, - pub program: Option, -} - -/// The validated output of executing or verifying a transaction, ready to be applied to the state. -/// -/// Can only be constructed by the transaction validation functions inside this crate, ensuring the -/// diff has been checked before any state mutation occurs. -pub struct ValidatedStateDiff(StateDiff); - -impl ValidatedStateDiff { - pub fn from_public_transaction( - tx: &PublicTransaction, - state: &V03State, - block_id: BlockId, - timestamp: Timestamp, - ) -> Result { - let message = tx.message(); - let witness_set = tx.witness_set(); - - ensure!( - !message.account_ids.is_empty(), - LeeError::InvalidInput("Public transaction must have at least one account".into()) - ); - - // All account_ids must be different - ensure!( - message.account_ids.iter().collect::>().len() == message.account_ids.len(), - LeeError::InvalidInput("Duplicate account_ids found in message".into(),) - ); - - // Check exactly one nonce is provided for each signature - ensure!( - message.nonces.len() == witness_set.signatures_and_public_keys.len(), - LeeError::InvalidInput( - "Mismatch between number of nonces and signatures/public keys".into(), - ) - ); - - // Check the signatures are valid - ensure!( - witness_set.is_valid_for(message), - LeeError::InvalidInput("Invalid signature for given message and public key".into()) - ); - - let signer_account_ids = tx.signer_account_ids(); - // Check nonces corresponds to the current nonces on the public state. - for (account_id, nonce) in signer_account_ids.iter().zip(&message.nonces) { - let current_nonce = state.get_account_by_id(*account_id).nonce; - ensure!( - current_nonce == *nonce, - LeeError::InvalidInput("Nonce mismatch".into()) - ); - } - - // Build pre_states for execution - let input_pre_states: Vec<_> = message - .account_ids - .iter() - .map(|account_id| { - AccountWithMetadata::new( - state.get_account_by_id(*account_id), - signer_account_ids.contains(account_id), - *account_id, - ) - }) - .collect(); - - let mut state_diff: HashMap = HashMap::new(); - - let initial_call = ChainedCall { - program_id: message.program_id, - instruction_data: message.instruction_data.clone(), - pre_states: input_pre_states, - pda_seeds: vec![], - }; - - #[expect( - clippy::items_after_statements, - reason = "More readable to keep it behind the place where it's used" - )] - #[derive(Debug)] - struct CallerData { - program_id: Option, - authorized_accounts: HashSet, - } - - let initial_caller_data = CallerData { - program_id: None, - authorized_accounts: signer_account_ids.iter().copied().collect(), - }; - - let mut chained_calls = - VecDeque::<(ChainedCall, CallerData)>::from_iter([(initial_call, initial_caller_data)]); - let mut chain_calls_counter = 0; - - while let Some((chained_call, caller_data)) = chained_calls.pop_front() { - ensure!( - chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS, - LeeError::MaxChainedCallsDepthExceeded - ); - - // Check that the `program_id` corresponds to a deployed program - let Some(program) = state.programs().get(&chained_call.program_id) else { - return Err(LeeError::InvalidInput("Unknown program".into())); - }; - - debug!( - "Program {:?} pre_states: {:?}, instruction_data: {:?}", - chained_call.program_id, chained_call.pre_states, chained_call.instruction_data - ); - let mut program_output = program.execute( - caller_data.program_id, - &chained_call.pre_states, - &chained_call.instruction_data, - )?; - debug!( - "Program {:?} output: {:?}", - chained_call.program_id, program_output - ); - - let authorized_pdas = - compute_public_authorized_pdas(caller_data.program_id, &chained_call.pda_seeds); - - // Account is authorized if it is either in the caller's authorized accounts or in the - // list of PDAs the caller has authorized. - let is_authorized = |account_id: &AccountId| { - authorized_pdas.contains(account_id) - || caller_data.authorized_accounts.contains(account_id) - }; - - for pre in &program_output.pre_states { - let account_id = pre.account_id; - // Check that the program output pre_states coincide with the values in the public - // state or with any modifications to those values during the chain of calls. - let expected_pre = state_diff - .get(&account_id) - .cloned() - .unwrap_or_else(|| state.get_account_by_id(account_id)); - ensure!( - pre.account == expected_pre, - InvalidProgramBehaviorError::InconsistentAccountPreState { - account_id, - expected: Box::new(expected_pre), - actual: Box::new(pre.account.clone()) - } - ); - - // Check that the program output pre_states marked as authorized are indeed - // authorized, and vice-versa. - let is_indeed_authorized = is_authorized(&account_id); - ensure!( - !pre.is_authorized || is_indeed_authorized, - InvalidProgramBehaviorError::InvalidAccountAuthorization { account_id } - ); - ensure!( - pre.is_authorized || !is_indeed_authorized, - InvalidProgramBehaviorError::AuthorizedAccountMarkedAsNotAuthorized { - account_id - } - ); - } - - // Verify that the program output's self_program_id matches the expected program ID. - ensure!( - program_output.self_program_id == chained_call.program_id, - InvalidProgramBehaviorError::MismatchedProgramId { - expected: chained_call.program_id, - actual: program_output.self_program_id - } - ); - - // Verify that the program output's caller_program_id matches the actual caller. - ensure!( - program_output.caller_program_id == caller_data.program_id, - InvalidProgramBehaviorError::MismatchedCallerProgramId { - expected: caller_data.program_id, - actual: program_output.caller_program_id, - } - ); - - // Verify execution corresponds to a well-behaved program. - // See the # Programs section for the definition of the `validate_execution` method. - validate_execution( - &program_output.pre_states, - &program_output.post_states, - chained_call.program_id, - ) - .map_err(InvalidProgramBehaviorError::ExecutionValidationFailed)?; - - // Verify validity window - ensure!( - program_output.block_validity_window.is_valid_for(block_id) - && program_output - .timestamp_validity_window - .is_valid_for(timestamp), - LeeError::OutOfValidityWindow - ); - - for (i, post) in program_output.post_states.iter_mut().enumerate() { - let Some(claim) = post.required_claim() else { - continue; - }; - let pre = &program_output.pre_states[i]; - let account_id = pre.account_id; - - // The invoked program can only claim accounts with default program id. - ensure!( - post.account().program_owner == DEFAULT_PROGRAM_ID, - InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id } - ); - - match claim { - Claim::Authorized => { - // The program can only claim accounts that were authorized by the signer. - ensure!( - pre.is_authorized, - InvalidProgramBehaviorError::ClaimedUnauthorizedAccount { account_id } - ); - } - Claim::Pda(seed) => { - // The program can only claim accounts that correspond to the PDAs it is - // authorized to claim. The public-execution path only sees public - // accounts, so the public-PDA derivation is the correct formula here. - let pda = AccountId::for_public_pda(&chained_call.program_id, &seed); - ensure!( - account_id == pda, - InvalidProgramBehaviorError::MismatchedPdaClaim { - expected: pda, - actual: account_id - } - ); - } - } - - post.account_mut().program_owner = chained_call.program_id; - } - - // Update the state diff - for (pre, post) in program_output - .pre_states - .iter() - .zip(program_output.post_states.iter()) - { - state_diff.insert(pre.account_id, post.account().clone()); - } - - // Source from `program_output.pre_states`, not `chained_call.pre_states`: - // the loop above already gates program_output's `is_authorized` via the - // `!pre.is_authorized || is_indeed_authorized` check, while `chained_call. - // pre_states` is caller-controlled and can be forged (audit-issue 91). - // - // Union with the caller's authorized set so that authorization is monotonically - // growing: once an account is authorized at any point in the chain it remains - // authorized for all subsequent calls. - let authorized_accounts: HashSet<_> = caller_data - .authorized_accounts - .into_iter() - .chain( - program_output - .pre_states - .iter() - .filter(|pre| pre.is_authorized) - .map(|pre| pre.account_id), - ) - .collect(); - for new_call in program_output.chained_calls.into_iter().rev() { - chained_calls.push_front(( - new_call, - CallerData { - program_id: Some(chained_call.program_id), - authorized_accounts: authorized_accounts.clone(), - }, - )); - } - - chain_calls_counter = chain_calls_counter - .checked_add(1) - .expect("we check the max depth at the beginning of the loop"); - } - - // Check that all modified uninitialized accounts where claimed - for (account_id, post) in state_diff.iter().filter_map(|(account_id, post)| { - let pre = state.get_account_by_id(*account_id); - if pre.program_owner != DEFAULT_PROGRAM_ID { - return None; - } - if pre == *post { - return None; - } - Some((*account_id, post)) - }) { - ensure!( - post.program_owner != DEFAULT_PROGRAM_ID, - InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim { account_id } - ); - } - - Ok(Self(StateDiff { - signer_account_ids, - public_diff: state_diff, - new_commitments: vec![], - new_nullifiers: vec![], - program: None, - })) - } - - pub fn from_privacy_preserving_transaction( - tx: &PrivacyPreservingTransaction, - state: &V03State, - block_id: BlockId, - timestamp: Timestamp, - ) -> Result { - let message = &tx.message; - let witness_set = &tx.witness_set; - - // 1. Commitments or nullifiers are non empty - ensure!( - !message.new_commitments.is_empty() || !message.new_nullifiers.is_empty(), - LeeError::InvalidInput( - "Empty commitments and empty nullifiers found in message".into(), - ) - ); - - // 2. Check there are no duplicate account_ids in the public_account_ids list. - ensure!( - n_unique(&message.public_account_ids) == message.public_account_ids.len(), - LeeError::InvalidInput("Duplicate account_ids found in message".into()) - ); - - // Check there are no duplicate nullifiers in the new_nullifiers list - ensure!( - n_unique( - &message - .new_nullifiers - .iter() - .map(|(n, _)| n) - .collect::>() - ) == message.new_nullifiers.len(), - LeeError::InvalidInput("Duplicate nullifiers found in message".into()) - ); - - // Check there are no duplicate commitments in the new_commitments list - ensure!( - n_unique(&message.new_commitments) == message.new_commitments.len(), - LeeError::InvalidInput("Duplicate commitments found in message".into()) - ); - - // 3. Nonce checks and Valid signatures - // Check exactly one nonce is provided for each signature - ensure!( - message.nonces.len() == witness_set.signatures_and_public_keys.len(), - LeeError::InvalidInput( - "Mismatch between number of nonces and signatures/public keys".into(), - ) - ); - - // Check the signatures are valid - ensure!( - witness_set.signatures_are_valid_for(message), - LeeError::InvalidInput("Invalid signature for given message and public key".into()) - ); - - let signer_account_ids = tx.signer_account_ids(); - // Check nonces corresponds to the current nonces on the public state. - for (account_id, nonce) in signer_account_ids.iter().zip(&message.nonces) { - let current_nonce = state.get_account_by_id(*account_id).nonce; - ensure!( - current_nonce == *nonce, - LeeError::InvalidInput("Nonce mismatch".into()) - ); - } - - // Verify validity window - ensure!( - message.block_validity_window.is_valid_for(block_id) - && message.timestamp_validity_window.is_valid_for(timestamp), - LeeError::OutOfValidityWindow - ); - - // Build pre_states for proof verification - let public_pre_states: Vec<_> = message - .public_account_ids - .iter() - .map(|account_id| { - AccountWithMetadata::new( - state.get_account_by_id(*account_id), - signer_account_ids.contains(account_id), - *account_id, - ) - }) - .collect(); - - // 4. Proof verification - check_privacy_preserving_circuit_proof_is_valid( - &witness_set.proof, - &public_pre_states, - message, - )?; - - // 5. Commitment freshness - state.check_commitments_are_new(&message.new_commitments)?; - - // 6. Nullifier uniqueness - state.check_nullifiers_are_valid(&message.new_nullifiers)?; - - let public_diff = message - .public_account_ids - .iter() - .copied() - .zip(message.public_post_states.clone()) - .collect(); - let new_nullifiers = message - .new_nullifiers - .iter() - .copied() - .map(|(nullifier, _)| nullifier) - .collect(); - - Ok(Self(StateDiff { - signer_account_ids, - public_diff, - new_commitments: message.new_commitments.clone(), - new_nullifiers, - program: None, - })) - } - - pub fn from_program_deployment_transaction( - tx: &ProgramDeploymentTransaction, - state: &V03State, - ) -> Result { - // TODO: remove clone - let program = Program::new(tx.message.bytecode.clone().into())?; - if state.programs().contains_key(&program.id()) { - return Err(LeeError::ProgramAlreadyExists); - } - Ok(Self(StateDiff { - signer_account_ids: vec![], - public_diff: HashMap::new(), - new_commitments: vec![], - new_nullifiers: vec![], - program: Some(program), - })) - } - - /// Returns the public account changes produced by this transaction. - /// - /// Used by callers (e.g. the sequencer) to inspect the diff before committing it, for example - /// to enforce that system accounts are not modified by user transactions. - #[must_use] - pub fn public_diff(&self) -> HashMap { - self.0.public_diff.clone() - } - - pub(crate) fn into_state_diff(self) -> StateDiff { - self.0 - } -} - -fn check_privacy_preserving_circuit_proof_is_valid( - proof: &Proof, - public_pre_states: &[AccountWithMetadata], - message: &Message, -) -> Result<(), LeeError> { - let output = PrivacyPreservingCircuitOutput { - public_pre_states: public_pre_states.to_vec(), - public_post_states: message.public_post_states.clone(), - encrypted_private_post_states: message.encrypted_private_post_states.clone(), - new_commitments: message.new_commitments.clone(), - new_nullifiers: message.new_nullifiers.clone(), - block_validity_window: message.block_validity_window, - timestamp_validity_window: message.timestamp_validity_window, - }; - proof - .is_valid_for(&output) - .then_some(()) - .ok_or(LeeError::InvalidPrivacyPreservingProof) -} - -fn n_unique(data: &[T]) -> usize { - let set: HashSet<&T> = data.iter().collect(); - set.len() -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use lee_core::account::{Account, AccountId, Nonce}; - - use crate::{ - PrivateKey, PublicKey, V03State, - error::{InvalidProgramBehaviorError, LeeError}, - program::Program, - public_transaction::{Message, WitnessSet}, - validated_state_diff::ValidatedStateDiff, - }; - - fn public_state_from_balances( - initial_data: &[(AccountId, u128)], - ) -> HashMap { - initial_data - .iter() - .copied() - .map(|(account_id, balance)| { - ( - account_id, - Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance, - ..Account::default() - }, - ) - }) - .collect() - } - - #[test] - fn public_diff_reflects_a_successful_transfer() { - // A successful native transfer must record the debited sender in - // `public_diff()`. Catches the mutation that replaces `public_diff` with - // `HashMap::new()` (which would hide every account change). - let from_key = PrivateKey::try_new([1_u8; 32]).unwrap(); - let from = AccountId::from(&PublicKey::new_from_private_key(&from_key)); - let to_key = PrivateKey::try_new([2_u8; 32]).unwrap(); - let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); - - let state = V03State::new() - .with_public_accounts(public_state_from_balances(&[(from, 100)])) - .with_programs(std::iter::once( - crate::test_methods::simple_balance_transfer(), - )); - let program_id = crate::test_methods::simple_balance_transfer().id(); - let message = - Message::try_new(program_id, vec![from, to], vec![Nonce(0), Nonce(0)], 5_u128).unwrap(); - let witness_set = WitnessSet::for_message(&message, &[&from_key, &to_key]); - let tx = crate::PublicTransaction::new(message, witness_set); - - let diff = ValidatedStateDiff::from_public_transaction(&tx, &state, 1, 0) - .expect("a valid native transfer must validate"); - let public_diff = diff.public_diff(); - - assert!( - public_diff.contains_key(&from), - "public_diff must contain the debited sender", - ); - assert_eq!( - public_diff[&from].balance, 95, - "sender balance in the diff must reflect the debit", - ); - } - - /// Privacy-path version of the authorization-injection attack. The test passes when the - /// attack is rejected and the victim's balance is left untouched. - /// - /// `execute_and_prove` succeeds because each inner receipt is individually valid and the - /// outer circuit faithfully commits whatever the attacker's program output says, including - /// `victim(is_authorized=true)`. The circuit has no access to chain state and cannot know - /// the victim never signed. - /// - /// The host-side validator is what catches the attack: it independently reconstructs - /// `public_pre_states` from chain state using `signer_account_ids.contains(victim_id) = false`, - /// so it expects `victim(is_authorized=false)`. The committed journal and the reconstructed - /// expected output diverge, `receipt.verify` fails, and `from_privacy_preserving_transaction` - /// returns an error before any state is applied. - #[test] - fn privacy_malicious_programs_cannot_drain_public_victim() { - use lee_core::{ - Commitment, EncryptedAccountData, InputAccountIdentity, SharedSecretKey, - account::{Account, AccountWithMetadata}, - }; - - use crate::{ - PrivacyPreservingTransaction, - privacy_preserving_transaction::{ - circuit::{ProgramWithDependencies, execute_and_prove}, - message::Message, - witness_set::WitnessSet, - }, - state::{CommitmentSet, tests::test_private_account_keys_1}, - }; - - type InjectorInstruction = ( - lee_core::program::ProgramId, // p2_id - lee_core::program::ProgramId, // simple_balance_transfer_id - [u8; 32], // victim_id_raw - u128, // victim_balance - u128, // victim_nonce - lee_core::program::ProgramId, // victim_program_owner - [u8; 32], // recipient_id_raw - u128, // amount - ); - - // Attacker controls a private account. - let attacker_keys = test_private_account_keys_1(); - let attacker_id = AccountId::for_regular_private_account(&attacker_keys.npk(), 0); - let (attacker_ssk, attacker_epk) = SharedSecretKey::encapsulate(&attacker_keys.vpk()); - - let victim_id = AccountId::new([20_u8; 32]); - let recipient_id = AccountId::new([42_u8; 32]); - let victim_balance = 5_000_u128; - - // genesis sets program_owner = simple_balance_transfer_program.id() on all accounts. - let state = V03State::new() - .with_public_accounts(public_state_from_balances(&[ - (victim_id, victim_balance), - (recipient_id, 0), - ])) - .with_programs([ - crate::test_methods::simple_balance_transfer(), - crate::test_methods::malicious_injector(), - crate::test_methods::malicious_launderer(), - ]); - - // Build attacker's private account and its local commitment tree. - let attacker_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - ..Account::default() - }; - let attacker_commitment = Commitment::new(&attacker_id, &attacker_account); - let mut commitment_set = CommitmentSet::with_capacity(1); - commitment_set.extend(std::slice::from_ref(&attacker_commitment)); - let membership_proof = commitment_set - .get_proof_for(&attacker_commitment) - .expect("attacker commitment must be in the set"); - - let attacker_pre = AccountWithMetadata::new(attacker_account, true, attacker_id); - - let victim_account = state.get_account_by_id(victim_id); - let instruction: InjectorInstruction = ( - crate::test_methods::malicious_launderer().id(), - crate::test_methods::simple_balance_transfer().id(), - *victim_id.value(), - victim_account.balance, - victim_account.nonce.0, - victim_account.program_owner, - *recipient_id.value(), - victim_balance, - ); - let instruction_data = Program::serialize_instruction(instruction).unwrap(); - - let p2 = crate::test_methods::malicious_launderer(); - let at = crate::test_methods::simple_balance_transfer(); - let program_with_deps = ProgramWithDependencies::new( - crate::test_methods::malicious_injector(), - [(p2.id(), p2), (at.id(), at)].into(), - ); - - // account_identities order must match self.pre_states as built by the circuit: - // [0] attacker — first seen in P1's program_output.pre_states - // [1] victim — first seen in simple_balance_transfer's program_output.pre_states - // [2] recipient — first seen in simple_balance_transfer's program_output.pre_states - let account_identities = vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: attacker_epk, - view_tag: EncryptedAccountData::compute_view_tag( - &attacker_keys.npk(), - &attacker_keys.vpk(), - ), - ssk: attacker_ssk, - nsk: attacker_keys.nsk, - membership_proof, - identifier: 0, - }, - InputAccountIdentity::Public, // victim - InputAccountIdentity::Public, // recipient - ]; - - // execute_and_prove succeeds: all inner receipts are valid. - // The outer circuit commits victim(is_authorized=true) to its journal. - let (circuit_output, proof) = execute_and_prove( - vec![attacker_pre], - instruction_data, - account_identities, - &program_with_deps, - ) - .expect("execute_and_prove should succeed \u{2014} the programs execute correctly"); - - // public_account_ids lists the Public entries from account_identities, in order. - // The single ciphertext belongs to attacker's private account update. - let message = Message::try_from_circuit_output( - vec![victim_id, recipient_id], - vec![], // no public signers, no nonces - circuit_output, - ) - .unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures - let tx = PrivacyPreservingTransaction::new(message, witness_set); - - let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0); - - assert!( - matches!(result, Err(LeeError::InvalidPrivacyPreservingProof)), - "attack privacy transaction should be rejected with InvalidPrivacyPreservingProof" - ); - assert_eq!(state.get_account_by_id(victim_id).balance, victim_balance); - assert_eq!(state.get_account_by_id(recipient_id).balance, 0); - } - - /// Private-victim variant of the authorization-injection attack. The test passes when the - /// attack is rejected and the recipient's balance remains zero. - /// - /// After the circuit's Vacant branch accepts the injected `victim(is_authorized=true)` - /// verbatim, the attacker must choose how to declare the victim in `account_identities`. - /// There are two routes, both closed: - /// - /// - **mask=1 (`PrivateAuthorizedUpdate`)**: the circuit derives `account_id = - /// AccountId::for_regular_private_account(&npk_from(nsk), identifier)` and asserts it matches - /// `pre_state.account_id`. Passing this check requires the victim's `nsk`, which the attacker - /// does not have. `execute_and_prove` panics inside the ZKVM and no proof is produced. - /// - /// - **mask=0 (`Public`)**: the circuit places the account in `public_pre_states` and - /// `execute_and_prove` succeeds. The host-side validator then reconstructs - /// `public_pre_states` from chain state; `state.get_account_by_id(victim_id)` returns the - /// default account (balance=0) because the victim has no public state entry. The committed - /// journal and the reconstructed expected output diverge, `receipt.verify` fails, and - /// `from_privacy_preserving_transaction` returns an error before any state is applied. This - /// test exercises this route. - #[test] - fn privacy_malicious_programs_cannot_drain_private_victim() { - use lee_core::{ - Commitment, EncryptedAccountData, InputAccountIdentity, SharedSecretKey, - account::{Account, AccountWithMetadata}, - }; - - use crate::{ - PrivacyPreservingTransaction, - privacy_preserving_transaction::{ - circuit::{ProgramWithDependencies, execute_and_prove}, - message::Message, - witness_set::WitnessSet, - }, - state::{ - CommitmentSet, - tests::{test_private_account_keys_1, test_private_account_keys_2}, - }, - }; - - type InjectorInstruction = ( - lee_core::program::ProgramId, // p2_id - lee_core::program::ProgramId, // simple_balance_transfer_id - [u8; 32], // victim_id_raw - u128, // victim_balance - u128, // victim_nonce - lee_core::program::ProgramId, // victim_program_owner - [u8; 32], // recipient_id_raw - u128, // amount - ); - - // Attacker controls a private account. - let attacker_keys = test_private_account_keys_1(); - let attacker_id = AccountId::for_regular_private_account(&attacker_keys.npk(), 0); - let (attacker_ssk, attacker_epk) = SharedSecretKey::encapsulate(&attacker_keys.vpk()); - - // Victim is a private account — not registered in public chain state. - let victim_keys = test_private_account_keys_2(); - let victim_id = AccountId::for_regular_private_account(&victim_keys.npk(), 0); - let victim_balance = 5_000_u128; - - let recipient_id = AccountId::new([42_u8; 32]); - - // Victim has no public state entry; only recipient is registered at genesis. - let state = V03State::new() - .with_public_accounts(public_state_from_balances(&[(recipient_id, 0)])) - .with_programs([ - crate::test_methods::simple_balance_transfer(), - crate::test_methods::malicious_injector(), - crate::test_methods::malicious_launderer(), - ]); - - // Build attacker's private account and its local commitment tree. - let attacker_account = Account { - program_owner: crate::test_methods::simple_balance_transfer().id(), - balance: 100, - ..Account::default() - }; - let attacker_commitment = Commitment::new(&attacker_id, &attacker_account); - let mut commitment_set = CommitmentSet::with_capacity(1); - commitment_set.extend(std::slice::from_ref(&attacker_commitment)); - let membership_proof = commitment_set - .get_proof_for(&attacker_commitment) - .expect("attacker commitment must be in the set"); - - let attacker_pre = AccountWithMetadata::new(attacker_account, true, attacker_id); - - // The attacker supplies the victim's account data directly — it cannot be read from - // public state. The injected balance and program_owner allow simple_balance_transfer - // to succeed inside the circuit, which has no access to chain state and cannot detect - // that these values are fabricated. - let instruction: InjectorInstruction = ( - crate::test_methods::malicious_launderer().id(), - crate::test_methods::simple_balance_transfer().id(), - *victim_id.value(), - victim_balance, - 0_u128, // nonce - crate::test_methods::simple_balance_transfer().id(), // program_owner - *recipient_id.value(), - victim_balance, - ); - let instruction_data = Program::serialize_instruction(instruction).unwrap(); - - let p2 = crate::test_methods::malicious_launderer(); - let at = crate::test_methods::simple_balance_transfer(); - let program_with_deps = ProgramWithDependencies::new( - crate::test_methods::malicious_injector(), - [(p2.id(), p2), (at.id(), at)].into(), - ); - - // account_identities order must match self.pre_states as built by the circuit: - // [0] attacker — first seen in P1's program_output.pre_states - // [1] victim — first seen in simple_balance_transfer's program_output.pre_states - // [2] recipient — first seen in simple_balance_transfer's program_output.pre_states - // - // Victim is marked Public: the attacker has no nsk for the victim's private account, - // so PrivateAuthorizedUpdate is not an option. - let account_identities = vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: attacker_epk, - view_tag: EncryptedAccountData::compute_view_tag( - &attacker_keys.npk(), - &attacker_keys.vpk(), - ), - ssk: attacker_ssk, - nsk: attacker_keys.nsk, - membership_proof, - identifier: 0, - }, - InputAccountIdentity::Public, // victim — attacker lacks victim's nsk - InputAccountIdentity::Public, // recipient - ]; - - // execute_and_prove succeeds: simple_balance_transfer runs against the injected - // victim(balance=5000, is_authorized=true) and produces valid inner receipts. - // The outer circuit commits victim(is_authorized=true) to public_pre_states. - let (circuit_output, proof) = execute_and_prove( - vec![attacker_pre], - instruction_data, - account_identities, - &program_with_deps, - ) - .expect("execute_and_prove should succeed \u{2014} the programs execute correctly"); - - // public_account_ids lists the Public entries from account_identities, in order. - // The single ciphertext belongs to attacker's private account update. - let message = Message::try_from_circuit_output( - vec![victim_id, recipient_id], - vec![], // no public signers, no nonces - circuit_output, - ) - .unwrap(); - - let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures - let tx = PrivacyPreservingTransaction::new(message, witness_set); - - let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0); - - assert!( - matches!(result, Err(LeeError::InvalidPrivacyPreservingProof)), - "attack on private victim should be rejected with InvalidPrivacyPreservingProof" - ); - // Victim has no public balance to check; confirming the recipient received nothing - // is sufficient to show no funds moved. - assert_eq!(state.get_account_by_id(recipient_id).balance, 0); - } - - /// Two malicious programs (injector + launderer) attempt to drain a victim's balance - /// without the victim signing anything. The test passes when the attack is rejected - /// and the victim's balance is left untouched. - /// - /// Attack flow: - /// Transaction (attacker signs) → P1 (`malicious_injector`) - /// → injects `victim(is_authorized=true)` into chained-call `pre_states` for P2 - /// P2 (`malicious_launderer`) - /// → outputs empty pre/post states, forwarding the forged flag to `simple_balance_transfer` - /// → if `authorized_accounts` were built from the injected `pre_states`, - /// `{victim}.contains(victim)` would pass and the transfer would execute. - /// - /// The validator must reject this: `authorized_accounts` must be derived from the - /// parent program's own validated `program_output.pre_states`, not from the chained-call - /// input, so a forged `is_authorized=true` flag is never trusted. - #[test] - fn malicious_programs_cannot_drain_victim_without_signature() { - // p2_id, simple_balance_transfer_id, victim_id_raw, victim_balance, victim_nonce, - // victim_program_owner, recipient_id_raw, amount. - // Primitives only — AccountId/Account cannot round-trip through instruction_data - // via risc0_zkvm::serde (SerializeDisplay issue). - type InjectorInstruction = ( - lee_core::program::ProgramId, // p2_id - lee_core::program::ProgramId, // simple_balance_transfer_id - [u8; 32], // victim_id_raw - u128, // victim_balance - u128, // victim_nonce - lee_core::program::ProgramId, // victim_program_owner - [u8; 32], // recipient_id_raw - u128, // amount - ); - - let attacker_key = PrivateKey::try_new([10; 32]).unwrap(); - let attacker_id = AccountId::from(&PublicKey::new_from_private_key(&attacker_key)); - - let victim_key = PrivateKey::try_new([20; 32]).unwrap(); - let victim_id = AccountId::from(&PublicKey::new_from_private_key(&victim_key)); - - let recipient_id = AccountId::new([42; 32]); - - let victim_balance = 5_000_u128; - let state = V03State::new() - .with_public_accounts(public_state_from_balances(&[ - (attacker_id, 100), - (victim_id, victim_balance), - (recipient_id, 0), - ])) - .with_programs([ - crate::test_methods::simple_balance_transfer(), - crate::test_methods::malicious_injector(), - crate::test_methods::malicious_launderer(), - ]); - - // Read victim state from chain, exactly as the attacker would. - let victim_account = state.get_account_by_id(victim_id); - - let instruction: InjectorInstruction = ( - crate::test_methods::malicious_launderer().id(), - crate::test_methods::simple_balance_transfer().id(), - *victim_id.value(), - victim_account.balance, - victim_account.nonce.0, - victim_account.program_owner, - *recipient_id.value(), - victim_balance, - ); - - let message = Message::try_new( - crate::test_methods::malicious_injector().id(), - vec![attacker_id], - vec![Nonce(0)], - instruction, - ) - .unwrap(); - - let witness_set = WitnessSet::for_message(&message, &[&attacker_key]); - let tx = crate::PublicTransaction::new(message, witness_set); - - let result = ValidatedStateDiff::from_public_transaction(&tx, &state, 1, 0); - - assert!( - matches!( - result, - Err(LeeError::InvalidProgramBehavior( - InvalidProgramBehaviorError::InvalidAccountAuthorization { account_id } - )) if account_id == victim_id - ), - "attack transaction should be rejected with InvalidAccountAuthorization for the victim" - ); - - // Confirm the victim's balance is untouched. - let victim_balance_after = state.get_account_by_id(victim_id).balance; - let recipient_balance_after = state.get_account_by_id(recipient_id).balance; - - assert_eq!( - victim_balance_after, victim_balance, - "victim balance should be unchanged" - ); - assert_eq!( - recipient_balance_after, 0, - "recipient should receive nothing" - ); - } - - /// Regression test: a `PrivacyPreservingTransaction` carrying a structurally invalid - /// proof must be rejected with a clean `Err`. - #[test] - fn privacy_garbage_proof_is_rejected() { - use lee_core::{ - Commitment, - account::Account, - program::{BlockValidityWindow, TimestampValidityWindow}, - }; - - use crate::{ - PrivacyPreservingTransaction, - privacy_preserving_transaction::{ - circuit::Proof, message::Message, witness_set::WitnessSet, - }, - }; - - let state = V03State::new(); - - // Minimal message that passes every check up to proof verification: a single - // commitment satisfies the non-empty requirement, no signers makes the - // nonce/signature checks vacuously true, and unbounded validity windows are valid - // for any block/timestamp. - let account_id = AccountId::from(&PublicKey::new_from_private_key( - &PrivateKey::try_new([1_u8; 32]).unwrap(), - )); - let commitment = Commitment::new(&account_id, &Account::default()); - let message = Message { - public_account_ids: vec![], - nonces: vec![], - public_post_states: vec![], - encrypted_private_post_states: vec![], - new_commitments: vec![commitment], - new_nullifiers: vec![], - block_validity_window: BlockValidityWindow::new_unbounded(), - timestamp_validity_window: TimestampValidityWindow::new_unbounded(), - }; - - // Garbage proof bytes: not a valid borsh-encoded `InnerReceipt`. - let garbage_proof = Proof::from_inner(vec![0xff_u8; 64]); - let witness_set = WitnessSet::for_message(&message, garbage_proof, &[]); - let tx = PrivacyPreservingTransaction::new(message, witness_set); - - let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0); - - match result { - Err(LeeError::InvalidPrivacyPreservingProof) => {} - Err(other) => panic!("expected InvalidPrivacyPreservingProof, got {other:?}"), - Ok(_) => panic!("garbage proof was accepted instead of rejected"), - } - } -} diff --git a/lee/state_machine/src/validated_state_diff/mod.rs b/lee/state_machine/src/validated_state_diff/mod.rs new file mode 100644 index 00000000..8fae4fee --- /dev/null +++ b/lee/state_machine/src/validated_state_diff/mod.rs @@ -0,0 +1,521 @@ +use std::{ + collections::{HashMap, HashSet, VecDeque}, + hash::Hash, +}; + +use lee_core::{ + BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, Timestamp, + account::{Account, AccountId, AccountWithMetadata}, + program::{ + ChainedCall, Claim, DEFAULT_PROGRAM_ID, ProgramId, compute_public_authorized_pdas, + validate_execution, + }, +}; +use log::debug; + +use crate::{ + V03State, ensure, + error::{InvalidProgramBehaviorError, LeeError}, + privacy_preserving_transaction::{ + PrivacyPreservingTransaction, circuit::Proof, message::Message, + }, + program::Program, + program_deployment_transaction::ProgramDeploymentTransaction, + public_transaction::PublicTransaction, + state::MAX_NUMBER_CHAINED_CALLS, +}; + +pub struct StateDiff { + pub signer_account_ids: Vec, + pub public_diff: HashMap, + pub new_commitments: Vec, + pub new_nullifiers: Vec, + pub program: Option, +} + +/// The validated output of executing or verifying a transaction, ready to be applied to the state. +/// +/// Can only be constructed by the transaction validation functions inside this crate, ensuring the +/// diff has been checked before any state mutation occurs. +pub struct ValidatedStateDiff(StateDiff); + +impl ValidatedStateDiff { + pub fn from_public_transaction( + tx: &PublicTransaction, + state: &V03State, + block_id: BlockId, + timestamp: Timestamp, + ) -> Result { + let signer_account_ids = authenticate_public_transaction_signers(tx, state)?; + let message = tx.message(); + + ensure!( + !message.account_ids.is_empty(), + LeeError::InvalidInput("Public transaction must have at least one account".into()) + ); + + // All account_ids must be different + ensure!( + message.account_ids.iter().collect::>().len() == message.account_ids.len(), + LeeError::InvalidInput("Duplicate account_ids found in message".into(),) + ); + + // Build pre_states for execution + let input_pre_states: Vec<_> = message + .account_ids + .iter() + .map(|account_id| { + AccountWithMetadata::new( + state.get_account_by_id(*account_id), + signer_account_ids.contains(account_id), + *account_id, + ) + }) + .collect(); + + let mut state_diff: HashMap = HashMap::new(); + + let initial_call = ChainedCall { + program_id: message.program_id, + instruction_data: message.instruction_data.clone(), + pre_states: input_pre_states, + pda_seeds: vec![], + }; + + let initial_caller_data = CallerData { + program_id: None, + authorized_accounts: signer_account_ids.iter().copied().collect(), + }; + + let mut chained_calls = + VecDeque::<(ChainedCall, CallerData)>::from_iter([(initial_call, initial_caller_data)]); + let mut chain_calls_counter = 0; + + while let Some((chained_call, caller_data)) = chained_calls.pop_front() { + ensure!( + chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS, + LeeError::MaxChainedCallsDepthExceeded + ); + + // Check that the `program_id` corresponds to a deployed program + let Some(program) = state.programs().get(&chained_call.program_id) else { + return Err(LeeError::InvalidInput("Unknown program".into())); + }; + + debug!( + "Program {:?} pre_states: {:?}, instruction_data: {:?}", + chained_call.program_id, chained_call.pre_states, chained_call.instruction_data + ); + let mut program_output = program.execute( + caller_data.program_id, + &chained_call.pre_states, + &chained_call.instruction_data, + )?; + debug!( + "Program {:?} output: {:?}", + chained_call.program_id, program_output + ); + + let authorized_pdas = + compute_public_authorized_pdas(caller_data.program_id, &chained_call.pda_seeds); + + // Account is authorized if it is either in the caller's authorized accounts or in the + // list of PDAs the caller has authorized. + let is_authorized = |account_id: &AccountId| { + authorized_pdas.contains(account_id) + || caller_data.authorized_accounts.contains(account_id) + }; + + for pre in &program_output.pre_states { + let account_id = pre.account_id; + // Check that the program output pre_states coincide with the values in the public + // state or with any modifications to those values during the chain of calls. + let expected_pre = state_diff + .get(&account_id) + .cloned() + .unwrap_or_else(|| state.get_account_by_id(account_id)); + ensure!( + pre.account == expected_pre, + InvalidProgramBehaviorError::InconsistentAccountPreState { + account_id, + expected: Box::new(expected_pre), + actual: Box::new(pre.account.clone()) + } + ); + + // Check that the program output pre_states marked as authorized are indeed + // authorized, and vice-versa. + let is_indeed_authorized = is_authorized(&account_id); + ensure!( + !pre.is_authorized || is_indeed_authorized, + InvalidProgramBehaviorError::InvalidAccountAuthorization { account_id } + ); + ensure!( + pre.is_authorized || !is_indeed_authorized, + InvalidProgramBehaviorError::AuthorizedAccountMarkedAsNotAuthorized { + account_id + } + ); + } + + // Verify that the program output's self_program_id matches the expected program ID. + ensure!( + program_output.self_program_id == chained_call.program_id, + InvalidProgramBehaviorError::MismatchedProgramId { + expected: chained_call.program_id, + actual: program_output.self_program_id + } + ); + + // Verify that the program output's caller_program_id matches the actual caller. + ensure!( + program_output.caller_program_id == caller_data.program_id, + InvalidProgramBehaviorError::MismatchedCallerProgramId { + expected: caller_data.program_id, + actual: program_output.caller_program_id, + } + ); + + // Verify execution corresponds to a well-behaved program. + // See the # Programs section for the definition of the `validate_execution` method. + validate_execution( + &program_output.pre_states, + &program_output.post_states, + chained_call.program_id, + ) + .map_err(InvalidProgramBehaviorError::ExecutionValidationFailed)?; + + // Verify validity window + ensure!( + program_output.block_validity_window.is_valid_for(block_id) + && program_output + .timestamp_validity_window + .is_valid_for(timestamp), + LeeError::OutOfValidityWindow + ); + + for (i, post) in program_output.post_states.iter_mut().enumerate() { + let Some(claim) = post.required_claim() else { + continue; + }; + let pre = &program_output.pre_states[i]; + let account_id = pre.account_id; + + // The invoked program can only claim accounts with default program id. + ensure!( + post.account().program_owner == DEFAULT_PROGRAM_ID, + InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id } + ); + + match claim { + Claim::Authorized => { + // The program can only claim accounts that were authorized by the signer. + ensure!( + pre.is_authorized, + InvalidProgramBehaviorError::ClaimedUnauthorizedAccount { account_id } + ); + } + Claim::Pda(seed) => { + // The program can only claim accounts that correspond to the PDAs it is + // authorized to claim. The public-execution path only sees public + // accounts, so the public-PDA derivation is the correct formula here. + let pda = AccountId::for_public_pda(&chained_call.program_id, &seed); + ensure!( + account_id == pda, + InvalidProgramBehaviorError::MismatchedPdaClaim { + expected: pda, + actual: account_id + } + ); + } + } + + post.account_mut().program_owner = chained_call.program_id; + } + + // Update the state diff + for (pre, post) in program_output + .pre_states + .iter() + .zip(program_output.post_states.iter()) + { + state_diff.insert(pre.account_id, post.account().clone()); + } + + // Source from `program_output.pre_states`, not `chained_call.pre_states`: + // the loop above already gates program_output's `is_authorized` via the + // `!pre.is_authorized || is_indeed_authorized` check, while `chained_call. + // pre_states` is caller-controlled and can be forged (audit-issue 91). + // + // Union with the caller's authorized set so that authorization is monotonically + // growing: once an account is authorized at any point in the chain it remains + // authorized for all subsequent calls. + let authorized_accounts: HashSet<_> = caller_data + .authorized_accounts + .into_iter() + .chain( + program_output + .pre_states + .iter() + .filter(|pre| pre.is_authorized) + .map(|pre| pre.account_id), + ) + .collect(); + for new_call in program_output.chained_calls.into_iter().rev() { + chained_calls.push_front(( + new_call, + CallerData { + program_id: Some(chained_call.program_id), + authorized_accounts: authorized_accounts.clone(), + }, + )); + } + + chain_calls_counter = chain_calls_counter + .checked_add(1) + .expect("we check the max depth at the beginning of the loop"); + } + + // Check that all modified uninitialized accounts where claimed + for (account_id, post) in state_diff.iter().filter_map(|(account_id, post)| { + let pre = state.get_account_by_id(*account_id); + if pre.program_owner != DEFAULT_PROGRAM_ID { + return None; + } + if pre == *post { + return None; + } + Some((*account_id, post)) + }) { + ensure!( + post.program_owner != DEFAULT_PROGRAM_ID, + InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim { account_id } + ); + } + + Ok(Self(StateDiff { + signer_account_ids, + public_diff: state_diff, + new_commitments: vec![], + new_nullifiers: vec![], + program: None, + })) + } + + pub fn from_privacy_preserving_transaction( + tx: &PrivacyPreservingTransaction, + state: &V03State, + block_id: BlockId, + timestamp: Timestamp, + ) -> Result { + let message = &tx.message; + let witness_set = &tx.witness_set; + + // 1. Commitments or nullifiers are non empty + ensure!( + !message.new_commitments.is_empty() || !message.new_nullifiers.is_empty(), + LeeError::InvalidInput( + "Empty commitments and empty nullifiers found in message".into(), + ) + ); + + // 2. Check there are no duplicate account_ids in the public_account_ids list. + ensure!( + n_unique(&message.public_account_ids) == message.public_account_ids.len(), + LeeError::InvalidInput("Duplicate account_ids found in message".into()) + ); + + // Check there are no duplicate nullifiers in the new_nullifiers list + ensure!( + n_unique( + &message + .new_nullifiers + .iter() + .map(|(n, _)| n) + .collect::>() + ) == message.new_nullifiers.len(), + LeeError::InvalidInput("Duplicate nullifiers found in message".into()) + ); + + // Check there are no duplicate commitments in the new_commitments list + ensure!( + n_unique(&message.new_commitments) == message.new_commitments.len(), + LeeError::InvalidInput("Duplicate commitments found in message".into()) + ); + + // 3. Nonce checks and Valid signatures + // Check exactly one nonce is provided for each signature + ensure!( + message.nonces.len() == witness_set.signatures_and_public_keys.len(), + LeeError::InvalidInput( + "Mismatch between number of nonces and signatures/public keys".into(), + ) + ); + + // Check the signatures are valid + ensure!( + witness_set.signatures_are_valid_for(message), + LeeError::InvalidInput("Invalid signature for given message and public key".into()) + ); + + let signer_account_ids = tx.signer_account_ids(); + // Check nonces corresponds to the current nonces on the public state. + for (account_id, nonce) in signer_account_ids.iter().zip(&message.nonces) { + let current_nonce = state.get_account_by_id(*account_id).nonce; + ensure!( + current_nonce == *nonce, + LeeError::InvalidInput("Nonce mismatch".into()) + ); + } + + // Verify validity window + ensure!( + message.block_validity_window.is_valid_for(block_id) + && message.timestamp_validity_window.is_valid_for(timestamp), + LeeError::OutOfValidityWindow + ); + + // Build pre_states for proof verification + let public_pre_states: Vec<_> = message + .public_account_ids + .iter() + .map(|account_id| { + AccountWithMetadata::new( + state.get_account_by_id(*account_id), + signer_account_ids.contains(account_id), + *account_id, + ) + }) + .collect(); + + // 4. Proof verification + check_privacy_preserving_circuit_proof_is_valid( + &witness_set.proof, + &public_pre_states, + message, + )?; + + // 5. Commitment freshness + state.check_commitments_are_new(&message.new_commitments)?; + + // 6. Nullifier uniqueness + state.check_nullifiers_are_valid(&message.new_nullifiers)?; + + let public_diff = message + .public_account_ids + .iter() + .copied() + .zip(message.public_post_states.clone()) + .collect(); + let new_nullifiers = message + .new_nullifiers + .iter() + .copied() + .map(|(nullifier, _)| nullifier) + .collect(); + + Ok(Self(StateDiff { + signer_account_ids, + public_diff, + new_commitments: message.new_commitments.clone(), + new_nullifiers, + program: None, + })) + } + + pub fn from_program_deployment_transaction( + tx: &ProgramDeploymentTransaction, + state: &V03State, + ) -> Result { + // TODO: remove clone + let program = Program::new(tx.message.bytecode.clone().into())?; + if state.programs().contains_key(&program.id()) { + return Err(LeeError::ProgramAlreadyExists); + } + Ok(Self(StateDiff { + signer_account_ids: vec![], + public_diff: HashMap::new(), + new_commitments: vec![], + new_nullifiers: vec![], + program: Some(program), + })) + } + + /// Returns the public account changes produced by this transaction. + /// + /// Used by callers (e.g. the sequencer) to inspect the diff before committing it, for example + /// to enforce that system accounts are not modified by user transactions. + #[must_use] + pub fn public_diff(&self) -> HashMap { + self.0.public_diff.clone() + } + + pub(crate) fn into_state_diff(self) -> StateDiff { + self.0 + } +} + +#[derive(Debug)] +struct CallerData { + program_id: Option, + authorized_accounts: HashSet, +} + +fn authenticate_public_transaction_signers( + tx: &PublicTransaction, + state: &V03State, +) -> Result, LeeError> { + let message = tx.message(); + let witness_set = tx.witness_set(); + + ensure!( + message.nonces.len() == witness_set.signatures_and_public_keys.len(), + LeeError::InvalidInput( + "Mismatch between number of nonces and signatures/public keys".into(), + ) + ); + + ensure!( + witness_set.is_valid_for(message), + LeeError::InvalidInput("Invalid signature for given message and public key".into()) + ); + + let signer_account_ids = tx.signer_account_ids(); + for (account_id, nonce) in signer_account_ids.iter().zip(&message.nonces) { + let current_nonce = state.get_account_by_id(*account_id).nonce; + ensure!( + current_nonce == *nonce, + LeeError::InvalidInput("Nonce mismatch".into()) + ); + } + + Ok(signer_account_ids) +} + +fn check_privacy_preserving_circuit_proof_is_valid( + proof: &Proof, + public_pre_states: &[AccountWithMetadata], + message: &Message, +) -> Result<(), LeeError> { + let output = PrivacyPreservingCircuitOutput { + public_pre_states: public_pre_states.to_vec(), + public_post_states: message.public_post_states.clone(), + encrypted_private_post_states: message.encrypted_private_post_states.clone(), + new_commitments: message.new_commitments.clone(), + new_nullifiers: message.new_nullifiers.clone(), + block_validity_window: message.block_validity_window, + timestamp_validity_window: message.timestamp_validity_window, + }; + proof + .is_valid_for(&output) + .then_some(()) + .ok_or(LeeError::InvalidPrivacyPreservingProof) +} + +fn n_unique(data: &[T]) -> usize { + let set: HashSet<&T> = data.iter().collect(); + set.len() +} + +#[cfg(test)] +mod tests; diff --git a/lee/state_machine/src/validated_state_diff/tests.rs b/lee/state_machine/src/validated_state_diff/tests.rs new file mode 100644 index 00000000..3cb61e7d --- /dev/null +++ b/lee/state_machine/src/validated_state_diff/tests.rs @@ -0,0 +1,533 @@ +use std::collections::HashMap; + +use lee_core::account::{Account, AccountId, Nonce}; + +use crate::{ + PrivateKey, PublicKey, V03State, + error::{InvalidProgramBehaviorError, LeeError}, + program::Program, + public_transaction::{Message, WitnessSet}, + validated_state_diff::ValidatedStateDiff, +}; + +fn public_state_from_balances(initial_data: &[(AccountId, u128)]) -> HashMap { + initial_data + .iter() + .copied() + .map(|(account_id, balance)| { + ( + account_id, + Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance, + ..Account::default() + }, + ) + }) + .collect() +} + +#[test] +fn public_diff_reflects_a_successful_transfer() { + // A successful native transfer must record the debited sender in + // `public_diff()`. Catches the mutation that replaces `public_diff` with + // `HashMap::new()` (which would hide every account change). + let from_key = PrivateKey::try_new([1_u8; 32]).unwrap(); + let from = AccountId::from(&PublicKey::new_from_private_key(&from_key)); + let to_key = PrivateKey::try_new([2_u8; 32]).unwrap(); + let to = AccountId::from(&PublicKey::new_from_private_key(&to_key)); + + let state = V03State::new() + .with_public_accounts(public_state_from_balances(&[(from, 100)])) + .with_programs(std::iter::once( + crate::test_methods::simple_balance_transfer(), + )); + let program_id = crate::test_methods::simple_balance_transfer().id(); + let message = + Message::try_new(program_id, vec![from, to], vec![Nonce(0), Nonce(0)], 5_u128).unwrap(); + let witness_set = WitnessSet::for_message(&message, &[&from_key, &to_key]); + let tx = crate::PublicTransaction::new(message, witness_set); + + let diff = ValidatedStateDiff::from_public_transaction(&tx, &state, 1, 0) + .expect("a valid native transfer must validate"); + let public_diff = diff.public_diff(); + + assert!( + public_diff.contains_key(&from), + "public_diff must contain the debited sender", + ); + assert_eq!( + public_diff[&from].balance, 95, + "sender balance in the diff must reflect the debit", + ); +} + +/// Privacy-path version of the authorization-injection attack. The test passes when the +/// attack is rejected and the victim's balance is left untouched. +/// +/// `execute_and_prove` succeeds because each inner receipt is individually valid and the +/// outer circuit faithfully commits whatever the attacker's program output says, including +/// `victim(is_authorized=true)`. The circuit has no access to chain state and cannot know +/// the victim never signed. +/// +/// The host-side validator is what catches the attack: it independently reconstructs +/// `public_pre_states` from chain state using `signer_account_ids.contains(victim_id) = false`, +/// so it expects `victim(is_authorized=false)`. The committed journal and the reconstructed +/// expected output diverge, `receipt.verify` fails, and `from_privacy_preserving_transaction` +/// returns an error before any state is applied. +#[test] +fn privacy_malicious_programs_cannot_drain_public_victim() { + use lee_core::{ + Commitment, EncryptedAccountData, InputAccountIdentity, SharedSecretKey, + account::{Account, AccountWithMetadata}, + }; + + use crate::{ + PrivacyPreservingTransaction, + privacy_preserving_transaction::{ + circuit::{ProgramWithDependencies, execute_and_prove}, + message::Message, + witness_set::WitnessSet, + }, + state::{CommitmentSet, tests::test_private_account_keys_1}, + }; + + type InjectorInstruction = ( + lee_core::program::ProgramId, // p2_id + lee_core::program::ProgramId, // simple_balance_transfer_id + [u8; 32], // victim_id_raw + u128, // victim_balance + u128, // victim_nonce + lee_core::program::ProgramId, // victim_program_owner + [u8; 32], // recipient_id_raw + u128, // amount + ); + + // Attacker controls a private account. + let attacker_keys = test_private_account_keys_1(); + let attacker_id = AccountId::for_regular_private_account(&attacker_keys.npk(), 0); + let (attacker_ssk, attacker_epk) = SharedSecretKey::encapsulate(&attacker_keys.vpk()); + + let victim_id = AccountId::new([20_u8; 32]); + let recipient_id = AccountId::new([42_u8; 32]); + let victim_balance = 5_000_u128; + + // genesis sets program_owner = simple_balance_transfer_program.id() on all accounts. + let state = V03State::new() + .with_public_accounts(public_state_from_balances(&[ + (victim_id, victim_balance), + (recipient_id, 0), + ])) + .with_programs([ + crate::test_methods::simple_balance_transfer(), + crate::test_methods::malicious_injector(), + crate::test_methods::malicious_launderer(), + ]); + + // Build attacker's private account and its local commitment tree. + let attacker_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + ..Account::default() + }; + let attacker_commitment = Commitment::new(&attacker_id, &attacker_account); + let mut commitment_set = CommitmentSet::with_capacity(1); + commitment_set.extend(std::slice::from_ref(&attacker_commitment)); + let membership_proof = commitment_set + .get_proof_for(&attacker_commitment) + .expect("attacker commitment must be in the set"); + + let attacker_pre = AccountWithMetadata::new(attacker_account, true, attacker_id); + + let victim_account = state.get_account_by_id(victim_id); + let instruction: InjectorInstruction = ( + crate::test_methods::malicious_launderer().id(), + crate::test_methods::simple_balance_transfer().id(), + *victim_id.value(), + victim_account.balance, + victim_account.nonce.0, + victim_account.program_owner, + *recipient_id.value(), + victim_balance, + ); + let instruction_data = Program::serialize_instruction(instruction).unwrap(); + + let p2 = crate::test_methods::malicious_launderer(); + let at = crate::test_methods::simple_balance_transfer(); + let program_with_deps = ProgramWithDependencies::new( + crate::test_methods::malicious_injector(), + [(p2.id(), p2), (at.id(), at)].into(), + ); + + // account_identities order must match self.pre_states as built by the circuit: + // [0] attacker — first seen in P1's program_output.pre_states + // [1] victim — first seen in simple_balance_transfer's program_output.pre_states + // [2] recipient — first seen in simple_balance_transfer's program_output.pre_states + let account_identities = vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: attacker_epk, + view_tag: EncryptedAccountData::compute_view_tag( + &attacker_keys.npk(), + &attacker_keys.vpk(), + ), + ssk: attacker_ssk, + nsk: attacker_keys.nsk, + membership_proof, + identifier: 0, + }, + InputAccountIdentity::Public, // victim + InputAccountIdentity::Public, // recipient + ]; + + // execute_and_prove succeeds: all inner receipts are valid. + // The outer circuit commits victim(is_authorized=true) to its journal. + let (circuit_output, proof) = execute_and_prove( + vec![attacker_pre], + instruction_data, + account_identities, + &program_with_deps, + ) + .expect("execute_and_prove should succeed \u{2014} the programs execute correctly"); + + // public_account_ids lists the Public entries from account_identities, in order. + // The single ciphertext belongs to attacker's private account update. + let message = Message::try_from_circuit_output( + vec![victim_id, recipient_id], + vec![], // no public signers, no nonces + circuit_output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0); + + assert!( + matches!(result, Err(LeeError::InvalidPrivacyPreservingProof)), + "attack privacy transaction should be rejected with InvalidPrivacyPreservingProof" + ); + assert_eq!(state.get_account_by_id(victim_id).balance, victim_balance); + assert_eq!(state.get_account_by_id(recipient_id).balance, 0); +} + +/// Private-victim variant of the authorization-injection attack. The test passes when the +/// attack is rejected and the recipient's balance remains zero. +/// +/// After the circuit's Vacant branch accepts the injected `victim(is_authorized=true)` +/// verbatim, the attacker must choose how to declare the victim in `account_identities`. +/// There are two routes, both closed: +/// +/// - **mask=1 (`PrivateAuthorizedUpdate`)**: the circuit derives `account_id = +/// AccountId::for_regular_private_account(&npk_from(nsk), identifier)` and asserts it matches +/// `pre_state.account_id`. Passing this check requires the victim's `nsk`, which the attacker +/// does not have. `execute_and_prove` panics inside the ZKVM and no proof is produced. +/// +/// - **mask=0 (`Public`)**: the circuit places the account in `public_pre_states` and +/// `execute_and_prove` succeeds. The host-side validator then reconstructs `public_pre_states` +/// from chain state; `state.get_account_by_id(victim_id)` returns the default account (balance=0) +/// because the victim has no public state entry. The committed journal and the reconstructed +/// expected output diverge, `receipt.verify` fails, and `from_privacy_preserving_transaction` +/// returns an error before any state is applied. This test exercises this route. +#[test] +fn privacy_malicious_programs_cannot_drain_private_victim() { + use lee_core::{ + Commitment, EncryptedAccountData, InputAccountIdentity, SharedSecretKey, + account::{Account, AccountWithMetadata}, + }; + + use crate::{ + PrivacyPreservingTransaction, + privacy_preserving_transaction::{ + circuit::{ProgramWithDependencies, execute_and_prove}, + message::Message, + witness_set::WitnessSet, + }, + state::{ + CommitmentSet, + tests::{test_private_account_keys_1, test_private_account_keys_2}, + }, + }; + + type InjectorInstruction = ( + lee_core::program::ProgramId, // p2_id + lee_core::program::ProgramId, // simple_balance_transfer_id + [u8; 32], // victim_id_raw + u128, // victim_balance + u128, // victim_nonce + lee_core::program::ProgramId, // victim_program_owner + [u8; 32], // recipient_id_raw + u128, // amount + ); + + // Attacker controls a private account. + let attacker_keys = test_private_account_keys_1(); + let attacker_id = AccountId::for_regular_private_account(&attacker_keys.npk(), 0); + let (attacker_ssk, attacker_epk) = SharedSecretKey::encapsulate(&attacker_keys.vpk()); + + // Victim is a private account — not registered in public chain state. + let victim_keys = test_private_account_keys_2(); + let victim_id = AccountId::for_regular_private_account(&victim_keys.npk(), 0); + let victim_balance = 5_000_u128; + + let recipient_id = AccountId::new([42_u8; 32]); + + // Victim has no public state entry; only recipient is registered at genesis. + let state = V03State::new() + .with_public_accounts(public_state_from_balances(&[(recipient_id, 0)])) + .with_programs([ + crate::test_methods::simple_balance_transfer(), + crate::test_methods::malicious_injector(), + crate::test_methods::malicious_launderer(), + ]); + + // Build attacker's private account and its local commitment tree. + let attacker_account = Account { + program_owner: crate::test_methods::simple_balance_transfer().id(), + balance: 100, + ..Account::default() + }; + let attacker_commitment = Commitment::new(&attacker_id, &attacker_account); + let mut commitment_set = CommitmentSet::with_capacity(1); + commitment_set.extend(std::slice::from_ref(&attacker_commitment)); + let membership_proof = commitment_set + .get_proof_for(&attacker_commitment) + .expect("attacker commitment must be in the set"); + + let attacker_pre = AccountWithMetadata::new(attacker_account, true, attacker_id); + + // The attacker supplies the victim's account data directly — it cannot be read from + // public state. The injected balance and program_owner allow simple_balance_transfer + // to succeed inside the circuit, which has no access to chain state and cannot detect + // that these values are fabricated. + let instruction: InjectorInstruction = ( + crate::test_methods::malicious_launderer().id(), + crate::test_methods::simple_balance_transfer().id(), + *victim_id.value(), + victim_balance, + 0_u128, // nonce + crate::test_methods::simple_balance_transfer().id(), // program_owner + *recipient_id.value(), + victim_balance, + ); + let instruction_data = Program::serialize_instruction(instruction).unwrap(); + + let p2 = crate::test_methods::malicious_launderer(); + let at = crate::test_methods::simple_balance_transfer(); + let program_with_deps = ProgramWithDependencies::new( + crate::test_methods::malicious_injector(), + [(p2.id(), p2), (at.id(), at)].into(), + ); + + // account_identities order must match self.pre_states as built by the circuit: + // [0] attacker — first seen in P1's program_output.pre_states + // [1] victim — first seen in simple_balance_transfer's program_output.pre_states + // [2] recipient — first seen in simple_balance_transfer's program_output.pre_states + // + // Victim is marked Public: the attacker has no nsk for the victim's private account, + // so PrivateAuthorizedUpdate is not an option. + let account_identities = vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: attacker_epk, + view_tag: EncryptedAccountData::compute_view_tag( + &attacker_keys.npk(), + &attacker_keys.vpk(), + ), + ssk: attacker_ssk, + nsk: attacker_keys.nsk, + membership_proof, + identifier: 0, + }, + InputAccountIdentity::Public, // victim — attacker lacks victim's nsk + InputAccountIdentity::Public, // recipient + ]; + + // execute_and_prove succeeds: simple_balance_transfer runs against the injected + // victim(balance=5000, is_authorized=true) and produces valid inner receipts. + // The outer circuit commits victim(is_authorized=true) to public_pre_states. + let (circuit_output, proof) = execute_and_prove( + vec![attacker_pre], + instruction_data, + account_identities, + &program_with_deps, + ) + .expect("execute_and_prove should succeed \u{2014} the programs execute correctly"); + + // public_account_ids lists the Public entries from account_identities, in order. + // The single ciphertext belongs to attacker's private account update. + let message = Message::try_from_circuit_output( + vec![victim_id, recipient_id], + vec![], // no public signers, no nonces + circuit_output, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0); + + assert!( + matches!(result, Err(LeeError::InvalidPrivacyPreservingProof)), + "attack on private victim should be rejected with InvalidPrivacyPreservingProof" + ); + // Victim has no public balance to check; confirming the recipient received nothing + // is sufficient to show no funds moved. + assert_eq!(state.get_account_by_id(recipient_id).balance, 0); +} + +/// Two malicious programs (injector + launderer) attempt to drain a victim's balance +/// without the victim signing anything. The test passes when the attack is rejected +/// and the victim's balance is left untouched. +/// +/// Attack flow: +/// Transaction (attacker signs) → P1 (`malicious_injector`) +/// → injects `victim(is_authorized=true)` into chained-call `pre_states` for P2 +/// P2 (`malicious_launderer`) +/// → outputs empty pre/post states, forwarding the forged flag to `simple_balance_transfer` +/// → if `authorized_accounts` were built from the injected `pre_states`, +/// `{victim}.contains(victim)` would pass and the transfer would execute. +/// +/// The validator must reject this: `authorized_accounts` must be derived from the +/// parent program's own validated `program_output.pre_states`, not from the chained-call +/// input, so a forged `is_authorized=true` flag is never trusted. +#[test] +fn malicious_programs_cannot_drain_victim_without_signature() { + // p2_id, simple_balance_transfer_id, victim_id_raw, victim_balance, victim_nonce, + // victim_program_owner, recipient_id_raw, amount. + // Primitives only — AccountId/Account cannot round-trip through instruction_data + // via risc0_zkvm::serde (SerializeDisplay issue). + type InjectorInstruction = ( + lee_core::program::ProgramId, // p2_id + lee_core::program::ProgramId, // simple_balance_transfer_id + [u8; 32], // victim_id_raw + u128, // victim_balance + u128, // victim_nonce + lee_core::program::ProgramId, // victim_program_owner + [u8; 32], // recipient_id_raw + u128, // amount + ); + + let attacker_key = PrivateKey::try_new([10; 32]).unwrap(); + let attacker_id = AccountId::from(&PublicKey::new_from_private_key(&attacker_key)); + + let victim_key = PrivateKey::try_new([20; 32]).unwrap(); + let victim_id = AccountId::from(&PublicKey::new_from_private_key(&victim_key)); + + let recipient_id = AccountId::new([42; 32]); + + let victim_balance = 5_000_u128; + let state = V03State::new() + .with_public_accounts(public_state_from_balances(&[ + (attacker_id, 100), + (victim_id, victim_balance), + (recipient_id, 0), + ])) + .with_programs([ + crate::test_methods::simple_balance_transfer(), + crate::test_methods::malicious_injector(), + crate::test_methods::malicious_launderer(), + ]); + + // Read victim state from chain, exactly as the attacker would. + let victim_account = state.get_account_by_id(victim_id); + + let instruction: InjectorInstruction = ( + crate::test_methods::malicious_launderer().id(), + crate::test_methods::simple_balance_transfer().id(), + *victim_id.value(), + victim_account.balance, + victim_account.nonce.0, + victim_account.program_owner, + *recipient_id.value(), + victim_balance, + ); + + let message = Message::try_new( + crate::test_methods::malicious_injector().id(), + vec![attacker_id], + vec![Nonce(0)], + instruction, + ) + .unwrap(); + + let witness_set = WitnessSet::for_message(&message, &[&attacker_key]); + let tx = crate::PublicTransaction::new(message, witness_set); + + let result = ValidatedStateDiff::from_public_transaction(&tx, &state, 1, 0); + + assert!( + matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::InvalidAccountAuthorization { account_id } + )) if account_id == victim_id + ), + "attack transaction should be rejected with InvalidAccountAuthorization for the victim" + ); + + // Confirm the victim's balance is untouched. + let victim_balance_after = state.get_account_by_id(victim_id).balance; + let recipient_balance_after = state.get_account_by_id(recipient_id).balance; + + assert_eq!( + victim_balance_after, victim_balance, + "victim balance should be unchanged" + ); + assert_eq!( + recipient_balance_after, 0, + "recipient should receive nothing" + ); +} + +/// Regression test: a `PrivacyPreservingTransaction` carrying a structurally invalid +/// proof must be rejected with a clean `Err`. +#[test] +fn privacy_garbage_proof_is_rejected() { + use lee_core::{ + Commitment, + account::Account, + program::{BlockValidityWindow, TimestampValidityWindow}, + }; + + use crate::{ + PrivacyPreservingTransaction, + privacy_preserving_transaction::{ + circuit::Proof, message::Message, witness_set::WitnessSet, + }, + }; + + let state = V03State::new(); + + // Minimal message that passes every check up to proof verification: a single + // commitment satisfies the non-empty requirement, no signers makes the + // nonce/signature checks vacuously true, and unbounded validity windows are valid + // for any block/timestamp. + let account_id = AccountId::from(&PublicKey::new_from_private_key( + &PrivateKey::try_new([1_u8; 32]).unwrap(), + )); + let commitment = Commitment::new(&account_id, &Account::default()); + let message = Message { + public_account_ids: vec![], + nonces: vec![], + public_post_states: vec![], + encrypted_private_post_states: vec![], + new_commitments: vec![commitment], + new_nullifiers: vec![], + block_validity_window: BlockValidityWindow::new_unbounded(), + timestamp_validity_window: TimestampValidityWindow::new_unbounded(), + }; + + // Garbage proof bytes: not a valid borsh-encoded `InnerReceipt`. + let garbage_proof = Proof::from_inner(vec![0xff_u8; 64]); + let witness_set = WitnessSet::for_message(&message, garbage_proof, &[]); + let tx = PrivacyPreservingTransaction::new(message, witness_set); + + let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0); + + match result { + Err(LeeError::InvalidPrivacyPreservingProof) => {} + Err(other) => panic!("expected InvalidPrivacyPreservingProof, got {other:?}"), + Ok(_) => panic!("garbage proof was accepted instead of rejected"), + } +} diff --git a/lez/explorer_service/src/components/account_nonce_list.rs b/lez/explorer_service/src/components/account_nonce_list.rs new file mode 100644 index 00000000..d8b30aff --- /dev/null +++ b/lez/explorer_service/src/components/account_nonce_list.rs @@ -0,0 +1,58 @@ +use indexer_service_protocol::AccountId; +use itertools::{EitherOrBoth, Itertools as _}; +use leptos::prelude::*; +use leptos_router::components::A; + +#[component] +pub fn AccountNonceList(account_ids: Vec, nonces: Vec) -> impl IntoView { + view! { +
+ {account_ids + .into_iter() + .zip_longest(nonces.into_iter()) + .map(|maybe_pair| { + match maybe_pair { + EitherOrBoth::Both(account_id, nonce) => { + let account_id_str = account_id.to_string(); + view! { + + } + } + EitherOrBoth::Left(account_id) => { + let account_id_str = account_id.to_string(); + view! { + + } + } + EitherOrBoth::Right(_) => { + view! { + + } + } + } + }) + .collect::>()} +
+ } +} diff --git a/lez/explorer_service/src/components/mod.rs b/lez/explorer_service/src/components/mod.rs index 306c79a8..3d0a4dae 100644 --- a/lez/explorer_service/src/components/mod.rs +++ b/lez/explorer_service/src/components/mod.rs @@ -1,7 +1,15 @@ +pub use account_nonce_list::AccountNonceList; pub use account_preview::AccountPreview; pub use block_preview::BlockPreview; +pub use search_results::SearchResultsView; +pub use transaction_details::{ + PrivacyPreservingTxDetails, ProgramDeploymentTxDetails, PublicTxDetails, +}; pub use transaction_preview::TransactionPreview; +pub mod account_nonce_list; pub mod account_preview; pub mod block_preview; +pub mod search_results; +pub mod transaction_details; pub mod transaction_preview; diff --git a/lez/explorer_service/src/components/search_results.rs b/lez/explorer_service/src/components/search_results.rs new file mode 100644 index 00000000..1033e515 --- /dev/null +++ b/lez/explorer_service/src/components/search_results.rs @@ -0,0 +1,93 @@ +use leptos::prelude::*; + +use super::{AccountPreview, BlockPreview, TransactionPreview}; +use crate::api::SearchResults; + +/// Search results view component +#[component] +pub fn SearchResultsView(results: SearchResults) -> impl IntoView { + let SearchResults { + blocks, + transactions, + accounts, + } = results; + let has_results = !blocks.is_empty() || !transactions.is_empty() || !accounts.is_empty(); + + view! { +
+

"Search Results"

+ {if has_results { + view! { +
+ {if blocks.is_empty() { + ().into_any() + } else { + view! { +
+

"Blocks"

+
+ {blocks + .into_iter() + .map(|block| { + view! { } + }) + .collect::>()} +
+
+ } + .into_any() + }} + + {if transactions.is_empty() { + ().into_any() + } else { + view! { +
+

"Transactions"

+
+ {transactions + .into_iter() + .map(|tx| { + view! { } + }) + .collect::>()} +
+
+ } + .into_any() + }} + + {if accounts.is_empty() { + ().into_any() + } else { + view! { +
+

"Accounts"

+
+ {accounts + .into_iter() + .map(|(id, account)| { + view! { + + } + }) + .collect::>()} +
+
+ } + .into_any() + }} + +
+ } + .into_any() + } else { + view! {
"No results found"
} + .into_any() + }} +
+ } +} diff --git a/lez/explorer_service/src/components/transaction_details.rs b/lez/explorer_service/src/components/transaction_details.rs new file mode 100644 index 00000000..c82f7d80 --- /dev/null +++ b/lez/explorer_service/src/components/transaction_details.rs @@ -0,0 +1,150 @@ +use indexer_service_protocol::{ + PrivacyPreservingMessage, PrivacyPreservingTransaction, ProgramDeploymentMessage, + ProgramDeploymentTransaction, PublicMessage, PublicTransaction, WitnessSet, +}; +use leptos::prelude::*; + +use super::AccountNonceList; + +/// Public transaction details component +#[component] +pub fn PublicTxDetails(tx: PublicTransaction) -> impl IntoView { + let PublicTransaction { + hash: _, + message, + witness_set, + } = tx; + let PublicMessage { + program_id, + account_ids, + nonces, + instruction_data, + } = message; + let WitnessSet { + signatures_and_public_keys, + proof, + } = witness_set; + + let program_id_str = program_id.to_string(); + let proof_len = proof.map_or(0, |p| p.0.len()); + let signatures_count = signatures_and_public_keys.len(); + + view! { +
+

"Public Transaction Details"

+
+
+ "Program ID:" + {program_id_str} +
+
+ "Instruction Data:" + + {format!("{} u32 values", instruction_data.len())} + +
+
+ "Proof Size:" + {format!("{proof_len} bytes")} +
+
+ "Signatures:" + {signatures_count.to_string()} +
+
+ +

"Accounts"

+ +
+ } +} + +/// Privacy-preserving transaction details component +#[component] +pub fn PrivacyPreservingTxDetails(tx: PrivacyPreservingTransaction) -> impl IntoView { + let PrivacyPreservingTransaction { + hash: _, + message, + witness_set, + } = tx; + let PrivacyPreservingMessage { + public_account_ids, + nonces, + public_post_states: _, + encrypted_private_post_states, + new_commitments, + new_nullifiers, + block_validity_window, + timestamp_validity_window, + } = message; + let WitnessSet { + signatures_and_public_keys: _, + proof, + } = witness_set; + let proof_len = proof.map_or(0, |p| p.0.len()); + + view! { +
+

"Privacy-Preserving Transaction Details"

+
+
+ "Public Accounts:" + + {public_account_ids.len().to_string()} + +
+
+ "New Commitments:" + {new_commitments.len().to_string()} +
+
+ "Nullifiers:" + {new_nullifiers.len().to_string()} +
+
+ "Encrypted States:" + + {encrypted_private_post_states.len().to_string()} + +
+
+ "Proof Size:" + {format!("{proof_len} bytes")} +
+
+ "Block Validity Window:" + {block_validity_window.to_string()} +
+
+ "Timestamp Validity Window:" + {timestamp_validity_window.to_string()} +
+
+ +

"Public Accounts"

+ +
+ } +} + +/// Program deployment transaction details component +#[component] +pub fn ProgramDeploymentTxDetails(tx: ProgramDeploymentTransaction) -> impl IntoView { + let ProgramDeploymentTransaction { hash: _, message } = tx; + let ProgramDeploymentMessage { bytecode } = message; + + let bytecode_len = bytecode.len(); + view! { +
+

"Program Deployment Transaction Details"

+
+
+ "Bytecode Size:" + + {format!("{bytecode_len} bytes")} + +
+
+
+ } +} diff --git a/lez/explorer_service/src/pages/main_page.rs b/lez/explorer_service/src/pages/main_page.rs index 7e26e794..831182ce 100644 --- a/lez/explorer_service/src/pages/main_page.rs +++ b/lez/explorer_service/src/pages/main_page.rs @@ -6,8 +6,8 @@ use leptos_router::{ use web_sys::SubmitEvent; use crate::{ - api::{self, SearchResults}, - components::{AccountPreview, BlockPreview, TransactionPreview}, + api, + components::{BlockPreview, SearchResultsView}, }; const RECENT_BLOCKS_LIMIT: u64 = 10; @@ -138,93 +138,8 @@ pub fn MainPage() -> impl IntoView { .get() .and_then(|opt_results| opt_results) .map(|results| { - let SearchResults { - blocks, - transactions, - accounts, - } = results; - let has_results = !blocks.is_empty() - || !transactions.is_empty() - || !accounts.is_empty(); - view! { -
-

"Search Results"

- {if has_results { - view! { -
- {if blocks.is_empty() { - ().into_any() - } else { - view! { -
-

"Blocks"

-
- {blocks - .into_iter() - .map(|block| { - view! { } - }) - .collect::>()} -
-
- } - .into_any() - }} - - {if transactions.is_empty() { - ().into_any() - } else { - view! { -
-

"Transactions"

-
- {transactions - .into_iter() - .map(|tx| { - view! { } - }) - .collect::>()} -
-
- } - .into_any() - }} - - {if accounts.is_empty() { - ().into_any() - } else { - view! { -
-

"Accounts"

-
- {accounts - .into_iter() - .map(|(id, account)| { - view! { - - } - }) - .collect::>()} -
-
- } - .into_any() - }} - -
- } - .into_any() - } else { - view! {
"No results found"
} - .into_any() - }} -
- } - .into_any() - }) + view! { }.into_any() + }) }} diff --git a/lez/explorer_service/src/pages/transaction_page.rs b/lez/explorer_service/src/pages/transaction_page.rs index 0a3fc8e2..a99c269f 100644 --- a/lez/explorer_service/src/pages/transaction_page.rs +++ b/lez/explorer_service/src/pages/transaction_page.rs @@ -1,14 +1,13 @@ use std::str::FromStr as _; -use indexer_service_protocol::{ - HashType, PrivacyPreservingMessage, PrivacyPreservingTransaction, ProgramDeploymentMessage, - ProgramDeploymentTransaction, PublicMessage, PublicTransaction, Transaction, WitnessSet, -}; -use itertools::{EitherOrBoth, Itertools as _}; +use indexer_service_protocol::{HashType, Transaction}; use leptos::prelude::*; -use leptos_router::{components::A, hooks::use_params_map}; +use leptos_router::hooks::use_params_map; -use crate::api; +use crate::{ + api, + components::{PrivacyPreservingTxDetails, ProgramDeploymentTxDetails, PublicTxDetails}, +}; /// Transaction page component #[component] @@ -66,244 +65,21 @@ pub fn TransactionPage() -> impl IntoView { { match tx { - Transaction::Public(ptx) => { - let PublicTransaction { - hash: _, - message, - witness_set, - } = ptx; - let PublicMessage { - program_id, - account_ids, - nonces, - instruction_data, - } = message; - let WitnessSet { - signatures_and_public_keys, - proof, - } = witness_set; + Transaction::Public(ptx) => { + view! { }.into_any() + } + Transaction::PrivacyPreserving(pptx) => { + view! { }.into_any() + } + Transaction::ProgramDeployment(pdtx) => { + view! { }.into_any() + } + } + } - let program_id_str = program_id.to_string(); - let proof_len = proof.map_or(0, |p| p.0.len()); - let signatures_count = signatures_and_public_keys.len(); - - view! { -
-

"Public Transaction Details"

-
-
- "Program ID:" - {program_id_str} -
-
- "Instruction Data:" - - {format!("{} u32 values", instruction_data.len())} - -
-
- "Proof Size:" - {format!("{proof_len} bytes")} -
-
- "Signatures:" - {signatures_count.to_string()} -
-
- -

"Accounts"

-
- {account_ids - .into_iter() - .zip_longest(nonces.into_iter()) - .map(|maybe_pair| { - match maybe_pair { - EitherOrBoth::Both(account_id, nonce) => { - let account_id_str = account_id.to_string(); - view! { - - } - } - EitherOrBoth::Left(account_id) => { - let account_id_str = account_id.to_string(); - view! { - - } - } - EitherOrBoth::Right(_) => { - view! { - - } - } - } - }) - .collect::>()} -
-
- } - .into_any() + } - Transaction::PrivacyPreserving(pptx) => { - let PrivacyPreservingTransaction { - hash: _, - message, - witness_set, - } = pptx; - let PrivacyPreservingMessage { - public_account_ids, - nonces, - public_post_states: _, - encrypted_private_post_states, - new_commitments, - new_nullifiers, - block_validity_window, - timestamp_validity_window, - } = message; - let WitnessSet { - signatures_and_public_keys: _, - proof, - } = witness_set; - let proof_len = proof.map_or(0, |p| p.0.len()); - view! { -
-

"Privacy-Preserving Transaction Details"

-
-
- "Public Accounts:" - - {public_account_ids.len().to_string()} - -
-
- "New Commitments:" - {new_commitments.len().to_string()} -
-
- "Nullifiers:" - {new_nullifiers.len().to_string()} -
-
- "Encrypted States:" - - {encrypted_private_post_states.len().to_string()} - -
-
- "Proof Size:" - {format!("{proof_len} bytes")} -
-
- "Block Validity Window:" - {block_validity_window.to_string()} -
-
- "Timestamp Validity Window:" - {timestamp_validity_window.to_string()} -
-
- -

"Public Accounts"

-
- {public_account_ids - .into_iter() - .zip_longest(nonces.into_iter()) - .map(|maybe_pair| { - match maybe_pair { - EitherOrBoth::Both(account_id, nonce) => { - let account_id_str = account_id.to_string(); - view! { - - } - } - EitherOrBoth::Left(account_id) => { - let account_id_str = account_id.to_string(); - view! { - - } - } - EitherOrBoth::Right(_) => { - view! { - - } - } - } - }) - .collect::>()} -
-
- } - .into_any() - } - Transaction::ProgramDeployment(pdtx) => { - let ProgramDeploymentTransaction { - hash: _, - message, - } = pdtx; - let ProgramDeploymentMessage { bytecode } = message; - - let bytecode_len = bytecode.len(); - view! { -
-

"Program Deployment Transaction Details"

-
-
- "Bytecode Size:" - - {format!("{bytecode_len} bytes")} - -
-
-
- } - .into_any() - } - }} - - - } - .into_any() + .into_any() } Err(e) => { view! { diff --git a/lez/indexer/core/src/block_store.rs b/lez/indexer/core/src/block_store.rs index f00c94c5..c67148bd 100644 --- a/lez/indexer/core/src/block_store.rs +++ b/lez/indexer/core/src/block_store.rs @@ -211,6 +211,60 @@ mod tests { use super::*; + struct TestFixture { + storage: IndexerStore, + from: AccountId, + to: AccountId, + _home: tempfile::TempDir, + } + + #[expect( + clippy::arithmetic_side_effects, + reason = "test helper with bounded inputs" + )] + async fn store_with_transfer_blocks( + block_count: u64, + prev_hash: Option, + ) -> TestFixture { + let home = tempdir().unwrap(); + let storage = IndexerStore::open_db(home.path()).unwrap(); + + let initial_accounts = initial_pub_accounts_private_keys(); + let from = initial_accounts[0].account_id; + let to = initial_accounts[1].account_id; + let sign_key = initial_accounts[0].pub_sign_key.clone(); + + let mut prev_hash = prev_hash; + for i in 0..block_count { + let tx = common::test_utils::create_transaction_native_token_transfer( + from, + u128::from(i), + to, + 10, + &sign_key, + ); + let block_id = i + 1; + + let next_block = common::test_utils::produce_dummy_block(block_id, prev_hash, vec![tx]); + prev_hash = Some(next_block.header.hash); + + storage + .put_block( + next_block, + HeaderId::from([u8::try_from(i + 1).unwrap(); 32]), + ) + .await + .unwrap(); + } + + TestFixture { + storage, + from, + to, + _home: home, + } + } + #[test] fn correct_startup() { let home = tempdir().unwrap(); @@ -225,7 +279,6 @@ mod tests { #[tokio::test] async fn state_transition() { let home = tempdir().unwrap(); - let storage = IndexerStore::open_db(home.as_ref()).unwrap(); let initial_accounts = initial_pub_accounts_private_keys(); @@ -233,7 +286,6 @@ mod tests { let to = initial_accounts[1].account_id; let sign_key = initial_accounts[0].pub_sign_key.clone(); - // Submit genesis block let clock_tx = LeeTransaction::Public(clock_invocation(0)); let genesis_block_data = HashableBlockData { block_id: 1, @@ -249,15 +301,13 @@ mod tests { .await .unwrap(); - for i in 0..10 { + for i in 0..10_u128 { let tx = common::test_utils::create_transaction_native_token_transfer( from, i, to, 10, &sign_key, ); let block_id = u64::try_from(i + 1).unwrap(); - let next_block = common::test_utils::produce_dummy_block(block_id, prev_hash, vec![tx]); prev_hash = Some(next_block.header.hash); - storage .put_block( next_block, @@ -276,48 +326,23 @@ mod tests { #[tokio::test] async fn account_state_at_block() { - let home = tempdir().unwrap(); + let TestFixture { + storage, + from, + to, + _home, + } = store_with_transfer_blocks(10, None).await; - let storage = IndexerStore::open_db(home.as_ref()).unwrap(); - - let mut prev_hash = None; - - let initial_accounts = initial_pub_accounts_private_keys(); - let from = initial_accounts[0].account_id; - let to = initial_accounts[1].account_id; - let sign_key = initial_accounts[0].pub_sign_key.clone(); - - for i in 0..10 { - let tx = common::test_utils::create_transaction_native_token_transfer( - from, i, to, 10, &sign_key, - ); - let block_id = u64::try_from(i + 1).unwrap(); - - let next_block = common::test_utils::produce_dummy_block(block_id, prev_hash, vec![tx]); - prev_hash = Some(next_block.header.hash); - - storage - .put_block( - next_block, - HeaderId::from([u8::try_from(i + 1).unwrap(); 32]), - ) - .await - .unwrap(); - } - - // Genesis block: no transfers applied yet. let acc1_at_1 = storage.account_state_at_block(&from, 1).unwrap(); let acc2_at_1 = storage.account_state_at_block(&to, 1).unwrap(); assert_eq!(acc1_at_1.balance, 9990); assert_eq!(acc2_at_1.balance, 20010); - // After block 5: 4 transfers of 10 applied (one each in blocks 2..=5). let acc1_at_5 = storage.account_state_at_block(&from, 5).unwrap(); let acc2_at_5 = storage.account_state_at_block(&to, 5).unwrap(); assert_eq!(acc1_at_5.balance, 9950); assert_eq!(acc2_at_5.balance, 20050); - // After final block 9: 8 transfers applied; should match current state. let acc1_at_9 = storage.account_state_at_block(&from, 9).unwrap(); let acc2_at_9 = storage.account_state_at_block(&to, 9).unwrap(); assert_eq!(acc1_at_9.balance, 9910); diff --git a/lez/indexer/service/src/mock_service.rs b/lez/indexer/service/src/mock_service.rs index d9ab9484..7bf6c528 100644 --- a/lez/indexer/service/src/mock_service.rs +++ b/lez/indexer/service/src/mock_service.rs @@ -330,6 +330,76 @@ impl indexer_service_rpc::RpcServer for MockIndexerService { } } +fn mock_public_tx( + tx_hash: HashType, + block_id: BlockId, + tx_idx: u64, + account_ids: &[AccountId], +) -> Transaction { + Transaction::Public(PublicTransaction { + hash: tx_hash, + message: PublicMessage { + program_id: ProgramId([1_u32; 8]), + account_ids: vec![ + account_ids[tx_idx as usize % account_ids.len()], + account_ids[(tx_idx as usize + 1) % account_ids.len()], + ], + nonces: vec![block_id as u128, (block_id + 1) as u128], + instruction_data: vec![1, 2, 3, 4], + }, + witness_set: WitnessSet { + signatures_and_public_keys: vec![], + proof: None, + }, + }) +} + +fn mock_privacy_preserving_tx( + tx_hash: HashType, + block_id: BlockId, + tx_idx: u64, + account_ids: &[AccountId], +) -> Transaction { + Transaction::PrivacyPreserving(PrivacyPreservingTransaction { + hash: tx_hash, + message: PrivacyPreservingMessage { + public_account_ids: vec![account_ids[tx_idx as usize % account_ids.len()]], + nonces: vec![block_id as u128], + public_post_states: vec![Account { + program_owner: ProgramId([1_u32; 8]), + balance: 500, + data: Data(vec![0xdd, 0xee]), + nonce: block_id as u128, + }], + encrypted_private_post_states: vec![EncryptedAccountData { + ciphertext: indexer_service_protocol::Ciphertext(vec![0x01, 0x02, 0x03, 0x04]), + epk: indexer_service_protocol::EphemeralPublicKey(vec![0xaa; 32]), + view_tag: 42, + }], + new_commitments: vec![Commitment([block_id as u8; 32])], + new_nullifiers: vec![( + indexer_service_protocol::Nullifier([tx_idx as u8; 32]), + CommitmentSetDigest([0xff; 32]), + )], + block_validity_window: ValidityWindow((None, None)), + timestamp_validity_window: ValidityWindow((None, None)), + }, + witness_set: WitnessSet { + signatures_and_public_keys: vec![], + proof: Some(indexer_service_protocol::Proof(vec![0; 32])), + }, + }) +} + +fn mock_program_deployment_tx(tx_hash: HashType) -> Transaction { + Transaction::ProgramDeployment(ProgramDeploymentTransaction { + hash: tx_hash, + message: ProgramDeploymentMessage { + bytecode: vec![0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00], + }, + }) +} + fn build_mock_block( block_id: BlockId, prev_hash: HashType, @@ -344,7 +414,6 @@ fn build_mock_block( HashType(hash) }; - // Create 2-4 transactions per block (mix of Public, PrivacyPreserving, and ProgramDeployment) let num_txs = 2 + (block_id % 3); let mut block_transactions = Vec::new(); @@ -356,65 +425,10 @@ fn build_mock_block( HashType(hash) }; - // Vary transaction types: Public, PrivacyPreserving, or ProgramDeployment let tx = match (block_id + tx_idx) % 5 { - // Public transactions (most common) - 0 | 1 => Transaction::Public(PublicTransaction { - hash: tx_hash, - message: PublicMessage { - program_id: ProgramId([1_u32; 8]), - account_ids: vec![ - account_ids[tx_idx as usize % account_ids.len()], - account_ids[(tx_idx as usize + 1) % account_ids.len()], - ], - nonces: vec![block_id as u128, (block_id + 1) as u128], - instruction_data: vec![1, 2, 3, 4], - }, - witness_set: WitnessSet { - signatures_and_public_keys: vec![], - proof: None, - }, - }), - // PrivacyPreserving transactions - 2 | 3 => Transaction::PrivacyPreserving(PrivacyPreservingTransaction { - hash: tx_hash, - message: PrivacyPreservingMessage { - public_account_ids: vec![account_ids[tx_idx as usize % account_ids.len()]], - nonces: vec![block_id as u128], - public_post_states: vec![Account { - program_owner: ProgramId([1_u32; 8]), - balance: 500, - data: Data(vec![0xdd, 0xee]), - nonce: block_id as u128, - }], - encrypted_private_post_states: vec![EncryptedAccountData { - ciphertext: indexer_service_protocol::Ciphertext(vec![ - 0x01, 0x02, 0x03, 0x04, - ]), - epk: indexer_service_protocol::EphemeralPublicKey(vec![0xaa; 32]), - view_tag: 42, - }], - new_commitments: vec![Commitment([block_id as u8; 32])], - new_nullifiers: vec![( - indexer_service_protocol::Nullifier([tx_idx as u8; 32]), - CommitmentSetDigest([0xff; 32]), - )], - block_validity_window: ValidityWindow((None, None)), - timestamp_validity_window: ValidityWindow((None, None)), - }, - witness_set: WitnessSet { - signatures_and_public_keys: vec![], - proof: Some(indexer_service_protocol::Proof(vec![0; 32])), - }, - }), - // ProgramDeployment transactions (rare) - _ => Transaction::ProgramDeployment(ProgramDeploymentTransaction { - hash: tx_hash, - message: ProgramDeploymentMessage { - bytecode: vec![0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00], /* WASM magic - * number */ - }, - }), + 0 | 1 => mock_public_tx(tx_hash, block_id, tx_idx, account_ids), + 2 | 3 => mock_privacy_preserving_tx(tx_hash, block_id, tx_idx, account_ids), + _ => mock_program_deployment_tx(tx_hash), }; block_transactions.push(tx); diff --git a/lez/keycard_wallet/src/lib.rs b/lez/keycard_wallet/src/lib.rs index 73486392..a3ecd140 100644 --- a/lez/keycard_wallet/src/lib.rs +++ b/lez/keycard_wallet/src/lib.rs @@ -134,10 +134,6 @@ impl KeycardWallet { }) } - #[expect( - clippy::arithmetic_side_effects, - reason = "64 - s_stripped.len() is safe: s_stripped.len() ≤ 31 because py_signature.len() is in [32, 63]" - )] pub fn sign_message_for_path( &self, py: Python, @@ -150,33 +146,9 @@ impl KeycardWallet { .call_method1("sign_message_for_path", (message, path))? .extract()?; - // The keycard Python library strips leading zeros from S when S < 2^(8k) for some k. - // Left-pad S back to 32 bytes so the full signature is always 64 bytes (R || S). - let py_signature = if py_signature.len() < 64 { - if py_signature.len() < 32 { - return Err(PyErr::new::(format!( - "signature from keycard too short: {} bytes", - py_signature.len() - ))); - } - let s_stripped = &py_signature[32..]; - let mut padded = [0_u8; 64]; - padded[..32].copy_from_slice(&py_signature[..32]); - padded[(64 - s_stripped.len())..].copy_from_slice(s_stripped); - padded.to_vec() - } else { - py_signature + let sig = Signature { + value: normalize_keycard_signature(py_signature)?, }; - - let signature: [u8; 64] = py_signature.try_into().map_err(|vec: Vec| { - PyErr::new::(format!( - "Invalid signature length: expected 64 bytes, got {} (bytes: {:02x?})", - vec.len(), - vec - )) - })?; - - let sig = Signature { value: signature }; let pub_key = self.get_public_key_for_path(py, path)?; if !sig.is_valid_for(message, &pub_key) { return Err(PyErr::new::( @@ -224,32 +196,8 @@ impl KeycardWallet { .call_method1("get_private_keys_for_path", (path,))? .extract()?; - let raw_nsk = Zeroizing::new(raw_nsk); - let raw_vsk = Zeroizing::new(raw_vsk); - - let nsk = { - if raw_nsk.len() != 32 { - return Err(PyErr::new::(format!( - "expected 32-byte NSK from keycard, got {} bytes", - raw_nsk.len() - ))); - } - let mut arr = Zeroizing::new([0_u8; 32]); - arr.copy_from_slice(&raw_nsk); - arr - }; - - let vsk = { - if raw_vsk.len() != 64 { - return Err(PyErr::new::(format!( - "expected 64-byte VSK from keycard, got {} bytes", - raw_vsk.len() - ))); - } - let mut arr = Zeroizing::new([0_u8; 64]); - arr.copy_from_slice(&raw_vsk); - arr - }; + let nsk = zeroizing_fixed_bytes::<32>("nullifier secret key", Zeroizing::new(raw_nsk))?; + let vsk = zeroizing_fixed_bytes::<64>("view secret key", Zeroizing::new(raw_vsk))?; Ok((nsk, vsk)) } @@ -269,6 +217,55 @@ impl KeycardWallet { } } +/// The keycard Python library strips leading zeros from S when S < 2^(8k) for some k. +/// Left-pad S back to 32 bytes so the full signature is always 64 bytes (R || S). +#[expect( + clippy::arithmetic_side_effects, + reason = "64 - s_stripped.len() is safe: s_stripped.len() ≤ 31 because py_signature.len() is in [32, 63]" +)] +fn normalize_keycard_signature(py_signature: Vec) -> PyResult<[u8; 64]> { + if py_signature.len() < 64 { + if py_signature.len() < 32 { + return Err(PyErr::new::(format!( + "signature from keycard too short: {} bytes", + py_signature.len() + ))); + } + let s_stripped = &py_signature[32..]; + let mut padded = [0_u8; 64]; + padded[..32].copy_from_slice(&py_signature[..32]); + padded[(64 - s_stripped.len())..].copy_from_slice(s_stripped); + Ok(padded) + } else { + py_signature.try_into().map_err(|vec: Vec| { + PyErr::new::(format!( + "Invalid signature length: expected 64 bytes, got {} (bytes: {:02x?})", + vec.len(), + vec + )) + }) + } +} + +#[expect( + clippy::needless_pass_by_value, + reason = "Zeroizing> is consumed to ensure the source is zeroed on drop" +)] +fn zeroizing_fixed_bytes( + label: &str, + raw: Zeroizing>, +) -> PyResult> { + if raw.len() != N { + return Err(PyErr::new::(format!( + "expected {N}-byte {label} from keycard, got {} bytes", + raw.len() + ))); + } + let mut arr = Zeroizing::new([0_u8; N]); + arr.copy_from_slice(&raw); + Ok(arr) +} + fn pairing_file_path() -> Option { let home = std::env::var("LEE_WALLET_HOME_DIR") .map(PathBuf::from) diff --git a/lez/keycard_wallet/src/python_path.rs b/lez/keycard_wallet/src/python_path.rs index 99ed936e..61196ad5 100644 --- a/lez/keycard_wallet/src/python_path.rs +++ b/lez/keycard_wallet/src/python_path.rs @@ -2,8 +2,7 @@ use std::{env, path::PathBuf}; use pyo3::{prelude::*, types::PyList}; -/// Adds the project's `python/` directory and venv site-packages to Python's sys.path. -pub fn add_python_path(py: Python<'_>) -> PyResult<()> { +fn collect_python_paths() -> Vec { let current_dir = env::current_dir().expect("Failed to get current working directory"); let python_base = env::var("VIRTUAL_ENV") @@ -11,7 +10,7 @@ pub fn add_python_path(py: Python<'_>) -> PyResult<()> { .and_then(|v| PathBuf::from(v).parent().map(PathBuf::from)) .unwrap_or_else(|| current_dir.clone()); - let mut paths_to_add: Vec = vec![ + let mut paths = vec![ python_base .join("lez") .join("keycard_wallet") @@ -23,24 +22,28 @@ pub fn add_python_path(py: Python<'_>) -> PyResult<()> { .join("keycard-py"), ]; - // If a virtualenv is active, add its site-packages so that dependencies - // installed in the venv (e.g. smartcard, ecdsa) are importable by the - // pyo3 embedded interpreter, which does not inherit sys.path from the - // shell's `python3` executable. + // pyo3's embedded interpreter does not inherit sys.path from the shell, + // so venv site-packages must be added explicitly. if let Ok(venv) = env::var("VIRTUAL_ENV") { let lib = PathBuf::from(&venv).join("lib"); if let Ok(entries) = std::fs::read_dir(&lib) { for entry in entries.flatten() { let site_packages = entry.path().join("site-packages"); if site_packages.exists() { - paths_to_add.push(site_packages); + paths.push(site_packages); } } } } - // Sanity check — warns early if a path doesn't exist - for path in &paths_to_add { + paths +} + +/// Adds the project's `python/` directory and venv site-packages to Python's sys.path. +pub fn add_python_path(py: Python<'_>) -> PyResult<()> { + let paths = collect_python_paths(); + + for path in &paths { if !path.exists() { log::info!("Warning: Python path does not exist: {}", path.display()); } @@ -50,10 +53,9 @@ pub fn add_python_path(py: Python<'_>) -> PyResult<()> { let binding = sys.getattr("path")?; let sys_path = binding.cast::()?; - for path in &paths_to_add { + for path in &paths { let path_str = path.to_str().expect("Invalid path"); - // Avoid duplicating the path let already_present = sys_path .iter() .any(|p| p.extract::<&str>().is_ok_and(|s| s == path_str)); diff --git a/lez/sequencer/core/src/lib.rs b/lez/sequencer/core/src/lib.rs index 35c0c33b..75dcc1a9 100644 --- a/lez/sequencer/core/src/lib.rs +++ b/lez/sequencer/core/src/lib.rs @@ -36,6 +36,7 @@ pub mod config; pub mod mock; /// The origin of a transaction. +#[derive(Clone, Copy)] pub enum TransactionOrigin { /// Basic transactions submitted by users via RPC. User, @@ -69,24 +70,17 @@ impl SequencerCore { /// assumed to represent the correct latest state consistent with Bedrock-finalized data. /// If no database is found, the sequencer performs a fresh start from genesis, /// initializing its state with the accounts defined in the configuration file. - pub async fn start_from_config( - config: SequencerConfig, - ) -> (Self, MemPoolHandle<(TransactionOrigin, LeeTransaction)>) { + fn open_or_create_store(config: &SequencerConfig) -> (SequencerStore, lee::V03State, Block) { let signing_key = lee::PrivateKey::try_new(config.signing_key).unwrap(); - - let bedrock_signing_key = - load_or_create_signing_key(&config.home.join("bedrock_signing_key")) - .expect("Failed to load or create bedrock signing key"); - let db_path = config.home.join("rocksdb"); - let (store, state, genesis_block) = if db_path.exists() { - let store = - SequencerStore::open_db(&db_path, signing_key.clone()).unwrap_or_else(|err| { - panic!( - "Failed to open database at {} with error: {err}", - db_path.display() - ) - }); + + if db_path.exists() { + let store = SequencerStore::open_db(&db_path, signing_key).unwrap_or_else(|err| { + panic!( + "Failed to open database at {} with error: {err}", + db_path.display() + ) + }); let state = store .get_lee_state() .expect("Failed to read state from store"); @@ -101,7 +95,7 @@ impl SequencerCore { db_path.display() ); - let (genesis_state, genesis_txs) = build_genesis_state(&config); + let (genesis_state, genesis_txs) = build_genesis_state(config); let hashable_data = HashableBlockData { block_id: GENESIS_BLOCK_ID, @@ -120,7 +114,17 @@ impl SequencerCore { .expect("Failed to create database with genesis block"); (store, genesis_state, genesis_block) - }; + } + } + + pub async fn start_from_config( + config: SequencerConfig, + ) -> (Self, MemPoolHandle<(TransactionOrigin, LeeTransaction)>) { + let bedrock_signing_key = + load_or_create_signing_key(&config.home.join("bedrock_signing_key")) + .expect("Failed to load or create bedrock signing key"); + + let (store, state, genesis_block) = Self::open_or_create_store(&config); let latest_block_meta = store .latest_block_meta() @@ -331,8 +335,60 @@ impl SequencerCore { Ok(self.chain_height) } - /// Builds a new block from transactions in the mempool. - /// Does NOT publish or store the block — the caller is responsible for that. + /// Validates and applies a single mempool transaction to the current state. + /// Returns `Ok(true)` if the transaction was valid and applied, `Ok(false)` if + /// it was skipped due to validation failure. + fn apply_mempool_transaction( + &mut self, + origin: TransactionOrigin, + tx: &LeeTransaction, + block_height: u64, + timestamp: u64, + deposit_event_ids: &mut Vec, + withdrawals: &mut Vec, + ) -> Result { + let tx_hash = tx.hash(); + match origin { + TransactionOrigin::User => { + let validated_diff = match tx.validate_on_state( + &self.state, + block_height, + timestamp, + ) { + Ok(diff) => diff, + Err(err) => { + error!( + "Transaction with hash {tx_hash} failed execution check with error: {err:#?}, skipping it", + ); + return Ok(false); + } + }; + + if let Some(withdraw_data) = extract_bridge_withdraw_data(tx) { + withdrawals.push(withdraw_data); + } + + self.state.apply_state_diff(validated_diff); + } + TransactionOrigin::Sequencer => { + let LeeTransaction::Public(public_tx) = tx else { + panic!("Sequencer may only generate Public transactions, found {tx:#?}"); + }; + + if let Some(deposit_op_id) = extract_bridge_deposit_id(tx) { + deposit_event_ids.push(deposit_op_id); + } + + self.state + .transition_from_public_transaction(public_tx, block_height, timestamp) + .context("Failed to execute sequencer-generated transaction")?; + } + } + + info!("Validated transaction with hash {tx_hash}, including it in block"); + Ok(true) + } + fn build_block_from_mempool(&mut self) -> Result { let now = Instant::now(); @@ -353,14 +409,12 @@ impl SequencerCore { let new_block_timestamp = u64::try_from(chrono::Utc::now().timestamp_millis()) .expect("Timestamp must be positive"); - // Pre-create the mandatory clock tx so its size is included in the block size check. let clock_tx = clock_invocation(new_block_timestamp); let clock_lee_tx = LeeTransaction::Public(clock_tx.clone()); while let Some((origin, tx)) = self.mempool.pop() { let tx_hash = tx.hash(); - // Check if block size exceeds limit (including the mandatory clock tx). let temp_valid_transactions = [ valid_transactions.as_slice(), std::slice::from_ref(&tx), @@ -379,66 +433,30 @@ impl SequencerCore { .len(); if block_size > max_block_size { - // Block would exceed size limit, remove last transaction and push back warn!( "Transaction with hash {tx_hash} deferred to next block: \ block size {block_size} bytes would exceed limit of {max_block_size} bytes", ); - self.mempool.push_front((origin, tx)); break; } - match origin { - TransactionOrigin::User => { - let validated_diff = match tx.validate_on_state( - &self.state, - new_block_height, - new_block_timestamp, - ) { - Ok(diff) => diff, - Err(err) => { - error!( - "Transaction with hash {tx_hash} failed execution check with error: {err:#?}, skipping it", - ); - continue; - } - }; - - if let Some(withdraw_data) = extract_bridge_withdraw_data(&tx) { - withdrawals.push(withdraw_data); - } - - self.state.apply_state_diff(validated_diff); - } - TransactionOrigin::Sequencer => { - let LeeTransaction::Public(public_tx) = &tx else { - panic!("Sequencer may only generate Public transactions, found {tx:#?}"); - }; - - if let Some(deposit_op_id) = extract_bridge_deposit_id(&tx) { - deposit_event_ids.push(deposit_op_id); - } - - self.state - .transition_from_public_transaction( - public_tx, - new_block_height, - new_block_timestamp, - ) - .context("Failed to execute sequencer-generated transaction")?; - } + if self.apply_mempool_transaction( + origin, + &tx, + new_block_height, + new_block_timestamp, + &mut deposit_event_ids, + &mut withdrawals, + )? { + valid_transactions.push(tx); } - valid_transactions.push(tx); - - info!("Validated transaction with hash {tx_hash}, including it in block"); if valid_transactions.len() >= self.sequencer_config.max_num_tx_in_block { break; } } - // Append the Clock Program invocation as the mandatory last transaction. self.state .transition_from_public_transaction(&clock_tx, new_block_height, new_block_timestamp) .context("Clock transaction failed. Aborting block production.")?; @@ -803,1248 +821,4 @@ fn load_or_create_signing_key(path: &Path) -> Result { #[cfg(test)] #[cfg(feature = "mock")] -mod tests { - #![expect(clippy::shadow_unrelated, reason = "We don't care about it in tests")] - - use std::{pin::pin, time::Duration}; - - use common::{ - HashType, - block::HashableBlockData, - test_utils::sequencer_sign_key_for_testing, - transaction::{LeeTransaction, clock_invocation}, - }; - use key_protocol::key_management::KeyChain; - use lee::{ - Account, AccountId, Data, EphemeralPublicKey, PrivacyPreservingTransaction, PrivateKey, - PublicKey, PublicTransaction, SharedSecretKey, V03State, - error::LeeError, - execute_and_prove, - privacy_preserving_transaction::{Message, circuit::ProgramWithDependencies}, - program::Program, - }; - use lee_core::{ - Commitment, EncryptedAccountData, InputAccountIdentity, Nullifier, - account::{AccountWithMetadata, Nonce}, - program::PdaSeed, - }; - use logos_blockchain_core::mantle::ops::channel::ChannelId; - use mempool::MemPoolHandle; - use storage::sequencer::sequencer_cells::PendingDepositEventRecord; - use tempfile::tempdir; - use testnet_initial_state::{initial_pub_accounts_private_keys, initial_public_user_accounts}; - - use crate::{ - TransactionOrigin, - block_store::SequencerStore, - build_genesis_state, - config::{BedrockConfig, SequencerConfig}, - mock::SequencerCoreWithMockClients, - }; - - #[derive(borsh::BorshSerialize)] - struct DepositMetadataForEncoding { - recipient_id: lee::AccountId, - } - - fn setup_sequencer_config() -> SequencerConfig { - let tempdir = tempfile::tempdir().unwrap(); - let home = tempdir.path().to_path_buf(); - - SequencerConfig { - home, - max_num_tx_in_block: 10, - max_block_size: bytesize::ByteSize::mib(1), - mempool_max_size: 10000, - block_create_timeout: Duration::from_secs(1), - signing_key: *sequencer_sign_key_for_testing().value(), - bedrock_config: BedrockConfig { - channel_id: ChannelId::from([0; 32]), - node_url: "http://not-used-in-unit-tests".parse().unwrap(), - auth: None, - }, - retry_pending_blocks_timeout: Duration::from_mins(4), - genesis: vec![], - } - } - - fn create_signing_key_for_account1() -> lee::PrivateKey { - initial_pub_accounts_private_keys()[0].pub_sign_key.clone() - } - - fn create_signing_key_for_account2() -> lee::PrivateKey { - initial_pub_accounts_private_keys()[1].pub_sign_key.clone() - } - - async fn common_setup() -> ( - SequencerCoreWithMockClients, - MemPoolHandle<(TransactionOrigin, LeeTransaction)>, - ) { - let config = setup_sequencer_config(); - common_setup_with_config(config).await - } - - async fn common_setup_with_config( - config: SequencerConfig, - ) -> ( - SequencerCoreWithMockClients, - MemPoolHandle<(TransactionOrigin, LeeTransaction)>, - ) { - let (mut sequencer, mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config).await; - - let tx = common::test_utils::produce_dummy_empty_transaction(); - mempool_handle - .push((TransactionOrigin::User, tx)) - .await - .unwrap(); - - sequencer.produce_new_block().await.unwrap(); - - (sequencer, mempool_handle) - } - - fn tx_is_bridge_deposit( - tx: &LeeTransaction, - deposit_op_id: [u8; 32], - expected_amount: u64, - ) -> bool { - let LeeTransaction::Public(public_tx) = tx else { - return false; - }; - - if public_tx.message.program_id != programs::bridge().id() { - return false; - } - - let instruction: bridge_core::Instruction = - match risc0_zkvm::serde::from_slice(&public_tx.message.instruction_data) { - Ok(instruction) => instruction, - Err(_err) => return false, - }; - - matches!( - instruction, - bridge_core::Instruction::Deposit { - l1_deposit_op_id, - amount, - .. - } if l1_deposit_op_id == deposit_op_id && amount == expected_amount - ) - } - - #[tokio::test] - async fn start_from_config() { - let config = setup_sequencer_config(); - let (sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config.clone()).await; - - assert_eq!(sequencer.chain_height, 1); - assert_eq!(sequencer.sequencer_config.max_num_tx_in_block, 10); - - let acc1_account_id = initial_public_user_accounts()[0].account_id; - let acc2_account_id = initial_public_user_accounts()[1].account_id; - - let balance_acc_1 = sequencer.state.get_account_by_id(acc1_account_id).balance; - let balance_acc_2 = sequencer.state.get_account_by_id(acc2_account_id).balance; - - assert_eq!(10000, balance_acc_1); - assert_eq!(20000, balance_acc_2); - } - - #[tokio::test] - async fn start_from_config_opens_existing_db_if_it_exists() { - let config = setup_sequencer_config(); - let temp_dir = tempdir().unwrap(); - let mut config = config; - config.home = temp_dir.path().to_path_buf(); - - let signing_key = lee::PrivateKey::try_new(config.signing_key).unwrap(); - let (genesis_state, genesis_txs) = build_genesis_state(&config); - let genesis_hashable_data = HashableBlockData { - block_id: 1, - transactions: genesis_txs, - prev_block_hash: HashType([0; 32]), - timestamp: 0, - }; - let genesis_block = genesis_hashable_data.into_pending_block(&signing_key); - - SequencerStore::create_db_with_genesis( - &config.home.join("rocksdb"), - &genesis_block, - &genesis_state, - signing_key, - ) - .unwrap(); - - let (sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config).await; - assert_eq!(sequencer.chain_height, 1); - assert!(sequencer.store.latest_block_meta().is_ok()); - } - - #[should_panic(expected = "Failed to open database")] - #[tokio::test] - async fn start_from_config_panics_when_db_open_returns_non_not_found_error() { - let mut config = setup_sequencer_config(); - let temp_dir = tempdir().unwrap(); - config.home = temp_dir.path().to_path_buf(); - - let db_path = config.home.join("rocksdb"); - - std::fs::create_dir_all(&config.home).unwrap(); - // Force RocksDB open to fail with an IO error by placing a file at DB path. - std::fs::write(&db_path, b"not-a-directory").unwrap(); - - let _ = SequencerCoreWithMockClients::start_from_config(config).await; - } - - #[tokio::test] - async fn start_from_config_replays_unfulfilled_deposit_events_from_db() { - let config = setup_sequencer_config(); - let deposit_op_id = [13_u8; 32]; - let expected_amount = 1_u64; - let recipient_id = initial_public_user_accounts()[0].account_id; - - { - let (_sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config.clone()).await; - } - - let pending_event = PendingDepositEventRecord { - deposit_op_id: HashType(deposit_op_id), - source_tx_hash: HashType([7_u8; 32]), - amount: expected_amount, - metadata: borsh::to_vec(&DepositMetadataForEncoding { recipient_id }).unwrap(), - submitted_in_block_id: None, - }; - - { - let signing_key = lee::PrivateKey::try_new(config.signing_key).unwrap(); - let store = SequencerStore::open_db(&config.home.join("rocksdb"), signing_key).unwrap(); - - let inserted = store - .dbio() - .add_pending_deposit_event(pending_event) - .unwrap(); - assert!(inserted); - } - - let (mut sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config).await; - - let (origin, tx) = tokio::time::timeout(Duration::from_secs(5), async { - loop { - if let Some((origin, tx)) = sequencer.mempool.pop() { - return (origin, tx); - } - - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timed out waiting for pending deposit event to be replayed into mempool"); - - match origin { - TransactionOrigin::Sequencer => {} - TransactionOrigin::User => { - panic!("Unexpected user transaction in empty mempool replay test") - } - } - - assert!(tx_is_bridge_deposit(&tx, deposit_op_id, expected_amount)); - - let pending_events = sequencer.store.get_unfulfilled_deposit_events().unwrap(); - let replayed_event = pending_events - .into_iter() - .find(|event| event.deposit_op_id == HashType(deposit_op_id)) - .expect("Pending deposit event should remain in DB until included in a block"); - assert!(replayed_event.submitted_in_block_id.is_none()); - } - - #[test] - fn transaction_pre_check_pass() { - let tx = common::test_utils::produce_dummy_empty_transaction(); - let result = tx.transaction_stateless_check(); - - assert!(result.is_ok()); - } - - #[tokio::test] - async fn transaction_pre_check_native_transfer_valid() { - let (_sequencer, _mempool_handle) = common_setup().await; - - let acc1 = initial_public_user_accounts()[0].account_id; - let acc2 = initial_public_user_accounts()[1].account_id; - - let sign_key1 = create_signing_key_for_account1(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1, 0, acc2, 10, &sign_key1, - ); - let result = tx.transaction_stateless_check(); - - assert!(result.is_ok()); - } - - #[tokio::test] - async fn transaction_pre_check_native_transfer_other_signature() { - let (mut sequencer, _mempool_handle) = common_setup().await; - - let acc1 = initial_public_user_accounts()[0].account_id; - let acc2 = initial_public_user_accounts()[1].account_id; - - let sign_key2 = create_signing_key_for_account2(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1, 0, acc2, 10, &sign_key2, - ); - - // Signature is valid, stateless check pass - let tx = tx.transaction_stateless_check().unwrap(); - - // Signature is not from sender. Execution fails - let result = tx.execute_check_on_state(&mut sequencer.state, 0, 0); - - assert!(matches!( - result, - Err(lee::error::LeeError::ProgramExecutionFailed(_)) - )); - } - - #[tokio::test] - async fn transaction_pre_check_native_transfer_sent_too_much() { - let (mut sequencer, _mempool_handle) = common_setup().await; - - let acc1 = initial_public_user_accounts()[0].account_id; - let acc2 = initial_public_user_accounts()[1].account_id; - - let sign_key1 = create_signing_key_for_account1(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1, 0, acc2, 10_000_000, &sign_key1, - ); - - let result = tx.transaction_stateless_check(); - - // Passed pre-check - assert!(result.is_ok()); - - let result = result - .unwrap() - .execute_check_on_state(&mut sequencer.state, 0, 0); - let is_failed_at_balance_mismatch = matches!( - result.err().unwrap(), - lee::error::LeeError::ProgramExecutionFailed(_) - ); - - assert!(is_failed_at_balance_mismatch); - } - - #[tokio::test] - async fn transaction_execute_native_transfer() { - let (mut sequencer, _mempool_handle) = common_setup().await; - - let acc1 = initial_public_user_accounts()[0].account_id; - let acc2 = initial_public_user_accounts()[1].account_id; - - let sign_key1 = create_signing_key_for_account1(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1, 0, acc2, 100, &sign_key1, - ); - - tx.execute_check_on_state(&mut sequencer.state, 0, 0) - .unwrap(); - - let bal_from = sequencer.state.get_account_by_id(acc1).balance; - let bal_to = sequencer.state.get_account_by_id(acc2).balance; - - assert_eq!(bal_from, 9900); - assert_eq!(bal_to, 20100); - } - - #[tokio::test] - async fn push_tx_into_mempool_blocks_until_mempool_is_full() { - let config = SequencerConfig { - mempool_max_size: 1, - ..setup_sequencer_config() - }; - let (mut sequencer, mempool_handle) = common_setup_with_config(config).await; - - let tx = common::test_utils::produce_dummy_empty_transaction(); - - // Fill the mempool - mempool_handle - .push((TransactionOrigin::User, tx.clone())) - .await - .unwrap(); - - // Check that pushing another transaction will block - let mut push_fut = pin!(mempool_handle.push((TransactionOrigin::User, tx.clone()))); - let poll = futures::poll!(push_fut.as_mut()); - assert!(poll.is_pending()); - - // Empty the mempool by producing a block - sequencer.produce_new_block().await.unwrap(); - - // Resolve the pending push - assert!(push_fut.await.is_ok()); - } - - #[tokio::test] - async fn build_block_from_mempool() { - let (mut sequencer, mempool_handle) = common_setup().await; - let genesis_height = sequencer.chain_height; - - let tx = common::test_utils::produce_dummy_empty_transaction(); - mempool_handle - .push((TransactionOrigin::User, tx)) - .await - .unwrap(); - - let result = sequencer.build_block_from_mempool(); - assert!(result.is_ok()); - assert_eq!(sequencer.chain_height, genesis_height + 1); - } - - #[tokio::test] - async fn replay_transactions_are_rejected_in_the_same_block() { - let (mut sequencer, mempool_handle) = common_setup().await; - - let acc1 = initial_public_user_accounts()[0].account_id; - let acc2 = initial_public_user_accounts()[1].account_id; - - let sign_key1 = create_signing_key_for_account1(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1, 0, acc2, 100, &sign_key1, - ); - - let tx_original = tx.clone(); - let tx_replay = tx.clone(); - // Pushing two copies of the same tx to the mempool - mempool_handle - .push((TransactionOrigin::User, tx_original)) - .await - .unwrap(); - mempool_handle - .push((TransactionOrigin::User, tx_replay)) - .await - .unwrap(); - - // Create block - sequencer.produce_new_block().await.unwrap(); - let block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - - // Only one user tx should be included; the clock tx is always appended last. - assert_eq!( - block.body.transactions, - vec![ - tx.clone(), - LeeTransaction::Public(clock_invocation(block.header.timestamp)) - ] - ); - } - - #[tokio::test] - async fn replay_transactions_are_rejected_in_different_blocks() { - let (mut sequencer, mempool_handle) = common_setup().await; - - let acc1 = initial_public_user_accounts()[0].account_id; - let acc2 = initial_public_user_accounts()[1].account_id; - - let sign_key1 = create_signing_key_for_account1(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1, 0, acc2, 100, &sign_key1, - ); - - // The transaction should be included the first time - mempool_handle - .push((TransactionOrigin::User, tx.clone())) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - let block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - assert_eq!( - block.body.transactions, - vec![ - tx.clone(), - LeeTransaction::Public(clock_invocation(block.header.timestamp)) - ] - ); - - // Add same transaction should fail - mempool_handle - .push((TransactionOrigin::User, tx.clone())) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - let block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - // The replay is rejected, so only the clock tx is in the block. - assert_eq!( - block.body.transactions, - vec![LeeTransaction::Public(clock_invocation( - block.header.timestamp - ))] - ); - } - - #[tokio::test] - async fn restart_from_storage() { - let config = setup_sequencer_config(); - let acc1_account_id = initial_public_user_accounts()[0].account_id; - let acc2_account_id = initial_public_user_accounts()[1].account_id; - let balance_to_move = 13; - - // In the following code block a transaction will be processed that moves `balance_to_move` - // from `acc_1` to `acc_2`. The block created with that transaction will be kept stored in - // the temporary directory for the block storage of this test. - { - let (mut sequencer, mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config.clone()).await; - let signing_key = create_signing_key_for_account1(); - - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1_account_id, - 0, - acc2_account_id, - balance_to_move, - &signing_key, - ); - - mempool_handle - .push((TransactionOrigin::User, tx.clone())) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - let block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - assert_eq!( - block.body.transactions, - vec![ - tx.clone(), - LeeTransaction::Public(clock_invocation(block.header.timestamp)) - ] - ); - } - - // Instantiating a new sequencer from the same config. This should load the existing block - // with the above transaction and update the state to reflect that. - let (sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config.clone()).await; - let balance_acc_1 = sequencer.state.get_account_by_id(acc1_account_id).balance; - let balance_acc_2 = sequencer.state.get_account_by_id(acc2_account_id).balance; - - // Balances should be consistent with the stored block - assert_eq!( - balance_acc_1, - initial_public_user_accounts()[0].balance - balance_to_move - ); - assert_eq!( - balance_acc_2, - initial_public_user_accounts()[1].balance + balance_to_move - ); - } - - #[tokio::test] - async fn get_pending_blocks() { - let config = setup_sequencer_config(); - let (mut sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config).await; - sequencer.produce_new_block().await.unwrap(); - sequencer.produce_new_block().await.unwrap(); - sequencer.produce_new_block().await.unwrap(); - assert_eq!(sequencer.get_pending_blocks().unwrap().len(), 4); - } - - #[tokio::test] - async fn delete_blocks() { - let config = setup_sequencer_config(); - let (mut sequencer, _mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config).await; - sequencer.produce_new_block().await.unwrap(); - sequencer.produce_new_block().await.unwrap(); - sequencer.produce_new_block().await.unwrap(); - - let last_finalized_block = 3; - sequencer - .clean_finalized_blocks_from_db(last_finalized_block) - .unwrap(); - - assert_eq!(sequencer.get_pending_blocks().unwrap().len(), 1); - } - - #[tokio::test] - async fn produce_block_with_correct_prev_meta_after_restart() { - let config = setup_sequencer_config(); - let acc1_account_id = initial_public_user_accounts()[0].account_id; - let acc2_account_id = initial_public_user_accounts()[1].account_id; - - // Step 1: Create initial database with some block metadata - let expected_prev_meta = { - let (mut sequencer, mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config.clone()).await; - - let signing_key = create_signing_key_for_account1(); - - // Add a transaction and produce a block to set up block metadata - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1_account_id, - 0, - acc2_account_id, - 100, - &signing_key, - ); - - mempool_handle - .push((TransactionOrigin::User, tx)) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - - // Get the metadata of the last block produced - sequencer.store.latest_block_meta().unwrap() - }; - - // Step 2: Restart sequencer from the same storage - let (mut sequencer, mempool_handle) = - SequencerCoreWithMockClients::start_from_config(config.clone()).await; - - // Step 3: Submit a new transaction - let signing_key = create_signing_key_for_account1(); - let tx = common::test_utils::create_transaction_native_token_transfer( - acc1_account_id, - 1, // Next nonce - acc2_account_id, - 50, - &signing_key, - ); - - mempool_handle - .push((TransactionOrigin::User, tx.clone())) - .await - .unwrap(); - - // Step 4: Produce new block - sequencer.produce_new_block().await.unwrap(); - - // Step 5: Verify the new block has correct previous block metadata - let new_block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - - assert_eq!( - new_block.header.prev_block_hash, expected_prev_meta.hash, - "New block's prev_block_hash should match the stored metadata hash" - ); - assert_eq!( - new_block.body.transactions, - vec![ - tx, - LeeTransaction::Public(clock_invocation(new_block.header.timestamp)) - ], - "New block should contain the submitted transaction and the clock invocation" - ); - } - - #[tokio::test] - async fn transactions_touching_clock_account_are_dropped_from_block() { - let (mut sequencer, mempool_handle) = common_setup().await; - - // Canonical clock invocation and a crafted variant with a different timestamp — both must - // be dropped because their diffs touch the clock accounts. - let crafted_clock_tx = { - let message = lee::public_transaction::Message::try_new( - programs::clock().id(), - system_accounts::clock_account_ids().to_vec(), - vec![], - 42_u64, - ) - .unwrap(); - LeeTransaction::Public(lee::PublicTransaction::new( - message, - lee::public_transaction::WitnessSet::from_raw_parts(vec![]), - )) - }; - mempool_handle - .push(( - TransactionOrigin::User, - LeeTransaction::Public(clock_invocation(0)), - )) - .await - .unwrap(); - mempool_handle - .push((TransactionOrigin::User, crafted_clock_tx)) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - - let block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - - // Both transactions were dropped. Only the system-appended clock tx remains. - assert_eq!( - block.body.transactions, - vec![LeeTransaction::Public(clock_invocation( - block.header.timestamp - ))] - ); - } - - #[tokio::test] - async fn user_tx_that_chain_calls_clock_is_dropped() { - let (mut sequencer, mempool_handle) = common_setup().await; - - let clock_chain_caller = test_programs::clock_chain_caller(); - // Deploy the clock_chain_caller test program. - let deploy_tx = LeeTransaction::ProgramDeployment(lee::ProgramDeploymentTransaction::new( - lee::program_deployment_transaction::Message::new(clock_chain_caller.elf().to_owned()), - )); - mempool_handle - .push((TransactionOrigin::User, deploy_tx)) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - - // Build a user transaction that invokes clock_chain_caller, which in turn chain-calls the - // clock program with the clock accounts. The sequencer should detect that the resulting - // state diff modifies clock accounts and drop the transaction. - let clock_chain_caller_id = test_programs::clock_chain_caller().id(); - let clock_program_id = programs::clock().id(); - let timestamp: u64 = 0; - - let message = lee::public_transaction::Message::try_new( - clock_chain_caller_id, - system_accounts::clock_account_ids().to_vec(), - vec![], // no signers - (clock_program_id, timestamp), - ) - .unwrap(); - let user_tx = LeeTransaction::Public(lee::PublicTransaction::new( - message, - lee::public_transaction::WitnessSet::from_raw_parts(vec![]), - )); - - mempool_handle - .push((TransactionOrigin::User, user_tx)) - .await - .unwrap(); - sequencer.produce_new_block().await.unwrap(); - - let block = sequencer - .store - .get_block_at_id(sequencer.chain_height) - .unwrap() - .unwrap(); - - // The user tx must have been dropped; only the mandatory clock invocation remains. - assert_eq!( - block.body.transactions, - vec![LeeTransaction::Public(clock_invocation( - block.header.timestamp - ))] - ); - } - - #[tokio::test] - async fn block_production_aborts_when_clock_account_data_is_corrupted() { - let (mut sequencer, mempool_handle) = common_setup().await; - - // Corrupt the clock 01 account data so the clock program panics on deserialization. - let clock_account_id = system_accounts::clock_account_ids()[0]; - let mut corrupted = sequencer.state.get_account_by_id(clock_account_id); - corrupted.data = vec![0xff; 3].try_into().unwrap(); - sequencer - .state - .force_insert_account(clock_account_id, corrupted); - - // Push a dummy transaction so the mempool is non-empty. - let tx = common::test_utils::produce_dummy_empty_transaction(); - mempool_handle - .push((TransactionOrigin::User, tx)) - .await - .unwrap(); - - // Block production must fail because the appended clock tx cannot execute. - let result = sequencer.produce_new_block().await; - assert!( - result.is_err(), - "Block production should abort when clock account data is corrupted" - ); - } - - #[test] - fn private_bridge_withdraw_invocation_is_dropped() { - let sender_keys = KeyChain::new_os_random(); - let sender_account_id = - AccountId::for_regular_private_account(&sender_keys.nullifier_public_key, 0); - let sender_private_account = Account { - program_owner: programs::authenticated_transfer().id(), - balance: 100, - nonce: Nonce(0xdead_beef), - data: Data::default(), - }; - let bridge_account_id = system_accounts::bridge_account_id(); - - let mut state = V03State::new() - .with_public_accounts([(bridge_account_id, system_accounts::bridge_account())]) - .with_private_accounts([( - Commitment::new(&sender_account_id, &sender_private_account), - Nullifier::for_account_initialization(&sender_account_id), - )]); - - let sender_commitment = Commitment::new(&sender_account_id, &sender_private_account); - - let sender_pre = AccountWithMetadata::new( - sender_private_account, - true, - (&sender_keys.nullifier_public_key, 0), - ); - let bridge_pre = AccountWithMetadata::new( - state.get_account_by_id(bridge_account_id), - false, - bridge_account_id, - ); - - let shared_secret = SharedSecretKey::encapsulate(&sender_keys.viewing_public_key).0; - - let instruction = Program::serialize_instruction(bridge_core::Instruction::Withdraw { - amount: 1, - bedrock_account_pk: [0; 32], - }) - .unwrap(); - - let program_with_deps = ProgramWithDependencies::new( - programs::bridge(), - [( - programs::authenticated_transfer().id(), - programs::authenticated_transfer(), - )] - .into(), - ); - - let (output, proof) = execute_and_prove( - vec![sender_pre, bridge_pre], - instruction, - vec![ - InputAccountIdentity::PrivateAuthorizedUpdate { - epk: EphemeralPublicKey(vec![12_u8; 1088]), - view_tag: EncryptedAccountData::compute_view_tag( - &sender_keys.nullifier_public_key, - &sender_keys.viewing_public_key, - ), - ssk: shared_secret, - nsk: sender_keys.private_key_holder.nullifier_secret_key, - membership_proof: state - .get_proof_for_commitment(&sender_commitment) - .expect("sender commitment must be in state"), - identifier: 0, - }, - InputAccountIdentity::Public, - ], - &program_with_deps, - ) - .expect("Execution should succeed"); - - let message = Message::try_from_circuit_output(vec![bridge_account_id], vec![], output) - .expect("Message construction should succeed"); - let witness_set = - lee::privacy_preserving_transaction::WitnessSet::for_message(&message, proof, &[]); - let tx = LeeTransaction::PrivacyPreserving(PrivacyPreservingTransaction::new( - message, - witness_set, - )); - let res = tx.execute_check_on_state(&mut state, 1, 0); - - assert!( - matches!(res, Err(LeeError::InvalidInput(_))), - "Bridge withdraw invocation should be rejected in private execution" - ); - } - - /// Builds a [`V03State`] with the clock program and `program` registered, the three clock - /// accounts initialized, and the clock advanced to `clock_timestamp` so that reads of the - /// `CLOCK_01` account observe it. - fn state_with_clock_and_program(program: Program, clock_timestamp: u64) -> V03State { - let mut state = V03State::new().with_programs([programs::clock(), program]); - for clock_id in system_accounts::clock_account_ids() { - state.force_insert_account(clock_id, system_accounts::clock_account()); - } - state - .transition_from_public_transaction( - &clock_invocation(clock_timestamp), - 1, - clock_timestamp, - ) - .expect("Clock invocation should advance the clock"); - state - } - - fn time_locked_transfer_transaction( - from: AccountId, - from_key: &PrivateKey, - from_nonce: u128, - to: AccountId, - clock_account_id: AccountId, - amount: u128, - deadline: u64, - ) -> PublicTransaction { - let program_id = test_programs::time_locked_transfer().id(); - let message = lee::public_transaction::Message::try_new( - program_id, - vec![from, to, clock_account_id], - vec![Nonce(from_nonce)], - (amount, deadline), - ) - .unwrap(); - let witness_set = lee::public_transaction::WitnessSet::for_message(&message, &[from_key]); - PublicTransaction::new(message, witness_set) - } - - #[test] - fn time_locked_transfer_succeeds_when_deadline_has_passed() { - let clock_timestamp = 600; - let mut state = - state_with_clock_and_program(test_programs::time_locked_transfer(), clock_timestamp); - - // The recipient must be a non-default account so the program may credit it without - // claiming it. - let recipient_id = AccountId::new([42; 32]); - state.force_insert_account( - recipient_id, - Account { - program_owner: programs::authenticated_transfer().id(), - ..Account::default() - }, - ); - - let key1 = PrivateKey::try_new([1; 32]).unwrap(); - let sender_id = AccountId::from(&PublicKey::new_from_private_key(&key1)); - state.force_insert_account( - sender_id, - Account { - program_owner: test_programs::time_locked_transfer().id(), - balance: 100, - ..Account::default() - }, - ); - - let amount = 100; - // Deadline is in the past relative to the clock, so the transfer is unlocked. - let deadline = 0; - - let tx = time_locked_transfer_transaction( - sender_id, - &key1, - 0, - recipient_id, - system_accounts::clock_account_ids()[0], - amount, - deadline, - ); - - state - .transition_from_public_transaction(&tx, 2, clock_timestamp) - .unwrap(); - - // Balances changed. - assert_eq!(state.get_account_by_id(sender_id).balance, 0); - assert_eq!(state.get_account_by_id(recipient_id).balance, 100); - } - - #[test] - fn time_locked_transfer_fails_when_deadline_is_in_the_future() { - let clock_timestamp = 600; - let mut state = - state_with_clock_and_program(test_programs::time_locked_transfer(), clock_timestamp); - - let recipient_id = AccountId::new([42; 32]); - state.force_insert_account( - recipient_id, - Account { - program_owner: programs::authenticated_transfer().id(), - ..Account::default() - }, - ); - - let key1 = PrivateKey::try_new([1; 32]).unwrap(); - let sender_id = AccountId::from(&PublicKey::new_from_private_key(&key1)); - state.force_insert_account( - sender_id, - Account { - program_owner: test_programs::time_locked_transfer().id(), - balance: 100, - ..Account::default() - }, - ); - - let amount = 100; - // Far-future deadline: the program panics because the clock has not reached it. - let deadline = u64::MAX; - - let tx = time_locked_transfer_transaction( - sender_id, - &key1, - 0, - recipient_id, - system_accounts::clock_account_ids()[0], - amount, - deadline, - ); - - let result = state.transition_from_public_transaction(&tx, 2, clock_timestamp); - - assert!( - result.is_err(), - "Transfer should fail when deadline is in the future" - ); - // Balances unchanged. - assert_eq!(state.get_account_by_id(sender_id).balance, 100); - assert_eq!(state.get_account_by_id(recipient_id).balance, 0); - } - - fn pinata_cooldown_data(prize: u128, cooldown_ms: u64, last_claim_timestamp: u64) -> Vec { - let mut buf = Vec::with_capacity(32); - buf.extend_from_slice(&prize.to_le_bytes()); - buf.extend_from_slice(&cooldown_ms.to_le_bytes()); - buf.extend_from_slice(&last_claim_timestamp.to_le_bytes()); - buf - } - - fn pinata_cooldown_transaction( - pinata_id: AccountId, - winner_id: AccountId, - clock_account_id: AccountId, - ) -> PublicTransaction { - let program_id = test_programs::pinata_cooldown().id(); - let message = lee::public_transaction::Message::try_new( - program_id, - vec![pinata_id, winner_id, clock_account_id], - vec![], - (), - ) - .unwrap(); - let witness_set = lee::public_transaction::WitnessSet::for_message(&message, &[]); - PublicTransaction::new(message, witness_set) - } - - #[test] - fn pinata_cooldown_claim_succeeds_after_cooldown() { - let winner_id = AccountId::new([11; 32]); - let pinata_id = AccountId::new([99; 32]); - - let genesis_timestamp = 1000; - let prize = 50; - let cooldown_ms = 500; - // Last claim was at genesis, so any timestamp >= genesis + cooldown should work. - let last_claim_timestamp = genesis_timestamp; - - // Advance the clock so the cooldown check reads an updated timestamp. - let block_timestamp = genesis_timestamp + cooldown_ms; - let mut state = - state_with_clock_and_program(test_programs::pinata_cooldown(), block_timestamp); - - // The winner must be a non-default account so the program may credit it without claiming. - state.force_insert_account( - winner_id, - Account { - program_owner: programs::authenticated_transfer().id(), - ..Account::default() - }, - ); - state.force_insert_account( - pinata_id, - Account { - program_owner: test_programs::pinata_cooldown().id(), - balance: 1000, - data: pinata_cooldown_data(prize, cooldown_ms, last_claim_timestamp) - .try_into() - .unwrap(), - ..Account::default() - }, - ); - - let tx = pinata_cooldown_transaction( - pinata_id, - winner_id, - system_accounts::clock_account_ids()[0], - ); - - state - .transition_from_public_transaction(&tx, 2, block_timestamp) - .unwrap(); - - assert_eq!(state.get_account_by_id(pinata_id).balance, 1000 - prize); - assert_eq!(state.get_account_by_id(winner_id).balance, prize); - } - - #[test] - fn pinata_cooldown_claim_fails_during_cooldown() { - let winner_id = AccountId::new([11; 32]); - let pinata_id = AccountId::new([99; 32]); - - let genesis_timestamp = 1000; - let prize = 50; - let cooldown_ms = 500; - let last_claim_timestamp = genesis_timestamp; - - // Timestamp is only 100ms after the last claim, well within the 500ms cooldown. - let block_timestamp = genesis_timestamp + 100; - let mut state = - state_with_clock_and_program(test_programs::pinata_cooldown(), block_timestamp); - - state.force_insert_account( - winner_id, - Account { - program_owner: programs::authenticated_transfer().id(), - ..Account::default() - }, - ); - state.force_insert_account( - pinata_id, - Account { - program_owner: test_programs::pinata_cooldown().id(), - balance: 1000, - data: pinata_cooldown_data(prize, cooldown_ms, last_claim_timestamp) - .try_into() - .unwrap(), - ..Account::default() - }, - ); - - let tx = pinata_cooldown_transaction( - pinata_id, - winner_id, - system_accounts::clock_account_ids()[0], - ); - - let result = state.transition_from_public_transaction(&tx, 2, block_timestamp); - - assert!(result.is_err(), "Claim should fail during cooldown period"); - assert_eq!(state.get_account_by_id(pinata_id).balance, 1000); - assert_eq!(state.get_account_by_id(winner_id).balance, 0); - } - - #[test] - fn pda_mechanism_with_pinata_token_program() { - let pinata_token = programs::pinata_token(); - let token = programs::token(); - - let pinata_definition_id = AccountId::new([1; 32]); - let pinata_token_definition_id = AccountId::new([2; 32]); - // Total supply of pinata token will be in an account under a PDA. - let pinata_token_holding_id = - AccountId::for_public_pda(&pinata_token.id(), &PdaSeed::new([0; 32])); - let winner_token_holding_id = AccountId::new([3; 32]); - - let expected_winner_account_holding = token_core::TokenHolding::Fungible { - definition_id: pinata_token_definition_id, - balance: 150, - }; - let expected_winner_token_holding_post = Account { - program_owner: token.id(), - data: Data::from(&expected_winner_account_holding), - ..Account::default() - }; - - // Register the pinata-token and token programs and create the pinata definition account. - // This replaces the removed `add_pinata_token_program` helper. - let mut state = V03State::new().with_programs([pinata_token.clone(), token.clone()]); - state.force_insert_account( - pinata_definition_id, - Account { - program_owner: pinata_token.id(), - // Difficulty: 3 - data: vec![3; 33].try_into().unwrap(), - ..Account::default() - }, - ); - - // Set up the token accounts directly (bypassing public transactions which - // would require signers for Claim::Authorized). The focus of this test is - // the PDA mechanism in the pinata program's chained call, not token creation. - let total_supply: u128 = 10_000_000; - let token_definition = token_core::TokenDefinition::Fungible { - name: String::from("PINATA"), - total_supply, - metadata_id: None, - }; - let token_holding = token_core::TokenHolding::Fungible { - definition_id: pinata_token_definition_id, - balance: total_supply, - }; - let winner_holding = token_core::TokenHolding::Fungible { - definition_id: pinata_token_definition_id, - balance: 0, - }; - state.force_insert_account( - pinata_token_definition_id, - Account { - program_owner: token.id(), - data: Data::from(&token_definition), - ..Account::default() - }, - ); - state.force_insert_account( - pinata_token_holding_id, - Account { - program_owner: token.id(), - data: Data::from(&token_holding), - ..Account::default() - }, - ); - state.force_insert_account( - winner_token_holding_id, - Account { - program_owner: token.id(), - data: Data::from(&winner_holding), - ..Account::default() - }, - ); - - // Submit a solution to the pinata program to claim the prize - let solution: u128 = 989_106; - let message = lee::public_transaction::Message::try_new( - pinata_token.id(), - vec![ - pinata_definition_id, - pinata_token_holding_id, - winner_token_holding_id, - ], - vec![], - solution, - ) - .unwrap(); - let witness_set = lee::public_transaction::WitnessSet::for_message(&message, &[]); - let tx = PublicTransaction::new(message, witness_set); - state.transition_from_public_transaction(&tx, 1, 0).unwrap(); - - let winner_token_holding_post = state.get_account_by_id(winner_token_holding_id); - assert_eq!( - winner_token_holding_post, - expected_winner_token_holding_post - ); - } -} +mod tests; diff --git a/lez/sequencer/core/src/tests.rs b/lez/sequencer/core/src/tests.rs new file mode 100644 index 00000000..a5c535a0 --- /dev/null +++ b/lez/sequencer/core/src/tests.rs @@ -0,0 +1,1233 @@ +#![expect(clippy::shadow_unrelated, reason = "We don't care about it in tests")] + +use std::{pin::pin, time::Duration}; + +use common::{ + HashType, + block::HashableBlockData, + test_utils::sequencer_sign_key_for_testing, + transaction::{LeeTransaction, clock_invocation}, +}; +use key_protocol::key_management::KeyChain; +use lee::{ + Account, AccountId, Data, EphemeralPublicKey, PrivacyPreservingTransaction, PrivateKey, + PublicKey, PublicTransaction, SharedSecretKey, V03State, + error::LeeError, + execute_and_prove, + privacy_preserving_transaction::{Message, circuit::ProgramWithDependencies}, + program::Program, +}; +use lee_core::{ + Commitment, EncryptedAccountData, InputAccountIdentity, Nullifier, + account::{AccountWithMetadata, Nonce}, + program::PdaSeed, +}; +use logos_blockchain_core::mantle::ops::channel::ChannelId; +use mempool::MemPoolHandle; +use storage::sequencer::sequencer_cells::PendingDepositEventRecord; +use tempfile::tempdir; +use testnet_initial_state::{initial_pub_accounts_private_keys, initial_public_user_accounts}; + +use crate::{ + TransactionOrigin, + block_store::SequencerStore, + build_genesis_state, + config::{BedrockConfig, SequencerConfig}, + mock::SequencerCoreWithMockClients, +}; + +#[derive(borsh::BorshSerialize)] +struct DepositMetadataForEncoding { + recipient_id: lee::AccountId, +} + +fn setup_sequencer_config() -> SequencerConfig { + let tempdir = tempfile::tempdir().unwrap(); + let home = tempdir.path().to_path_buf(); + + SequencerConfig { + home, + max_num_tx_in_block: 10, + max_block_size: bytesize::ByteSize::mib(1), + mempool_max_size: 10000, + block_create_timeout: Duration::from_secs(1), + signing_key: *sequencer_sign_key_for_testing().value(), + bedrock_config: BedrockConfig { + channel_id: ChannelId::from([0; 32]), + node_url: "http://not-used-in-unit-tests".parse().unwrap(), + auth: None, + }, + retry_pending_blocks_timeout: Duration::from_mins(4), + genesis: vec![], + } +} + +fn create_signing_key_for_account1() -> lee::PrivateKey { + initial_pub_accounts_private_keys()[0].pub_sign_key.clone() +} + +fn create_signing_key_for_account2() -> lee::PrivateKey { + initial_pub_accounts_private_keys()[1].pub_sign_key.clone() +} + +async fn common_setup() -> ( + SequencerCoreWithMockClients, + MemPoolHandle<(TransactionOrigin, LeeTransaction)>, +) { + let config = setup_sequencer_config(); + common_setup_with_config(config).await +} + +async fn common_setup_with_config( + config: SequencerConfig, +) -> ( + SequencerCoreWithMockClients, + MemPoolHandle<(TransactionOrigin, LeeTransaction)>, +) { + let (mut sequencer, mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config).await; + + let tx = common::test_utils::produce_dummy_empty_transaction(); + mempool_handle + .push((TransactionOrigin::User, tx)) + .await + .unwrap(); + + sequencer.produce_new_block().await.unwrap(); + + (sequencer, mempool_handle) +} + +fn tx_is_bridge_deposit( + tx: &LeeTransaction, + deposit_op_id: [u8; 32], + expected_amount: u64, +) -> bool { + let LeeTransaction::Public(public_tx) = tx else { + return false; + }; + + if public_tx.message.program_id != programs::bridge().id() { + return false; + } + + let instruction: bridge_core::Instruction = + match risc0_zkvm::serde::from_slice(&public_tx.message.instruction_data) { + Ok(instruction) => instruction, + Err(_err) => return false, + }; + + matches!( + instruction, + bridge_core::Instruction::Deposit { + l1_deposit_op_id, + amount, + .. + } if l1_deposit_op_id == deposit_op_id && amount == expected_amount + ) +} + +#[tokio::test] +async fn start_from_config() { + let config = setup_sequencer_config(); + let (sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config.clone()).await; + + assert_eq!(sequencer.chain_height, 1); + assert_eq!(sequencer.sequencer_config.max_num_tx_in_block, 10); + + let acc1_account_id = initial_public_user_accounts()[0].account_id; + let acc2_account_id = initial_public_user_accounts()[1].account_id; + + let balance_acc_1 = sequencer.state.get_account_by_id(acc1_account_id).balance; + let balance_acc_2 = sequencer.state.get_account_by_id(acc2_account_id).balance; + + assert_eq!(10000, balance_acc_1); + assert_eq!(20000, balance_acc_2); +} + +#[tokio::test] +async fn start_from_config_opens_existing_db_if_it_exists() { + let config = setup_sequencer_config(); + let temp_dir = tempdir().unwrap(); + let mut config = config; + config.home = temp_dir.path().to_path_buf(); + + let signing_key = lee::PrivateKey::try_new(config.signing_key).unwrap(); + let (genesis_state, genesis_txs) = build_genesis_state(&config); + let genesis_hashable_data = HashableBlockData { + block_id: 1, + transactions: genesis_txs, + prev_block_hash: HashType([0; 32]), + timestamp: 0, + }; + let genesis_block = genesis_hashable_data.into_pending_block(&signing_key); + + SequencerStore::create_db_with_genesis( + &config.home.join("rocksdb"), + &genesis_block, + &genesis_state, + signing_key, + ) + .unwrap(); + + let (sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config).await; + assert_eq!(sequencer.chain_height, 1); + assert!(sequencer.store.latest_block_meta().is_ok()); +} + +#[should_panic(expected = "Failed to open database")] +#[tokio::test] +async fn start_from_config_panics_when_db_open_returns_non_not_found_error() { + let mut config = setup_sequencer_config(); + let temp_dir = tempdir().unwrap(); + config.home = temp_dir.path().to_path_buf(); + + let db_path = config.home.join("rocksdb"); + + std::fs::create_dir_all(&config.home).unwrap(); + // Force RocksDB open to fail with an IO error by placing a file at DB path. + std::fs::write(&db_path, b"not-a-directory").unwrap(); + + let _ = SequencerCoreWithMockClients::start_from_config(config).await; +} + +#[tokio::test] +async fn start_from_config_replays_unfulfilled_deposit_events_from_db() { + let config = setup_sequencer_config(); + let deposit_op_id = [13_u8; 32]; + let expected_amount = 1_u64; + let recipient_id = initial_public_user_accounts()[0].account_id; + + { + let (_sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config.clone()).await; + } + + let pending_event = PendingDepositEventRecord { + deposit_op_id: HashType(deposit_op_id), + source_tx_hash: HashType([7_u8; 32]), + amount: expected_amount, + metadata: borsh::to_vec(&DepositMetadataForEncoding { recipient_id }).unwrap(), + submitted_in_block_id: None, + }; + + { + let signing_key = lee::PrivateKey::try_new(config.signing_key).unwrap(); + let store = SequencerStore::open_db(&config.home.join("rocksdb"), signing_key).unwrap(); + + let inserted = store + .dbio() + .add_pending_deposit_event(pending_event) + .unwrap(); + assert!(inserted); + } + + let (mut sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config).await; + + let (origin, tx) = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Some((origin, tx)) = sequencer.mempool.pop() { + return (origin, tx); + } + + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timed out waiting for pending deposit event to be replayed into mempool"); + + match origin { + TransactionOrigin::Sequencer => {} + TransactionOrigin::User => { + panic!("Unexpected user transaction in empty mempool replay test") + } + } + + assert!(tx_is_bridge_deposit(&tx, deposit_op_id, expected_amount)); + + let pending_events = sequencer.store.get_unfulfilled_deposit_events().unwrap(); + let replayed_event = pending_events + .into_iter() + .find(|event| event.deposit_op_id == HashType(deposit_op_id)) + .expect("Pending deposit event should remain in DB until included in a block"); + assert!(replayed_event.submitted_in_block_id.is_none()); +} + +#[test] +fn transaction_pre_check_pass() { + let tx = common::test_utils::produce_dummy_empty_transaction(); + let result = tx.transaction_stateless_check(); + + assert!(result.is_ok()); +} + +#[tokio::test] +async fn transaction_pre_check_native_transfer_valid() { + let (_sequencer, _mempool_handle) = common_setup().await; + + let acc1 = initial_public_user_accounts()[0].account_id; + let acc2 = initial_public_user_accounts()[1].account_id; + + let sign_key1 = create_signing_key_for_account1(); + + let tx = + common::test_utils::create_transaction_native_token_transfer(acc1, 0, acc2, 10, &sign_key1); + let result = tx.transaction_stateless_check(); + + assert!(result.is_ok()); +} + +#[tokio::test] +async fn transaction_pre_check_native_transfer_other_signature() { + let (mut sequencer, _mempool_handle) = common_setup().await; + + let acc1 = initial_public_user_accounts()[0].account_id; + let acc2 = initial_public_user_accounts()[1].account_id; + + let sign_key2 = create_signing_key_for_account2(); + + let tx = + common::test_utils::create_transaction_native_token_transfer(acc1, 0, acc2, 10, &sign_key2); + + // Signature is valid, stateless check pass + let tx = tx.transaction_stateless_check().unwrap(); + + // Signature is not from sender. Execution fails + let result = tx.execute_check_on_state(&mut sequencer.state, 0, 0); + + assert!(matches!( + result, + Err(lee::error::LeeError::ProgramExecutionFailed(_)) + )); +} + +#[tokio::test] +async fn transaction_pre_check_native_transfer_sent_too_much() { + let (mut sequencer, _mempool_handle) = common_setup().await; + + let acc1 = initial_public_user_accounts()[0].account_id; + let acc2 = initial_public_user_accounts()[1].account_id; + + let sign_key1 = create_signing_key_for_account1(); + + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1, 0, acc2, 10_000_000, &sign_key1, + ); + + let result = tx.transaction_stateless_check(); + + // Passed pre-check + assert!(result.is_ok()); + + let result = result + .unwrap() + .execute_check_on_state(&mut sequencer.state, 0, 0); + let is_failed_at_balance_mismatch = matches!( + result.err().unwrap(), + lee::error::LeeError::ProgramExecutionFailed(_) + ); + + assert!(is_failed_at_balance_mismatch); +} + +#[tokio::test] +async fn transaction_execute_native_transfer() { + let (mut sequencer, _mempool_handle) = common_setup().await; + + let acc1 = initial_public_user_accounts()[0].account_id; + let acc2 = initial_public_user_accounts()[1].account_id; + + let sign_key1 = create_signing_key_for_account1(); + + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1, 0, acc2, 100, &sign_key1, + ); + + tx.execute_check_on_state(&mut sequencer.state, 0, 0) + .unwrap(); + + let bal_from = sequencer.state.get_account_by_id(acc1).balance; + let bal_to = sequencer.state.get_account_by_id(acc2).balance; + + assert_eq!(bal_from, 9900); + assert_eq!(bal_to, 20100); +} + +#[tokio::test] +async fn push_tx_into_mempool_blocks_until_mempool_is_full() { + let config = SequencerConfig { + mempool_max_size: 1, + ..setup_sequencer_config() + }; + let (mut sequencer, mempool_handle) = common_setup_with_config(config).await; + + let tx = common::test_utils::produce_dummy_empty_transaction(); + + // Fill the mempool + mempool_handle + .push((TransactionOrigin::User, tx.clone())) + .await + .unwrap(); + + // Check that pushing another transaction will block + let mut push_fut = pin!(mempool_handle.push((TransactionOrigin::User, tx.clone()))); + let poll = futures::poll!(push_fut.as_mut()); + assert!(poll.is_pending()); + + // Empty the mempool by producing a block + sequencer.produce_new_block().await.unwrap(); + + // Resolve the pending push + assert!(push_fut.await.is_ok()); +} + +#[tokio::test] +async fn build_block_from_mempool() { + let (mut sequencer, mempool_handle) = common_setup().await; + let genesis_height = sequencer.chain_height; + + let tx = common::test_utils::produce_dummy_empty_transaction(); + mempool_handle + .push((TransactionOrigin::User, tx)) + .await + .unwrap(); + + let result = sequencer.build_block_from_mempool(); + assert!(result.is_ok()); + assert_eq!(sequencer.chain_height, genesis_height + 1); +} + +#[tokio::test] +async fn replay_transactions_are_rejected_in_the_same_block() { + let (mut sequencer, mempool_handle) = common_setup().await; + + let acc1 = initial_public_user_accounts()[0].account_id; + let acc2 = initial_public_user_accounts()[1].account_id; + + let sign_key1 = create_signing_key_for_account1(); + + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1, 0, acc2, 100, &sign_key1, + ); + + let tx_original = tx.clone(); + let tx_replay = tx.clone(); + // Pushing two copies of the same tx to the mempool + mempool_handle + .push((TransactionOrigin::User, tx_original)) + .await + .unwrap(); + mempool_handle + .push((TransactionOrigin::User, tx_replay)) + .await + .unwrap(); + + // Create block + sequencer.produce_new_block().await.unwrap(); + let block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + + // Only one user tx should be included; the clock tx is always appended last. + assert_eq!( + block.body.transactions, + vec![ + tx.clone(), + LeeTransaction::Public(clock_invocation(block.header.timestamp)) + ] + ); +} + +#[tokio::test] +async fn replay_transactions_are_rejected_in_different_blocks() { + let (mut sequencer, mempool_handle) = common_setup().await; + + let acc1 = initial_public_user_accounts()[0].account_id; + let acc2 = initial_public_user_accounts()[1].account_id; + + let sign_key1 = create_signing_key_for_account1(); + + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1, 0, acc2, 100, &sign_key1, + ); + + // The transaction should be included the first time + mempool_handle + .push((TransactionOrigin::User, tx.clone())) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + let block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + assert_eq!( + block.body.transactions, + vec![ + tx.clone(), + LeeTransaction::Public(clock_invocation(block.header.timestamp)) + ] + ); + + // Add same transaction should fail + mempool_handle + .push((TransactionOrigin::User, tx.clone())) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + let block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + // The replay is rejected, so only the clock tx is in the block. + assert_eq!( + block.body.transactions, + vec![LeeTransaction::Public(clock_invocation( + block.header.timestamp + ))] + ); +} + +#[tokio::test] +async fn restart_from_storage() { + let config = setup_sequencer_config(); + let acc1_account_id = initial_public_user_accounts()[0].account_id; + let acc2_account_id = initial_public_user_accounts()[1].account_id; + let balance_to_move = 13; + + // In the following code block a transaction will be processed that moves `balance_to_move` + // from `acc_1` to `acc_2`. The block created with that transaction will be kept stored in + // the temporary directory for the block storage of this test. + { + let (mut sequencer, mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config.clone()).await; + let signing_key = create_signing_key_for_account1(); + + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1_account_id, + 0, + acc2_account_id, + balance_to_move, + &signing_key, + ); + + mempool_handle + .push((TransactionOrigin::User, tx.clone())) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + let block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + assert_eq!( + block.body.transactions, + vec![ + tx.clone(), + LeeTransaction::Public(clock_invocation(block.header.timestamp)) + ] + ); + } + + // Instantiating a new sequencer from the same config. This should load the existing block + // with the above transaction and update the state to reflect that. + let (sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config.clone()).await; + let balance_acc_1 = sequencer.state.get_account_by_id(acc1_account_id).balance; + let balance_acc_2 = sequencer.state.get_account_by_id(acc2_account_id).balance; + + // Balances should be consistent with the stored block + assert_eq!( + balance_acc_1, + initial_public_user_accounts()[0].balance - balance_to_move + ); + assert_eq!( + balance_acc_2, + initial_public_user_accounts()[1].balance + balance_to_move + ); +} + +#[tokio::test] +async fn get_pending_blocks() { + let config = setup_sequencer_config(); + let (mut sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config).await; + sequencer.produce_new_block().await.unwrap(); + sequencer.produce_new_block().await.unwrap(); + sequencer.produce_new_block().await.unwrap(); + assert_eq!(sequencer.get_pending_blocks().unwrap().len(), 4); +} + +#[tokio::test] +async fn delete_blocks() { + let config = setup_sequencer_config(); + let (mut sequencer, _mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config).await; + sequencer.produce_new_block().await.unwrap(); + sequencer.produce_new_block().await.unwrap(); + sequencer.produce_new_block().await.unwrap(); + + let last_finalized_block = 3; + sequencer + .clean_finalized_blocks_from_db(last_finalized_block) + .unwrap(); + + assert_eq!(sequencer.get_pending_blocks().unwrap().len(), 1); +} + +#[tokio::test] +async fn produce_block_with_correct_prev_meta_after_restart() { + let config = setup_sequencer_config(); + let acc1_account_id = initial_public_user_accounts()[0].account_id; + let acc2_account_id = initial_public_user_accounts()[1].account_id; + + // Step 1: Create initial database with some block metadata + let expected_prev_meta = { + let (mut sequencer, mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config.clone()).await; + + let signing_key = create_signing_key_for_account1(); + + // Add a transaction and produce a block to set up block metadata + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1_account_id, + 0, + acc2_account_id, + 100, + &signing_key, + ); + + mempool_handle + .push((TransactionOrigin::User, tx)) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + + // Get the metadata of the last block produced + sequencer.store.latest_block_meta().unwrap() + }; + + // Step 2: Restart sequencer from the same storage + let (mut sequencer, mempool_handle) = + SequencerCoreWithMockClients::start_from_config(config.clone()).await; + + // Step 3: Submit a new transaction + let signing_key = create_signing_key_for_account1(); + let tx = common::test_utils::create_transaction_native_token_transfer( + acc1_account_id, + 1, // Next nonce + acc2_account_id, + 50, + &signing_key, + ); + + mempool_handle + .push((TransactionOrigin::User, tx.clone())) + .await + .unwrap(); + + // Step 4: Produce new block + sequencer.produce_new_block().await.unwrap(); + + // Step 5: Verify the new block has correct previous block metadata + let new_block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + + assert_eq!( + new_block.header.prev_block_hash, expected_prev_meta.hash, + "New block's prev_block_hash should match the stored metadata hash" + ); + assert_eq!( + new_block.body.transactions, + vec![ + tx, + LeeTransaction::Public(clock_invocation(new_block.header.timestamp)) + ], + "New block should contain the submitted transaction and the clock invocation" + ); +} + +#[tokio::test] +async fn transactions_touching_clock_account_are_dropped_from_block() { + let (mut sequencer, mempool_handle) = common_setup().await; + + // Canonical clock invocation and a crafted variant with a different timestamp — both must + // be dropped because their diffs touch the clock accounts. + let crafted_clock_tx = { + let message = lee::public_transaction::Message::try_new( + programs::clock().id(), + system_accounts::clock_account_ids().to_vec(), + vec![], + 42_u64, + ) + .unwrap(); + LeeTransaction::Public(lee::PublicTransaction::new( + message, + lee::public_transaction::WitnessSet::from_raw_parts(vec![]), + )) + }; + mempool_handle + .push(( + TransactionOrigin::User, + LeeTransaction::Public(clock_invocation(0)), + )) + .await + .unwrap(); + mempool_handle + .push((TransactionOrigin::User, crafted_clock_tx)) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + + let block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + + // Both transactions were dropped. Only the system-appended clock tx remains. + assert_eq!( + block.body.transactions, + vec![LeeTransaction::Public(clock_invocation( + block.header.timestamp + ))] + ); +} + +#[tokio::test] +async fn user_tx_that_chain_calls_clock_is_dropped() { + let (mut sequencer, mempool_handle) = common_setup().await; + + let clock_chain_caller = test_programs::clock_chain_caller(); + // Deploy the clock_chain_caller test program. + let deploy_tx = LeeTransaction::ProgramDeployment(lee::ProgramDeploymentTransaction::new( + lee::program_deployment_transaction::Message::new(clock_chain_caller.elf().to_owned()), + )); + mempool_handle + .push((TransactionOrigin::User, deploy_tx)) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + + // Build a user transaction that invokes clock_chain_caller, which in turn chain-calls the + // clock program with the clock accounts. The sequencer should detect that the resulting + // state diff modifies clock accounts and drop the transaction. + let clock_chain_caller_id = test_programs::clock_chain_caller().id(); + let clock_program_id = programs::clock().id(); + let timestamp: u64 = 0; + + let message = lee::public_transaction::Message::try_new( + clock_chain_caller_id, + system_accounts::clock_account_ids().to_vec(), + vec![], // no signers + (clock_program_id, timestamp), + ) + .unwrap(); + let user_tx = LeeTransaction::Public(lee::PublicTransaction::new( + message, + lee::public_transaction::WitnessSet::from_raw_parts(vec![]), + )); + + mempool_handle + .push((TransactionOrigin::User, user_tx)) + .await + .unwrap(); + sequencer.produce_new_block().await.unwrap(); + + let block = sequencer + .store + .get_block_at_id(sequencer.chain_height) + .unwrap() + .unwrap(); + + // The user tx must have been dropped; only the mandatory clock invocation remains. + assert_eq!( + block.body.transactions, + vec![LeeTransaction::Public(clock_invocation( + block.header.timestamp + ))] + ); +} + +#[tokio::test] +async fn block_production_aborts_when_clock_account_data_is_corrupted() { + let (mut sequencer, mempool_handle) = common_setup().await; + + // Corrupt the clock 01 account data so the clock program panics on deserialization. + let clock_account_id = system_accounts::clock_account_ids()[0]; + let mut corrupted = sequencer.state.get_account_by_id(clock_account_id); + corrupted.data = vec![0xff; 3].try_into().unwrap(); + sequencer + .state + .force_insert_account(clock_account_id, corrupted); + + // Push a dummy transaction so the mempool is non-empty. + let tx = common::test_utils::produce_dummy_empty_transaction(); + mempool_handle + .push((TransactionOrigin::User, tx)) + .await + .unwrap(); + + // Block production must fail because the appended clock tx cannot execute. + let result = sequencer.produce_new_block().await; + assert!( + result.is_err(), + "Block production should abort when clock account data is corrupted" + ); +} + +#[test] +fn private_bridge_withdraw_invocation_is_dropped() { + let sender_keys = KeyChain::new_os_random(); + let sender_account_id = + AccountId::for_regular_private_account(&sender_keys.nullifier_public_key, 0); + let sender_private_account = Account { + program_owner: programs::authenticated_transfer().id(), + balance: 100, + nonce: Nonce(0xdead_beef), + data: Data::default(), + }; + let bridge_account_id = system_accounts::bridge_account_id(); + + let mut state = V03State::new() + .with_public_accounts([(bridge_account_id, system_accounts::bridge_account())]) + .with_private_accounts([( + Commitment::new(&sender_account_id, &sender_private_account), + Nullifier::for_account_initialization(&sender_account_id), + )]); + + let sender_commitment = Commitment::new(&sender_account_id, &sender_private_account); + + let sender_pre = AccountWithMetadata::new( + sender_private_account, + true, + (&sender_keys.nullifier_public_key, 0), + ); + let bridge_pre = AccountWithMetadata::new( + state.get_account_by_id(bridge_account_id), + false, + bridge_account_id, + ); + + let shared_secret = SharedSecretKey::encapsulate(&sender_keys.viewing_public_key).0; + + let instruction = Program::serialize_instruction(bridge_core::Instruction::Withdraw { + amount: 1, + bedrock_account_pk: [0; 32], + }) + .unwrap(); + + let program_with_deps = ProgramWithDependencies::new( + programs::bridge(), + [( + programs::authenticated_transfer().id(), + programs::authenticated_transfer(), + )] + .into(), + ); + + let (output, proof) = execute_and_prove( + vec![sender_pre, bridge_pre], + instruction, + vec![ + InputAccountIdentity::PrivateAuthorizedUpdate { + epk: EphemeralPublicKey(vec![12_u8; 1088]), + view_tag: EncryptedAccountData::compute_view_tag( + &sender_keys.nullifier_public_key, + &sender_keys.viewing_public_key, + ), + ssk: shared_secret, + nsk: sender_keys.private_key_holder.nullifier_secret_key, + membership_proof: state + .get_proof_for_commitment(&sender_commitment) + .expect("sender commitment must be in state"), + identifier: 0, + }, + InputAccountIdentity::Public, + ], + &program_with_deps, + ) + .expect("Execution should succeed"); + + let message = Message::try_from_circuit_output(vec![bridge_account_id], vec![], output) + .expect("Message construction should succeed"); + let witness_set = + lee::privacy_preserving_transaction::WitnessSet::for_message(&message, proof, &[]); + let tx = + LeeTransaction::PrivacyPreserving(PrivacyPreservingTransaction::new(message, witness_set)); + let res = tx.execute_check_on_state(&mut state, 1, 0); + + assert!( + matches!(res, Err(LeeError::InvalidInput(_))), + "Bridge withdraw invocation should be rejected in private execution" + ); +} + +/// Builds a [`V03State`] with the clock program and `program` registered, the three clock +/// accounts initialized, and the clock advanced to `clock_timestamp` so that reads of the +/// `CLOCK_01` account observe it. +fn state_with_clock_and_program(program: Program, clock_timestamp: u64) -> V03State { + let mut state = V03State::new().with_programs([programs::clock(), program]); + for clock_id in system_accounts::clock_account_ids() { + state.force_insert_account(clock_id, system_accounts::clock_account()); + } + state + .transition_from_public_transaction(&clock_invocation(clock_timestamp), 1, clock_timestamp) + .expect("Clock invocation should advance the clock"); + state +} + +fn time_locked_transfer_transaction( + from: AccountId, + from_key: &PrivateKey, + from_nonce: u128, + to: AccountId, + clock_account_id: AccountId, + amount: u128, + deadline: u64, +) -> PublicTransaction { + let program_id = test_programs::time_locked_transfer().id(); + let message = lee::public_transaction::Message::try_new( + program_id, + vec![from, to, clock_account_id], + vec![Nonce(from_nonce)], + (amount, deadline), + ) + .unwrap(); + let witness_set = lee::public_transaction::WitnessSet::for_message(&message, &[from_key]); + PublicTransaction::new(message, witness_set) +} + +#[test] +fn time_locked_transfer_succeeds_when_deadline_has_passed() { + let clock_timestamp = 600; + let mut state = + state_with_clock_and_program(test_programs::time_locked_transfer(), clock_timestamp); + + // The recipient must be a non-default account so the program may credit it without + // claiming it. + let recipient_id = AccountId::new([42; 32]); + state.force_insert_account( + recipient_id, + Account { + program_owner: programs::authenticated_transfer().id(), + ..Account::default() + }, + ); + + let key1 = PrivateKey::try_new([1; 32]).unwrap(); + let sender_id = AccountId::from(&PublicKey::new_from_private_key(&key1)); + state.force_insert_account( + sender_id, + Account { + program_owner: test_programs::time_locked_transfer().id(), + balance: 100, + ..Account::default() + }, + ); + + let amount = 100; + // Deadline is in the past relative to the clock, so the transfer is unlocked. + let deadline = 0; + + let tx = time_locked_transfer_transaction( + sender_id, + &key1, + 0, + recipient_id, + system_accounts::clock_account_ids()[0], + amount, + deadline, + ); + + state + .transition_from_public_transaction(&tx, 2, clock_timestamp) + .unwrap(); + + // Balances changed. + assert_eq!(state.get_account_by_id(sender_id).balance, 0); + assert_eq!(state.get_account_by_id(recipient_id).balance, 100); +} + +#[test] +fn time_locked_transfer_fails_when_deadline_is_in_the_future() { + let clock_timestamp = 600; + let mut state = + state_with_clock_and_program(test_programs::time_locked_transfer(), clock_timestamp); + + let recipient_id = AccountId::new([42; 32]); + state.force_insert_account( + recipient_id, + Account { + program_owner: programs::authenticated_transfer().id(), + ..Account::default() + }, + ); + + let key1 = PrivateKey::try_new([1; 32]).unwrap(); + let sender_id = AccountId::from(&PublicKey::new_from_private_key(&key1)); + state.force_insert_account( + sender_id, + Account { + program_owner: test_programs::time_locked_transfer().id(), + balance: 100, + ..Account::default() + }, + ); + + let amount = 100; + // Far-future deadline: the program panics because the clock has not reached it. + let deadline = u64::MAX; + + let tx = time_locked_transfer_transaction( + sender_id, + &key1, + 0, + recipient_id, + system_accounts::clock_account_ids()[0], + amount, + deadline, + ); + + let result = state.transition_from_public_transaction(&tx, 2, clock_timestamp); + + assert!( + result.is_err(), + "Transfer should fail when deadline is in the future" + ); + // Balances unchanged. + assert_eq!(state.get_account_by_id(sender_id).balance, 100); + assert_eq!(state.get_account_by_id(recipient_id).balance, 0); +} + +fn pinata_cooldown_data(prize: u128, cooldown_ms: u64, last_claim_timestamp: u64) -> Vec { + let mut buf = Vec::with_capacity(32); + buf.extend_from_slice(&prize.to_le_bytes()); + buf.extend_from_slice(&cooldown_ms.to_le_bytes()); + buf.extend_from_slice(&last_claim_timestamp.to_le_bytes()); + buf +} + +fn pinata_cooldown_transaction( + pinata_id: AccountId, + winner_id: AccountId, + clock_account_id: AccountId, +) -> PublicTransaction { + let program_id = test_programs::pinata_cooldown().id(); + let message = lee::public_transaction::Message::try_new( + program_id, + vec![pinata_id, winner_id, clock_account_id], + vec![], + (), + ) + .unwrap(); + let witness_set = lee::public_transaction::WitnessSet::for_message(&message, &[]); + PublicTransaction::new(message, witness_set) +} + +#[test] +fn pinata_cooldown_claim_succeeds_after_cooldown() { + let winner_id = AccountId::new([11; 32]); + let pinata_id = AccountId::new([99; 32]); + + let genesis_timestamp = 1000; + let prize = 50; + let cooldown_ms = 500; + // Last claim was at genesis, so any timestamp >= genesis + cooldown should work. + let last_claim_timestamp = genesis_timestamp; + + // Advance the clock so the cooldown check reads an updated timestamp. + let block_timestamp = genesis_timestamp + cooldown_ms; + let mut state = state_with_clock_and_program(test_programs::pinata_cooldown(), block_timestamp); + + // The winner must be a non-default account so the program may credit it without claiming. + state.force_insert_account( + winner_id, + Account { + program_owner: programs::authenticated_transfer().id(), + ..Account::default() + }, + ); + state.force_insert_account( + pinata_id, + Account { + program_owner: test_programs::pinata_cooldown().id(), + balance: 1000, + data: pinata_cooldown_data(prize, cooldown_ms, last_claim_timestamp) + .try_into() + .unwrap(), + ..Account::default() + }, + ); + + let tx = pinata_cooldown_transaction( + pinata_id, + winner_id, + system_accounts::clock_account_ids()[0], + ); + + state + .transition_from_public_transaction(&tx, 2, block_timestamp) + .unwrap(); + + assert_eq!(state.get_account_by_id(pinata_id).balance, 1000 - prize); + assert_eq!(state.get_account_by_id(winner_id).balance, prize); +} + +#[test] +fn pinata_cooldown_claim_fails_during_cooldown() { + let winner_id = AccountId::new([11; 32]); + let pinata_id = AccountId::new([99; 32]); + + let genesis_timestamp = 1000; + let prize = 50; + let cooldown_ms = 500; + let last_claim_timestamp = genesis_timestamp; + + // Timestamp is only 100ms after the last claim, well within the 500ms cooldown. + let block_timestamp = genesis_timestamp + 100; + let mut state = state_with_clock_and_program(test_programs::pinata_cooldown(), block_timestamp); + + state.force_insert_account( + winner_id, + Account { + program_owner: programs::authenticated_transfer().id(), + ..Account::default() + }, + ); + state.force_insert_account( + pinata_id, + Account { + program_owner: test_programs::pinata_cooldown().id(), + balance: 1000, + data: pinata_cooldown_data(prize, cooldown_ms, last_claim_timestamp) + .try_into() + .unwrap(), + ..Account::default() + }, + ); + + let tx = pinata_cooldown_transaction( + pinata_id, + winner_id, + system_accounts::clock_account_ids()[0], + ); + + let result = state.transition_from_public_transaction(&tx, 2, block_timestamp); + + assert!(result.is_err(), "Claim should fail during cooldown period"); + assert_eq!(state.get_account_by_id(pinata_id).balance, 1000); + assert_eq!(state.get_account_by_id(winner_id).balance, 0); +} + +#[test] +fn pda_mechanism_with_pinata_token_program() { + let pinata_token = programs::pinata_token(); + let token = programs::token(); + + let pinata_definition_id = AccountId::new([1; 32]); + let pinata_token_definition_id = AccountId::new([2; 32]); + // Total supply of pinata token will be in an account under a PDA. + let pinata_token_holding_id = + AccountId::for_public_pda(&pinata_token.id(), &PdaSeed::new([0; 32])); + let winner_token_holding_id = AccountId::new([3; 32]); + + let expected_winner_account_holding = token_core::TokenHolding::Fungible { + definition_id: pinata_token_definition_id, + balance: 150, + }; + let expected_winner_token_holding_post = Account { + program_owner: token.id(), + data: Data::from(&expected_winner_account_holding), + ..Account::default() + }; + + // Register the pinata-token and token programs and create the pinata definition account. + // This replaces the removed `add_pinata_token_program` helper. + let mut state = V03State::new().with_programs([pinata_token.clone(), token.clone()]); + state.force_insert_account( + pinata_definition_id, + Account { + program_owner: pinata_token.id(), + // Difficulty: 3 + data: vec![3; 33].try_into().unwrap(), + ..Account::default() + }, + ); + + // Set up the token accounts directly (bypassing public transactions which + // would require signers for Claim::Authorized). The focus of this test is + // the PDA mechanism in the pinata program's chained call, not token creation. + let total_supply: u128 = 10_000_000; + let token_definition = token_core::TokenDefinition::Fungible { + name: String::from("PINATA"), + total_supply, + metadata_id: None, + }; + let token_holding = token_core::TokenHolding::Fungible { + definition_id: pinata_token_definition_id, + balance: total_supply, + }; + let winner_holding = token_core::TokenHolding::Fungible { + definition_id: pinata_token_definition_id, + balance: 0, + }; + state.force_insert_account( + pinata_token_definition_id, + Account { + program_owner: token.id(), + data: Data::from(&token_definition), + ..Account::default() + }, + ); + state.force_insert_account( + pinata_token_holding_id, + Account { + program_owner: token.id(), + data: Data::from(&token_holding), + ..Account::default() + }, + ); + state.force_insert_account( + winner_token_holding_id, + Account { + program_owner: token.id(), + data: Data::from(&winner_holding), + ..Account::default() + }, + ); + + // Submit a solution to the pinata program to claim the prize + let solution: u128 = 989_106; + let message = lee::public_transaction::Message::try_new( + pinata_token.id(), + vec![ + pinata_definition_id, + pinata_token_holding_id, + winner_token_holding_id, + ], + vec![], + solution, + ) + .unwrap(); + let witness_set = lee::public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + let winner_token_holding_post = state.get_account_by_id(winner_token_holding_id); + assert_eq!( + winner_token_holding_post, + expected_winner_token_holding_post + ); +} diff --git a/lez/storage/src/indexer/mod.rs b/lez/storage/src/indexer/mod.rs index a753a71a..87d586fb 100644 --- a/lez/storage/src/indexer/mod.rs +++ b/lez/storage/src/indexer/mod.rs @@ -153,97 +153,19 @@ impl RocksDBIO { } let br_id = closest_breakpoint_id(block_id); - let mut breakpoint = self.get_breakpoint(br_id)?; + let mut state = self.get_breakpoint(br_id)?; let start = u64::from(BREAKPOINT_INTERVAL) .checked_mul(br_id) .expect("Reached maximum breakpoint id"); - for mut block in self.get_block_batch_seq( + for block in self.get_block_batch_seq( start.checked_add(1).expect("Will be lesser that u64::MAX")..=block_id, )? { - let expected_clock = LeeTransaction::Public(clock_invocation(block.header.timestamp)); - - let clock_tx = block.body.transactions.pop().ok_or_else(|| { - DbError::db_interaction_error( - "Block must contain clock transaction at the end".to_owned(), - ) - })?; - let user_txs = block.body.transactions; - - if clock_tx != expected_clock { - return Err(DbError::db_interaction_error( - "Last transaction in block must be the clock invocation for the block timestamp" - .to_owned(), - )); - } - for transaction in user_txs { - let is_genesis = block.header.block_id == GENESIS_BLOCK_ID; - if is_genesis { - let genesis_tx = match transaction { - LeeTransaction::Public(public_tx) => public_tx, - LeeTransaction::PrivacyPreserving(_) - | LeeTransaction::ProgramDeployment(_) => { - return Err(DbError::db_interaction_error( - "Genesis block should contain only public transactions".to_owned(), - )); - } - }; - breakpoint - .transition_from_public_transaction( - &genesis_tx, - block.header.block_id, - block.header.timestamp, - ) - .map_err(|err| { - DbError::db_interaction_error(format!( - "genesis transaction execution failed with err {err:?}" - )) - })?; - } else { - transaction - .transaction_stateless_check() - .map_err(|err| { - DbError::db_interaction_error(format!( - "transaction pre check failed with err {err:?}" - )) - })? - // FIXME: HOT FIX (testnet v0.2): does not check for system account updates due to - // sequencer-generated deposit tx'es; - // CHANGE ME back to `execute_check_on_state` when the indexer can authenticate deposit transactions - .execute_without_system_accounts_check_on_state( - &mut breakpoint, - block.header.block_id, - block.header.timestamp, - ) - .map_err(|err| { - DbError::db_interaction_error(format!( - "transaction execution failed with err {err:?}" - )) - })?; - } - } - - let LeeTransaction::Public(clock_public_tx) = clock_tx else { - return Err(DbError::db_interaction_error( - "Clock invocation must be a public transaction".to_owned(), - )); - }; - - breakpoint - .transition_from_public_transaction( - &clock_public_tx, - block.header.block_id, - block.header.timestamp, - ) - .map_err(|err| { - DbError::db_interaction_error(format!( - "clock transaction execution failed with err {err:?}" - )) - })?; + apply_block_transactions(block, &mut state)?; } - Ok(breakpoint) + Ok(state) } pub fn final_state(&self) -> DbResult { @@ -252,6 +174,86 @@ impl RocksDBIO { } } +fn apply_block_transactions(mut block: Block, state: &mut V03State) -> DbResult<()> { + let expected_clock = LeeTransaction::Public(clock_invocation(block.header.timestamp)); + + let clock_tx = block.body.transactions.pop().ok_or_else(|| { + DbError::db_interaction_error("Block must contain clock transaction at the end".to_owned()) + })?; + + if clock_tx != expected_clock { + return Err(DbError::db_interaction_error( + "Last transaction in block must be the clock invocation for the block timestamp" + .to_owned(), + )); + } + + for transaction in block.body.transactions { + if block.header.block_id == GENESIS_BLOCK_ID { + let genesis_tx = match transaction { + LeeTransaction::Public(public_tx) => public_tx, + LeeTransaction::PrivacyPreserving(_) | LeeTransaction::ProgramDeployment(_) => { + return Err(DbError::db_interaction_error( + "Genesis block should contain only public transactions".to_owned(), + )); + } + }; + state + .transition_from_public_transaction( + &genesis_tx, + block.header.block_id, + block.header.timestamp, + ) + .map_err(|err| { + DbError::db_interaction_error(format!( + "genesis transaction execution failed with err {err:?}" + )) + })?; + } else { + transaction + .transaction_stateless_check() + .map_err(|err| { + DbError::db_interaction_error(format!( + "transaction pre check failed with err {err:?}" + )) + })? + // FIXME: HOT FIX (testnet v0.2): does not check for system account updates due to + // sequencer-generated deposit tx'es; + // CHANGE ME back to `execute_check_on_state` when the indexer can authenticate deposit transactions + .execute_without_system_accounts_check_on_state( + state, + block.header.block_id, + block.header.timestamp, + ) + .map_err(|err| { + DbError::db_interaction_error(format!( + "transaction execution failed with err {err:?}" + )) + })?; + } + } + + let LeeTransaction::Public(clock_public_tx) = clock_tx else { + return Err(DbError::db_interaction_error( + "Clock invocation must be a public transaction".to_owned(), + )); + }; + + state + .transition_from_public_transaction( + &clock_public_tx, + block.header.block_id, + block.header.timestamp, + ) + .map_err(|err| { + DbError::db_interaction_error(format!( + "clock transaction execution failed with err {err:?}" + )) + })?; + + Ok(()) +} + fn closest_breakpoint_id(block_id: u64) -> u64 { block_id .saturating_sub(1) @@ -261,439 +263,4 @@ fn closest_breakpoint_id(block_id: u64) -> u64 { #[expect(clippy::shadow_unrelated, reason = "Fine for tests")] #[cfg(test)] -mod tests { - use common::test_utils::produce_dummy_block; - use lee::{Account, AccountId, PublicKey}; - use tempfile::tempdir; - - use super::*; - - fn genesis_block() -> Block { - produce_dummy_block(1, None, vec![]) - } - - fn acc1_sign_key() -> lee::PrivateKey { - lee::PrivateKey::try_new([1; 32]).unwrap() - } - - fn acc2_sign_key() -> lee::PrivateKey { - lee::PrivateKey::try_new([2; 32]).unwrap() - } - - fn acc1() -> AccountId { - AccountId::from(&PublicKey::new_from_private_key(&acc1_sign_key())) - } - - fn acc2() -> AccountId { - AccountId::from(&PublicKey::new_from_private_key(&acc2_sign_key())) - } - - fn initial_state() -> lee::V03State { - let mut public_accounts = [(acc1(), 10000), (acc2(), 20000)] - .into_iter() - .map(|(id, balance)| { - ( - id, - Account { - program_owner: programs::authenticated_transfer().id(), - balance, - ..Account::default() - }, - ) - }) - .collect::>(); - for clock_id in system_accounts::clock_account_ids() { - public_accounts.push((clock_id, system_accounts::clock_account())); - } - - lee::V03State::new() - .with_public_accounts(public_accounts) - .with_programs([programs::authenticated_transfer(), programs::clock()]) - } - - #[test] - fn start_db() { - let temp_dir = tempdir().unwrap(); - let temdir_path = temp_dir.path(); - - let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap(); - let first_id = dbio.get_meta_first_block_id_in_db().unwrap(); - let is_first_set = dbio.get_meta_is_first_block_set().unwrap(); - let last_observed_l1_header = dbio.get_meta_last_observed_l1_lib_header_in_db().unwrap(); - let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap(); - let last_block = dbio.get_block(1).unwrap(); - let breakpoint = dbio.get_breakpoint(0).unwrap(); - let final_state = dbio.final_state().unwrap(); - - assert_eq!(last_id, None); - assert_eq!(first_id, None); - assert_eq!(last_observed_l1_header, None); - assert!(!is_first_set); - assert_eq!(last_br_id, Some(0)); // TODO: Will be None after we remove hardcoded testnet state - assert!(last_block.is_none()); - assert_eq!( - breakpoint.get_account_by_id(acc1()), - final_state.get_account_by_id(acc1()) - ); - assert_eq!( - breakpoint.get_account_by_id(acc2()), - final_state.get_account_by_id(acc2()) - ); - } - - #[test] - fn one_block_insertion() { - let temp_dir = tempdir().unwrap(); - let temdir_path = temp_dir.path(); - - let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); - - let genesis_block = genesis_block(); - dbio.put_block(&genesis_block, [0; 32]).unwrap(); - - let prev_hash = genesis_block.header.hash; - let from = acc1(); - let to = acc2(); - let sign_key = acc1_sign_key(); - - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); - let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]); - - dbio.put_block(&block, [1; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let first_id = dbio.get_meta_first_block_id_in_db().unwrap(); - let last_observed_l1_header = dbio - .get_meta_last_observed_l1_lib_header_in_db() - .unwrap() - .unwrap(); - let is_first_set = dbio.get_meta_is_first_block_set().unwrap(); - let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - let breakpoint = dbio.get_breakpoint(0).unwrap(); - let final_state = dbio.final_state().unwrap(); - - assert_eq!(last_id, 2); - assert_eq!(first_id, Some(1)); - assert_eq!(last_observed_l1_header, [1; 32]); - assert!(is_first_set); - assert_eq!(last_br_id, Some(0)); - assert_eq!(last_block.header.hash, block.header.hash); - assert_eq!( - breakpoint.get_account_by_id(acc1()).balance - - final_state.get_account_by_id(acc1()).balance, - 1 - ); - assert_eq!( - final_state.get_account_by_id(acc2()).balance - - breakpoint.get_account_by_id(acc2()).balance, - 1 - ); - } - - #[test] - fn new_breakpoint() { - let temp_dir = tempdir().unwrap(); - let temdir_path = temp_dir.path(); - - let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); - - let from = acc1(); - let to = acc2(); - let sign_key = acc1_sign_key(); - - for i in 1..=BREAKPOINT_INTERVAL + 1 { - let prev_hash = dbio.get_meta_last_block_id_in_db().unwrap().map(|last_id| { - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - last_block.header.hash - }); - - let transfer_tx = common::test_utils::create_transaction_native_token_transfer( - from, - (i - 1).into(), - to, - 1, - &sign_key, - ); - let block = produce_dummy_block(i.into(), prev_hash, vec![transfer_tx]); - dbio.put_block(&block, [i; 32]).unwrap(); - } - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let first_id = dbio.get_meta_first_block_id_in_db().unwrap(); - let is_first_set = dbio.get_meta_is_first_block_set().unwrap(); - let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - let prev_breakpoint = dbio.get_breakpoint(0).unwrap(); - let breakpoint = dbio.get_breakpoint(1).unwrap(); - let final_state = dbio.final_state().unwrap(); - - assert_eq!(last_id, 101); - assert_eq!(first_id, Some(1)); - assert!(is_first_set); - assert_eq!(last_br_id, Some(1)); - assert_ne!(last_block.header.hash, genesis_block().header.hash); - assert_eq!( - prev_breakpoint.get_account_by_id(acc1()).balance - - final_state.get_account_by_id(acc1()).balance, - 101 - ); - assert_eq!( - final_state.get_account_by_id(acc2()).balance - - prev_breakpoint.get_account_by_id(acc2()).balance, - 101 - ); - assert_eq!( - breakpoint.get_account_by_id(acc1()).balance - - final_state.get_account_by_id(acc1()).balance, - 1 - ); - assert_eq!( - final_state.get_account_by_id(acc2()).balance - - breakpoint.get_account_by_id(acc2()).balance, - 1 - ); - } - - #[test] - fn simple_maps() { - let temp_dir = tempdir().unwrap(); - let temdir_path = temp_dir.path(); - - let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); - - let from = acc1(); - let to = acc2(); - let sign_key = acc1_sign_key(); - - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); - let block = produce_dummy_block(1, None, vec![transfer_tx]); - - let control_hash1 = block.header.hash; - - dbio.put_block(&block, [1; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key); - let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]); - - let control_hash2 = block.header.hash; - - dbio.put_block(&block, [2; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key); - - let control_tx_hash1 = transfer_tx.hash(); - - let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]); - dbio.put_block(&block, [3; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key); - - let control_tx_hash2 = transfer_tx.hash(); - - let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]); - dbio.put_block(&block, [4; 32]).unwrap(); - - let control_block_id1 = dbio.get_block_id_by_hash(control_hash1.0).unwrap().unwrap(); - let control_block_id2 = dbio.get_block_id_by_hash(control_hash2.0).unwrap().unwrap(); - let control_block_id3 = dbio - .get_block_id_by_tx_hash(control_tx_hash1.0) - .unwrap() - .unwrap(); - let control_block_id4 = dbio - .get_block_id_by_tx_hash(control_tx_hash2.0) - .unwrap() - .unwrap(); - - assert_eq!(control_block_id1, 1); - assert_eq!(control_block_id2, 2); - assert_eq!(control_block_id3, 3); - assert_eq!(control_block_id4, 4); - } - - #[test] - fn block_batch() { - let temp_dir = tempdir().unwrap(); - let temdir_path = temp_dir.path(); - - let mut block_res = vec![]; - - let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); - - let from = acc1(); - let to = acc2(); - let sign_key = acc1_sign_key(); - - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); - let block = produce_dummy_block(1, None, vec![transfer_tx]); - - block_res.push(block.clone()); - dbio.put_block(&block, [1; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key); - let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]); - - block_res.push(block.clone()); - dbio.put_block(&block, [2; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key); - - let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]); - block_res.push(block.clone()); - dbio.put_block(&block, [3; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key); - - let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]); - block_res.push(block.clone()); - dbio.put_block(&block, [4; 32]).unwrap(); - - let block_hashes_mem: Vec<[u8; 32]> = - block_res.into_iter().map(|bl| bl.header.hash.0).collect(); - - // Get blocks before ID 5 (i.e., starting from 4 going backwards), limit 4 - // This should return blocks 4, 3, 2, 1 in descending order - let mut batch_res = dbio.get_block_batch(Some(5), 4).unwrap(); - batch_res.reverse(); // Reverse to match ascending order for comparison - - let block_hashes_db: Vec<[u8; 32]> = - batch_res.into_iter().map(|bl| bl.header.hash.0).collect(); - - assert_eq!(block_hashes_mem, block_hashes_db); - - let block_hashes_mem_limited = &block_hashes_mem[1..]; - - // Get blocks before ID 5, limit 3 - // This should return blocks 4, 3, 2 in descending order - let mut batch_res_limited = dbio.get_block_batch(Some(5), 3).unwrap(); - batch_res_limited.reverse(); // Reverse to match ascending order for comparison - - let block_hashes_db_limited: Vec<[u8; 32]> = batch_res_limited - .into_iter() - .map(|bl| bl.header.hash.0) - .collect(); - - assert_eq!(block_hashes_mem_limited, block_hashes_db_limited.as_slice()); - - let block_batch_seq = dbio.get_block_batch_seq(1..=5).unwrap(); - let block_batch_ids = block_batch_seq - .into_iter() - .map(|block| block.header.block_id) - .collect::>(); - - assert_eq!(block_batch_ids, vec![1, 2, 3, 4]); - } - - #[test] - fn account_map() { - let temp_dir = tempdir().unwrap(); - let temdir_path = temp_dir.path(); - - let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); - - let from = acc1(); - let to = acc2(); - let sign_key = acc1_sign_key(); - - let mut tx_hash_res = vec![]; - - let transfer_tx1 = - common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); - let transfer_tx2 = - common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key); - tx_hash_res.push(transfer_tx1.hash().0); - tx_hash_res.push(transfer_tx2.hash().0); - - let block = produce_dummy_block(1, None, vec![transfer_tx1, transfer_tx2]); - - dbio.put_block(&block, [1; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx1 = - common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key); - let transfer_tx2 = - common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key); - tx_hash_res.push(transfer_tx1.hash().0); - tx_hash_res.push(transfer_tx2.hash().0); - - let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx1, transfer_tx2]); - - dbio.put_block(&block, [2; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx1 = - common::test_utils::create_transaction_native_token_transfer(from, 4, to, 1, &sign_key); - let transfer_tx2 = - common::test_utils::create_transaction_native_token_transfer(from, 5, to, 1, &sign_key); - tx_hash_res.push(transfer_tx1.hash().0); - tx_hash_res.push(transfer_tx2.hash().0); - - let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx1, transfer_tx2]); - - dbio.put_block(&block, [3; 32]).unwrap(); - - let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); - let last_block = dbio.get_block(last_id).unwrap().unwrap(); - - let prev_hash = last_block.header.hash; - let transfer_tx = - common::test_utils::create_transaction_native_token_transfer(from, 6, to, 1, &sign_key); - tx_hash_res.push(transfer_tx.hash().0); - - let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]); - - dbio.put_block(&block, [4; 32]).unwrap(); - - let acc1_tx = dbio.get_acc_transactions(*acc1().value(), 0, 7).unwrap(); - let acc1_tx_hashes: Vec<[u8; 32]> = acc1_tx.into_iter().map(|tx| tx.hash().0).collect(); - - assert_eq!(acc1_tx_hashes, tx_hash_res); - - let acc1_tx_limited = dbio.get_acc_transactions(*acc1().value(), 1, 4).unwrap(); - let acc1_tx_limited_hashes: Vec<[u8; 32]> = - acc1_tx_limited.into_iter().map(|tx| tx.hash().0).collect(); - - assert_eq!(acc1_tx_limited_hashes.as_slice(), &tx_hash_res[1..5]); - } -} +mod tests; diff --git a/lez/storage/src/indexer/tests.rs b/lez/storage/src/indexer/tests.rs new file mode 100644 index 00000000..44df23d5 --- /dev/null +++ b/lez/storage/src/indexer/tests.rs @@ -0,0 +1,433 @@ +use common::test_utils::produce_dummy_block; +use lee::{Account, AccountId, PublicKey}; +use tempfile::tempdir; + +use super::*; + +fn genesis_block() -> Block { + produce_dummy_block(1, None, vec![]) +} + +fn acc1_sign_key() -> lee::PrivateKey { + lee::PrivateKey::try_new([1; 32]).unwrap() +} + +fn acc2_sign_key() -> lee::PrivateKey { + lee::PrivateKey::try_new([2; 32]).unwrap() +} + +fn acc1() -> AccountId { + AccountId::from(&PublicKey::new_from_private_key(&acc1_sign_key())) +} + +fn acc2() -> AccountId { + AccountId::from(&PublicKey::new_from_private_key(&acc2_sign_key())) +} + +fn initial_state() -> lee::V03State { + let mut public_accounts = [(acc1(), 10000), (acc2(), 20000)] + .into_iter() + .map(|(id, balance)| { + ( + id, + Account { + program_owner: programs::authenticated_transfer().id(), + balance, + ..Account::default() + }, + ) + }) + .collect::>(); + for clock_id in system_accounts::clock_account_ids() { + public_accounts.push((clock_id, system_accounts::clock_account())); + } + + lee::V03State::new() + .with_public_accounts(public_accounts) + .with_programs([programs::authenticated_transfer(), programs::clock()]) +} + +#[test] +fn start_db() { + let temp_dir = tempdir().unwrap(); + let temdir_path = temp_dir.path(); + + let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap(); + let first_id = dbio.get_meta_first_block_id_in_db().unwrap(); + let is_first_set = dbio.get_meta_is_first_block_set().unwrap(); + let last_observed_l1_header = dbio.get_meta_last_observed_l1_lib_header_in_db().unwrap(); + let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap(); + let last_block = dbio.get_block(1).unwrap(); + let breakpoint = dbio.get_breakpoint(0).unwrap(); + let final_state = dbio.final_state().unwrap(); + + assert_eq!(last_id, None); + assert_eq!(first_id, None); + assert_eq!(last_observed_l1_header, None); + assert!(!is_first_set); + assert_eq!(last_br_id, Some(0)); // TODO: Will be None after we remove hardcoded testnet state + assert!(last_block.is_none()); + assert_eq!( + breakpoint.get_account_by_id(acc1()), + final_state.get_account_by_id(acc1()) + ); + assert_eq!( + breakpoint.get_account_by_id(acc2()), + final_state.get_account_by_id(acc2()) + ); +} + +#[test] +fn one_block_insertion() { + let temp_dir = tempdir().unwrap(); + let temdir_path = temp_dir.path(); + + let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); + + let genesis_block = genesis_block(); + dbio.put_block(&genesis_block, [0; 32]).unwrap(); + + let prev_hash = genesis_block.header.hash; + let from = acc1(); + let to = acc2(); + let sign_key = acc1_sign_key(); + + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); + let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]); + + dbio.put_block(&block, [1; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let first_id = dbio.get_meta_first_block_id_in_db().unwrap(); + let last_observed_l1_header = dbio + .get_meta_last_observed_l1_lib_header_in_db() + .unwrap() + .unwrap(); + let is_first_set = dbio.get_meta_is_first_block_set().unwrap(); + let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + let breakpoint = dbio.get_breakpoint(0).unwrap(); + let final_state = dbio.final_state().unwrap(); + + assert_eq!(last_id, 2); + assert_eq!(first_id, Some(1)); + assert_eq!(last_observed_l1_header, [1; 32]); + assert!(is_first_set); + assert_eq!(last_br_id, Some(0)); + assert_eq!(last_block.header.hash, block.header.hash); + assert_eq!( + breakpoint.get_account_by_id(acc1()).balance + - final_state.get_account_by_id(acc1()).balance, + 1 + ); + assert_eq!( + final_state.get_account_by_id(acc2()).balance + - breakpoint.get_account_by_id(acc2()).balance, + 1 + ); +} + +#[test] +fn new_breakpoint() { + let temp_dir = tempdir().unwrap(); + let temdir_path = temp_dir.path(); + + let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); + + let from = acc1(); + let to = acc2(); + let sign_key = acc1_sign_key(); + + for i in 1..=BREAKPOINT_INTERVAL + 1 { + let prev_hash = dbio.get_meta_last_block_id_in_db().unwrap().map(|last_id| { + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + last_block.header.hash + }); + + let transfer_tx = common::test_utils::create_transaction_native_token_transfer( + from, + (i - 1).into(), + to, + 1, + &sign_key, + ); + let block = produce_dummy_block(i.into(), prev_hash, vec![transfer_tx]); + dbio.put_block(&block, [i; 32]).unwrap(); + } + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let first_id = dbio.get_meta_first_block_id_in_db().unwrap(); + let is_first_set = dbio.get_meta_is_first_block_set().unwrap(); + let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + let prev_breakpoint = dbio.get_breakpoint(0).unwrap(); + let breakpoint = dbio.get_breakpoint(1).unwrap(); + let final_state = dbio.final_state().unwrap(); + + assert_eq!(last_id, 101); + assert_eq!(first_id, Some(1)); + assert!(is_first_set); + assert_eq!(last_br_id, Some(1)); + assert_ne!(last_block.header.hash, genesis_block().header.hash); + assert_eq!( + prev_breakpoint.get_account_by_id(acc1()).balance + - final_state.get_account_by_id(acc1()).balance, + 101 + ); + assert_eq!( + final_state.get_account_by_id(acc2()).balance + - prev_breakpoint.get_account_by_id(acc2()).balance, + 101 + ); + assert_eq!( + breakpoint.get_account_by_id(acc1()).balance + - final_state.get_account_by_id(acc1()).balance, + 1 + ); + assert_eq!( + final_state.get_account_by_id(acc2()).balance + - breakpoint.get_account_by_id(acc2()).balance, + 1 + ); +} + +#[test] +fn simple_maps() { + let temp_dir = tempdir().unwrap(); + let temdir_path = temp_dir.path(); + + let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); + + let from = acc1(); + let to = acc2(); + let sign_key = acc1_sign_key(); + + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); + let block = produce_dummy_block(1, None, vec![transfer_tx]); + + let control_hash1 = block.header.hash; + + dbio.put_block(&block, [1; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key); + let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]); + + let control_hash2 = block.header.hash; + + dbio.put_block(&block, [2; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key); + + let control_tx_hash1 = transfer_tx.hash(); + + let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]); + dbio.put_block(&block, [3; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key); + + let control_tx_hash2 = transfer_tx.hash(); + + let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]); + dbio.put_block(&block, [4; 32]).unwrap(); + + let control_block_id1 = dbio.get_block_id_by_hash(control_hash1.0).unwrap().unwrap(); + let control_block_id2 = dbio.get_block_id_by_hash(control_hash2.0).unwrap().unwrap(); + let control_block_id3 = dbio + .get_block_id_by_tx_hash(control_tx_hash1.0) + .unwrap() + .unwrap(); + let control_block_id4 = dbio + .get_block_id_by_tx_hash(control_tx_hash2.0) + .unwrap() + .unwrap(); + + assert_eq!(control_block_id1, 1); + assert_eq!(control_block_id2, 2); + assert_eq!(control_block_id3, 3); + assert_eq!(control_block_id4, 4); +} + +#[test] +fn block_batch() { + let temp_dir = tempdir().unwrap(); + let temdir_path = temp_dir.path(); + + let mut block_res = vec![]; + + let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); + + let from = acc1(); + let to = acc2(); + let sign_key = acc1_sign_key(); + + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); + let block = produce_dummy_block(1, None, vec![transfer_tx]); + + block_res.push(block.clone()); + dbio.put_block(&block, [1; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key); + let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]); + + block_res.push(block.clone()); + dbio.put_block(&block, [2; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key); + + let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]); + block_res.push(block.clone()); + dbio.put_block(&block, [3; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key); + + let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]); + block_res.push(block.clone()); + dbio.put_block(&block, [4; 32]).unwrap(); + + let block_hashes_mem: Vec<[u8; 32]> = + block_res.into_iter().map(|bl| bl.header.hash.0).collect(); + + // Get blocks before ID 5 (i.e., starting from 4 going backwards), limit 4 + // This should return blocks 4, 3, 2, 1 in descending order + let mut batch_res = dbio.get_block_batch(Some(5), 4).unwrap(); + batch_res.reverse(); // Reverse to match ascending order for comparison + + let block_hashes_db: Vec<[u8; 32]> = batch_res.into_iter().map(|bl| bl.header.hash.0).collect(); + + assert_eq!(block_hashes_mem, block_hashes_db); + + let block_hashes_mem_limited = &block_hashes_mem[1..]; + + // Get blocks before ID 5, limit 3 + // This should return blocks 4, 3, 2 in descending order + let mut batch_res_limited = dbio.get_block_batch(Some(5), 3).unwrap(); + batch_res_limited.reverse(); // Reverse to match ascending order for comparison + + let block_hashes_db_limited: Vec<[u8; 32]> = batch_res_limited + .into_iter() + .map(|bl| bl.header.hash.0) + .collect(); + + assert_eq!(block_hashes_mem_limited, block_hashes_db_limited.as_slice()); + + let block_batch_seq = dbio.get_block_batch_seq(1..=5).unwrap(); + let block_batch_ids = block_batch_seq + .into_iter() + .map(|block| block.header.block_id) + .collect::>(); + + assert_eq!(block_batch_ids, vec![1, 2, 3, 4]); +} + +#[test] +fn account_map() { + let temp_dir = tempdir().unwrap(); + let temdir_path = temp_dir.path(); + + let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap(); + + let from = acc1(); + let to = acc2(); + let sign_key = acc1_sign_key(); + + let mut tx_hash_res = vec![]; + + let transfer_tx1 = + common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key); + let transfer_tx2 = + common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key); + tx_hash_res.push(transfer_tx1.hash().0); + tx_hash_res.push(transfer_tx2.hash().0); + + let block = produce_dummy_block(1, None, vec![transfer_tx1, transfer_tx2]); + + dbio.put_block(&block, [1; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx1 = + common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key); + let transfer_tx2 = + common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key); + tx_hash_res.push(transfer_tx1.hash().0); + tx_hash_res.push(transfer_tx2.hash().0); + + let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx1, transfer_tx2]); + + dbio.put_block(&block, [2; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx1 = + common::test_utils::create_transaction_native_token_transfer(from, 4, to, 1, &sign_key); + let transfer_tx2 = + common::test_utils::create_transaction_native_token_transfer(from, 5, to, 1, &sign_key); + tx_hash_res.push(transfer_tx1.hash().0); + tx_hash_res.push(transfer_tx2.hash().0); + + let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx1, transfer_tx2]); + + dbio.put_block(&block, [3; 32]).unwrap(); + + let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap(); + let last_block = dbio.get_block(last_id).unwrap().unwrap(); + + let prev_hash = last_block.header.hash; + let transfer_tx = + common::test_utils::create_transaction_native_token_transfer(from, 6, to, 1, &sign_key); + tx_hash_res.push(transfer_tx.hash().0); + + let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]); + + dbio.put_block(&block, [4; 32]).unwrap(); + + let acc1_tx = dbio.get_acc_transactions(*acc1().value(), 0, 7).unwrap(); + let acc1_tx_hashes: Vec<[u8; 32]> = acc1_tx.into_iter().map(|tx| tx.hash().0).collect(); + + assert_eq!(acc1_tx_hashes, tx_hash_res); + + let acc1_tx_limited = dbio.get_acc_transactions(*acc1().value(), 1, 4).unwrap(); + let acc1_tx_limited_hashes: Vec<[u8; 32]> = + acc1_tx_limited.into_iter().map(|tx| tx.hash().0).collect(); + + assert_eq!(acc1_tx_limited_hashes.as_slice(), &tx_hash_res[1..5]); +} diff --git a/lez/storage/src/indexer/write_atomic.rs b/lez/storage/src/indexer/write_atomic.rs index a88e46ef..ec28f8ec 100644 --- a/lez/storage/src/indexer/write_atomic.rs +++ b/lez/storage/src/indexer/write_atomic.rs @@ -52,44 +52,8 @@ impl RocksDBIO { acc_id: [u8; 32], tx_hashes: &[[u8; 32]], ) -> DbResult<()> { - let acc_num_tx = self.get_acc_meta_num_tx(acc_id)?.unwrap_or(0); - let cf_att = self.account_id_to_tx_hash_column(); let mut write_batch = WriteBatch::new(); - - for (tx_id, tx_hash) in tx_hashes.iter().enumerate() { - let put_id = acc_num_tx - .checked_add(tx_id.try_into().expect("Must fit into u64")) - .expect("Tx count should be lesser that u64::MAX"); - - let mut prefix = borsh::to_vec(&acc_id).map_err(|berr| { - DbError::borsh_cast_message(berr, Some("Failed to serialize account id".to_owned())) - })?; - let suffix = borsh::to_vec(&put_id).map_err(|berr| { - DbError::borsh_cast_message(berr, Some("Failed to serialize tx id".to_owned())) - })?; - - prefix.extend_from_slice(&suffix); - - write_batch.put_cf( - &cf_att, - prefix, - borsh::to_vec(tx_hash).map_err(|berr| { - DbError::borsh_cast_message( - berr, - Some("Failed to serialize tx hash".to_owned()), - ) - })?, - ); - } - - self.update_acc_meta_batch( - acc_id, - acc_num_tx - .checked_add(tx_hashes.len().try_into().expect("Must fit into u64")) - .expect("Tx count should be lesser that u64::MAX"), - &mut write_batch, - )?; - + self.put_account_transactions_dependant(acc_id, tx_hashes, &mut write_batch)?; self.db.write(write_batch).map_err(|rerr| { DbError::rocksdb_cast_message(rerr, Some("Failed to write batch".to_owned())) }) diff --git a/lez/wallet-ffi/src/label.rs b/lez/wallet-ffi/src/label.rs new file mode 100644 index 00000000..1f6f1236 --- /dev/null +++ b/lez/wallet-ffi/src/label.rs @@ -0,0 +1,316 @@ +use std::{ + ffi::{c_char, CString}, + str::FromStr as _, +}; + +use crate::{ + c_str_to_string, + error::{print_error, WalletFfiError}, + wallet::get_wallet, + FfiAccountIdWithPrivacy, WalletHandle, +}; + +#[repr(C)] +pub struct LabelAvailability { + pub is_available: bool, + pub error: WalletFfiError, +} + +impl LabelAvailability { + #[must_use] + pub const fn availability(is_available: bool) -> Self { + Self { + is_available, + error: WalletFfiError::Success, + } + } + + #[must_use] + pub const fn error(error: WalletFfiError) -> Self { + Self { + is_available: false, + error, + } + } +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct AccountIdResolvedFromLabel { + pub account_id: FfiAccountIdWithPrivacy, + pub error: WalletFfiError, +} + +impl AccountIdResolvedFromLabel { + #[must_use] + pub const fn account_id(account_id: FfiAccountIdWithPrivacy) -> Self { + Self { + account_id, + error: WalletFfiError::Success, + } + } + + #[must_use] + pub fn error(error: WalletFfiError) -> Self { + Self { + account_id: FfiAccountIdWithPrivacy::default(), + error, + } + } +} + +#[repr(C)] +pub struct LabelList { + pub labels_data: *mut *const c_char, + pub labels_size: usize, + pub error: WalletFfiError, +} + +impl LabelList { + #[must_use] + pub fn from_labels(labels: Vec<*const c_char>) -> Self { + let labels_size = labels.len(); + let boxed_slice = labels.into_boxed_slice(); + let labels_data = Box::into_raw(boxed_slice).cast::<*const c_char>(); + + Self { + labels_data, + labels_size, + error: WalletFfiError::Success, + } + } + + #[must_use] + pub const fn error(error: WalletFfiError) -> Self { + Self { + labels_data: std::ptr::null_mut(), + labels_size: 0, + error, + } + } +} + +/// Check if label is available. +/// +/// # Parameters +/// - `handle`: Valid wallet handle +/// - `label`: Input null terminated C string for a label +/// +/// # Returns +/// - `LabelAvailability` struct +/// +/// # Safety +/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` +/// - `label` must be a valid pointer to a null-terminated C string +#[no_mangle] +pub unsafe extern "C" fn wallet_ffi_check_label_available( + handle: *mut WalletHandle, + label: *const c_char, +) -> LabelAvailability { + let wrapper = match get_wallet(handle) { + Ok(w) => w, + Err(e) => return LabelAvailability::error(e), + }; + + let label = match c_str_to_string(label, "label") { + Ok(value) => value, + Err(e) => return LabelAvailability::error(e), + }; + + let wallet = match wrapper.core.lock() { + Ok(w) => w, + Err(e) => { + print_error(format!("Failed to lock wallet: {e}")); + return LabelAvailability::error(WalletFfiError::InternalError); + } + }; + + let is_available = wallet + .storage() + .check_label_availability(&label.into()) + .is_ok(); + + LabelAvailability::availability(is_available) +} + +/// Add new label. +/// +/// # Parameters +/// - `handle`: Valid wallet handle +/// - `label`: Input null terminated C string for a label +/// - `account_id_with_privacy`: The account ID (32 bytes) and its privacy. +/// +/// # Returns +/// - `Success` on successful query +/// - Error code on failure +/// +/// # Safety +/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` +/// - `label` must be a valid pointer to a null-terminated C string +#[no_mangle] +pub unsafe extern "C" fn wallet_ffi_add_label( + handle: *mut WalletHandle, + label: *const c_char, + account_id_with_privacy: FfiAccountIdWithPrivacy, +) -> WalletFfiError { + let wrapper = match get_wallet(handle) { + Ok(w) => w, + Err(e) => return e, + }; + + let label = match c_str_to_string(label, "label") { + Ok(value) => value, + Err(e) => return e, + }; + + let mut wallet = match wrapper.core.lock() { + Ok(w) => w, + Err(e) => { + print_error(format!("Failed to lock wallet: {e}")); + return WalletFfiError::InternalError; + } + }; + + match wallet + .storage_mut() + .add_label(label.into(), account_id_with_privacy.into()) + { + Ok(()) => WalletFfiError::Success, + Err(err) => { + print_error(format!("Failed to add label : {err}")); + WalletFfiError::InternalError + } + } +} + +/// Resolve a label. +/// +/// # Parameters +/// - `handle`: Valid wallet handle +/// - `label`: Input null terminated C string for a label +/// +/// # Returns +/// - `AccountIdResolvedFromLabel` struct +/// +/// # Safety +/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` +/// - `label` must be a valid pointer to a null-terminated C string +#[no_mangle] +pub unsafe extern "C" fn wallet_ffi_resolve_label( + handle: *mut WalletHandle, + label: *const c_char, +) -> AccountIdResolvedFromLabel { + let wrapper = match get_wallet(handle) { + Ok(w) => w, + Err(e) => return AccountIdResolvedFromLabel::error(e), + }; + + let label = match c_str_to_string(label, "label") { + Ok(value) => value, + Err(e) => return AccountIdResolvedFromLabel::error(e), + }; + + let mut wallet = match wrapper.core.lock() { + Ok(w) => w, + Err(e) => { + print_error(format!("Failed to lock wallet: {e}")); + return AccountIdResolvedFromLabel::error(WalletFfiError::InternalError); + } + }; + + wallet + .storage_mut() + .resolve_label(&label.into()) + .map_or_else( + || { + print_error("Failed to resolve label"); + AccountIdResolvedFromLabel::error(WalletFfiError::InternalError) + }, + |acc_id| AccountIdResolvedFromLabel::account_id(acc_id.into()), + ) +} + +/// Get all labels for account. +/// +/// # Parameters +/// - `handle`: Valid wallet handle +/// - `account_id_with_privacy`: The account ID (32 bytes) and its privacy. +/// +/// # Returns +/// - `LabelList` struct +/// +/// # Safety +/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` +#[no_mangle] +pub unsafe extern "C" fn wallet_ffi_get_all_labels_for_account( + handle: *mut WalletHandle, + account_id_with_privacy: FfiAccountIdWithPrivacy, +) -> LabelList { + let wrapper = match get_wallet(handle) { + Ok(w) => w, + Err(e) => return LabelList::error(e), + }; + + let wallet = match wrapper.core.lock() { + Ok(w) => w, + Err(e) => { + print_error(format!("Failed to lock wallet: {e}")); + return LabelList::error(WalletFfiError::InternalError); + } + }; + + let mut labels = vec![]; + + for label in wallet + .storage() + .labels_for_account(account_id_with_privacy.into()) + { + let Ok(label_c) = CString::from_str(label.as_ref()) else { + print_error(format!("Failed to cast label into C string: {label}")); + return LabelList::error(WalletFfiError::InternalError); + }; + + let label_raw = label_c.into_raw().cast_const(); + + labels.push(label_raw); + } + + LabelList::from_labels(labels) +} + +/// Free label list. +/// +/// # Parameters +/// - `label_list`: Input list of labels +/// +/// # Returns +/// - `Success` on successful query +/// - Error code on failure +/// +/// # Safety +/// - `label_list` must be a valid pointer to `LabelList`, received from +/// `wallet_ffi_get_all_labels_for_account` +#[no_mangle] +pub unsafe extern "C" fn wallet_ffi_free_label_list(label_list: *mut LabelList) -> WalletFfiError { + if label_list.is_null() { + return WalletFfiError::NullPointer; + } + + let labels_raw = unsafe { &*label_list }; + + if !labels_raw.labels_data.is_null() && labels_raw.labels_size > 0 { + let labels_slice = + std::slice::from_raw_parts_mut(labels_raw.labels_data, labels_raw.labels_size); + + for label_ptr in labels_slice.iter() { + if !(*label_ptr).is_null() { + drop(CString::from_raw((*label_ptr).cast_mut())); + } + } + + let boxed_slice = Box::from_raw(std::ptr::from_mut::<[*const c_char]>(labels_slice)); + drop(boxed_slice); + } + + WalletFfiError::Success +} diff --git a/lez/wallet-ffi/src/lib.rs b/lez/wallet-ffi/src/lib.rs index 6c86b0c8..361ae672 100644 --- a/lez/wallet-ffi/src/lib.rs +++ b/lez/wallet-ffi/src/lib.rs @@ -46,6 +46,7 @@ pub mod bridge; pub mod error; pub mod generic_transaction; pub mod keys; +pub mod label; pub mod pinata; pub mod program_deployment; pub mod sync; diff --git a/lez/wallet-ffi/src/types.rs b/lez/wallet-ffi/src/types.rs index cbe9fab7..98590619 100644 --- a/lez/wallet-ffi/src/types.rs +++ b/lez/wallet-ffi/src/types.rs @@ -9,7 +9,7 @@ use std::{ use lee::{Data, ProgramId, SharedSecretKey}; use lee_core::{encryption::MlKem768EncapsulationKey, NullifierPublicKey}; -use wallet::AccountIdentity; +use wallet::{account::AccountIdWithPrivacy, AccountIdentity}; use crate::error::WalletFfiError; @@ -24,7 +24,7 @@ pub struct WalletHandle { /// 32-byte array type for `AccountId`, keys, hashes, etc. #[repr(C)] -#[derive(Clone, Copy, Default)] +#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)] pub struct FfiBytes32 { pub data: [u8; 32], } @@ -593,6 +593,38 @@ impl From for ProgramId { } } +#[repr(C)] +#[derive(Default, PartialEq, Eq, Debug, Clone, Copy)] +pub struct FfiAccountIdWithPrivacy { + pub account_id: FfiBytes32, + pub is_private: bool, +} + +impl From for FfiAccountIdWithPrivacy { + fn from(value: AccountIdWithPrivacy) -> Self { + match value { + AccountIdWithPrivacy::Public(acc) => Self { + account_id: acc.into(), + is_private: false, + }, + AccountIdWithPrivacy::Private(acc) => Self { + account_id: acc.into(), + is_private: true, + }, + } + } +} + +impl From for AccountIdWithPrivacy { + fn from(value: FfiAccountIdWithPrivacy) -> Self { + if value.is_private { + Self::Private(value.account_id.into()) + } else { + Self::Public(value.account_id.into()) + } + } +} + #[cfg(test)] mod tests { use lee::{AccountId, PrivateKey, PublicKey}; diff --git a/lez/wallet-ffi/wallet_ffi.h b/lez/wallet-ffi/wallet_ffi.h index d83c520e..8f15a888 100644 --- a/lez/wallet-ffi/wallet_ffi.h +++ b/lez/wallet-ffi/wallet_ffi.h @@ -299,6 +299,27 @@ typedef struct FfiPublicAccountKey { struct FfiBytes32 public_key; } FfiPublicAccountKey; +typedef struct LabelAvailability { + bool is_available; + enum WalletFfiError error; +} LabelAvailability; + +typedef struct FfiAccountIdWithPrivacy { + struct FfiBytes32 account_id; + bool is_private; +} FfiAccountIdWithPrivacy; + +typedef struct AccountIdResolvedFromLabel { + struct FfiAccountIdWithPrivacy account_id; + enum WalletFfiError error; +} AccountIdResolvedFromLabel; + +typedef struct LabelList { + const char **labels_data; + uintptr_t labels_size; + enum WalletFfiError error; +} LabelList; + typedef struct FfiCreateWalletOutput { struct WalletHandle *wallet; /** @@ -807,6 +828,92 @@ enum WalletFfiError wallet_ffi_resolve_private_account(struct WalletHandle *hand */ void wallet_ffi_free_account_identity(struct FfiAccountIdentity *account_identity); +/** + * Check if label is available. + * + * # Parameters + * - `handle`: Valid wallet handle + * - `label`: Input null terminated C string for a label + * + * # Returns + * - `LabelAvailability` struct + * + * # Safety + * - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` + * - `label` must be a valid pointer to a null-terminated C string + */ +struct LabelAvailability wallet_ffi_check_label_available(struct WalletHandle *handle, + const char *label); + +/** + * Add new label. + * + * # Parameters + * - `handle`: Valid wallet handle + * - `label`: Input null terminated C string for a label + * - `account_id_with_privacy`: The account ID (32 bytes) and its privacy. + * + * # Returns + * - `Success` on successful query + * - Error code on failure + * + * # Safety + * - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` + * - `label` must be a valid pointer to a null-terminated C string + */ +enum WalletFfiError wallet_ffi_add_label(struct WalletHandle *handle, + const char *label, + struct FfiAccountIdWithPrivacy account_id_with_privacy); + +/** + * Resolve a label. + * + * # Parameters + * - `handle`: Valid wallet handle + * - `label`: Input null terminated C string for a label + * + * # Returns + * - `AccountIdResolvedFromLabel` struct + * + * # Safety + * - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` + * - `label` must be a valid pointer to a null-terminated C string + */ +struct AccountIdResolvedFromLabel wallet_ffi_resolve_label(struct WalletHandle *handle, + const char *label); + +/** + * Get all labels for account. + * + * # Parameters + * - `handle`: Valid wallet handle + * - `account_id_with_privacy`: The account ID (32 bytes) and its privacy. + * + * # Returns + * - `LabelList` struct + * + * # Safety + * - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open` + */ +struct LabelList wallet_ffi_get_all_labels_for_account(struct WalletHandle *handle, + struct FfiAccountIdWithPrivacy account_id_with_privacy); + +/** + * Free label list. + * + * # Parameters + * - `label_list`: Input list of labels + * + * # Returns + * - `Success` on successful query + * - Error code on failure + * + * # Safety + * - `label_list` must be a valid pointer to `LabelList`, received from + * `wallet_ffi_get_all_labels_for_account` + */ +enum WalletFfiError wallet_ffi_free_label_list(struct LabelList *label_list); + /** * Claim a pinata reward using a public transaction. * diff --git a/lez/wallet/src/account.rs b/lez/wallet/src/account.rs index 64eee575..8caa7366 100644 --- a/lez/wallet/src/account.rs +++ b/lez/wallet/src/account.rs @@ -21,6 +21,12 @@ impl Label { } } +impl AsRef for Label { + fn as_ref(&self) -> &str { + &self.0 + } +} + impl FromStr for Label { type Err = std::convert::Infallible; diff --git a/lez/wallet/src/account_manager.rs b/lez/wallet/src/account_manager.rs index ce9d1833..dae8fd31 100644 --- a/lez/wallet/src/account_manager.rs +++ b/lez/wallet/src/account_manager.rs @@ -5,8 +5,8 @@ use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder; use keycard_wallet::{KeycardWallet, python_path}; use lee::{AccountId, PrivateKey, PublicKey, Signature}; use lee_core::{ - Identifier, InputAccountIdentity, MembershipProof, NullifierPublicKey, NullifierSecretKey, - SharedSecretKey, + CommitmentSetDigest, DUMMY_COMMITMENT_HASH, Identifier, InputAccountIdentity, MembershipProof, + NullifierPublicKey, NullifierSecretKey, SharedSecretKey, account::{AccountWithMetadata, Nonce}, encryption::{EncryptedAccountData, EphemeralPublicKey, ViewingPublicKey}, }; @@ -187,6 +187,7 @@ enum State { pub struct AccountManager { states: Vec, pin: Option, + dummy_commitment_root: CommitmentSetDigest, } impl AccountManager { @@ -200,38 +201,15 @@ impl AccountManager { for account in accounts { let state = match account { AccountIdentity::Public(account_id) => { - let acc = wallet - .get_account_public(account_id) - .await - .map_err(ExecutionFailureKind::SequencerError)?; - - let sk = wallet.get_account_public_signing_key(account_id).cloned(); - let account = AccountWithMetadata::new(acc.clone(), sk.is_some(), account_id); - - State::Public { account, sk } + prepare_public_state(wallet, account_id, true).await? } AccountIdentity::PublicNoSign(account_id) => { - let acc = wallet - .get_account_public(account_id) - .await - .map_err(ExecutionFailureKind::SequencerError)?; - - let sk = None; - let account = AccountWithMetadata::new(acc.clone(), sk.is_some(), account_id); - - State::Public { account, sk } + prepare_public_state(wallet, account_id, false).await? } AccountIdentity::PublicKeycard { account_id, key_path, } => { - let acc = wallet - .get_account_public(account_id) - .await - .map_err(ExecutionFailureKind::SequencerError)?; - - let account = AccountWithMetadata::new(acc.clone(), true, account_id); - if pin.is_none() { pin = Some( crate::helperfunctions::read_pin() @@ -247,66 +225,27 @@ impl AccountManager { .to_owned(), ); } - - State::PublicKeycard { account, key_path } - } - AccountIdentity::PrivateOwned(account_id) => { - let pre = private_key_tree_acc_preparation(wallet, account_id, false).await?; - - State::Private(pre) + 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::PrivateForeign { npk, vpk, identifier, - } => { - let acc = lee_core::account::Account::default(); - let auth_acc = AccountWithMetadata::new(acc, false, (&npk, identifier)); - let eph_holder = EphemeralKeyHolder::new(&vpk); - let ssk = eph_holder.calculate_shared_secret_sender(); - let epk = eph_holder.ephemeral_public_key().clone(); - let pre = AccountPreparedData { - nsk: None, - npk, - identifier, - vpk, - pre_state: auth_acc, - proof: None, - ssk, - epk, - is_pda: false, - }; - - State::Private(pre) - } - AccountIdentity::PrivatePdaOwned(account_id) => { - let pre = private_key_tree_acc_preparation(wallet, account_id, true).await?; - State::Private(pre) - } + } => 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::PrivatePdaForeign { account_id, npk, vpk, identifier, - } => { - let acc = lee_core::account::Account::default(); - let auth_acc = AccountWithMetadata::new(acc, false, account_id); - let eph_holder = EphemeralKeyHolder::new(&vpk); - let ssk = eph_holder.calculate_shared_secret_sender(); - let epk = eph_holder.ephemeral_public_key().clone(); - let pre = AccountPreparedData { - nsk: None, - npk, - identifier, - vpk, - pre_state: auth_acc, - proof: None, - ssk, - epk, - is_pda: true, - }; - State::Private(pre) - } + } => State::Private(private_pda_foreign_acc_preparation( + account_id, npk, vpk, identifier, + )), AccountIdentity::PrivateShared { nsk, npk, @@ -314,12 +253,12 @@ impl AccountManager { identifier, } => { let account_id = lee::AccountId::from((&npk, identifier)); - let pre = private_shared_acc_preparation( - wallet, account_id, nsk, npk, vpk, identifier, false, + State::Private( + private_shared_acc_preparation( + wallet, account_id, nsk, npk, vpk, identifier, false, + ) + .await?, ) - .await?; - - State::Private(pre) } AccountIdentity::PrivatePdaShared { account_id, @@ -327,20 +266,35 @@ impl AccountManager { npk, vpk, identifier, - } => { - let pre = private_shared_acc_preparation( + } => State::Private( + private_shared_acc_preparation( wallet, account_id, nsk, npk, vpk, identifier, true, ) - .await?; - - State::Private(pre) - } + .await?, + ), }; states.push(state); } - Ok(Self { states, pin }) + 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 + }; + + Ok(Self { + states, + pin, + dummy_commitment_root, + }) } pub fn pre_states(&self) -> Vec { @@ -404,6 +358,7 @@ impl AccountManager { npk: pre.npk, ssk: pre.ssk, identifier: pre.identifier, + commitment_root: self.dummy_commitment_root, seed: None, }, }, @@ -424,6 +379,7 @@ impl AccountManager { ssk: pre.ssk, nsk, identifier: pre.identifier, + commitment_root: self.dummy_commitment_root, }, (None, _) => InputAccountIdentity::PrivateUnauthorized { epk: pre.epk.clone(), @@ -431,6 +387,7 @@ impl AccountManager { npk: pre.npk, ssk: pre.ssk, identifier: pre.identifier, + commitment_root: self.dummy_commitment_root, }, }, }) @@ -517,6 +474,37 @@ struct AccountPreparedData { is_pda: bool, } +async fn prepare_public_state( + wallet: &WalletCore, + account_id: AccountId, + lookup_signing_key: bool, +) -> Result { + let acc = wallet + .get_account_public(account_id) + .await + .map_err(ExecutionFailureKind::SequencerError)?; + let sk = if lookup_signing_key { + wallet.get_account_public_signing_key(account_id).cloned() + } else { + None + }; + let account = AccountWithMetadata::new(acc, sk.is_some(), account_id); + Ok(State::Public { account, sk }) +} + +async fn prepare_public_keycard_state( + wallet: &WalletCore, + account_id: AccountId, + key_path: String, +) -> Result { + let acc = wallet + .get_account_public(account_id) + .await + .map_err(ExecutionFailureKind::SequencerError)?; + let account = AccountWithMetadata::new(acc, true, account_id); + Ok(State::PublicKeycard { account, key_path }) +} + async fn private_key_tree_acc_preparation( wallet: &WalletCore, account_id: AccountId, @@ -599,6 +587,55 @@ async fn private_shared_acc_preparation( }) } +fn private_foreign_acc_preparation( + npk: NullifierPublicKey, + vpk: ViewingPublicKey, + identifier: Identifier, +) -> AccountPreparedData { + let acc = lee_core::account::Account::default(); + let pre_state = AccountWithMetadata::new(acc, false, (&npk, identifier)); + let eph_holder = EphemeralKeyHolder::new(&vpk); + let ssk = eph_holder.calculate_shared_secret_sender(); + let epk = eph_holder.ephemeral_public_key().clone(); + + AccountPreparedData { + nsk: None, + npk, + identifier, + vpk, + pre_state, + proof: None, + ssk, + epk, + is_pda: false, + } +} + +fn private_pda_foreign_acc_preparation( + account_id: AccountId, + npk: NullifierPublicKey, + vpk: ViewingPublicKey, + identifier: Identifier, +) -> AccountPreparedData { + let acc = lee_core::account::Account::default(); + let pre_state = AccountWithMetadata::new(acc, false, account_id); + let eph_holder = EphemeralKeyHolder::new(&vpk); + let ssk = eph_holder.calculate_shared_secret_sender(); + let epk = eph_holder.ephemeral_public_key().clone(); + + AccountPreparedData { + nsk: None, + npk, + identifier, + vpk, + pre_state, + proof: None, + ssk, + epk, + is_pda: true, + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/lez/wallet/src/cli/account.rs b/lez/wallet/src/cli/account.rs index fb0db2cb..b49cfb25 100644 --- a/lez/wallet/src/cli/account.rs +++ b/lez/wallet/src/cli/account.rs @@ -125,75 +125,169 @@ pub enum NewSubcommand { }, } +impl NewSubcommand { + fn handle_public( + cci: Option, + label: Option