537 lines
22 KiB
Rust

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<AccountId, Account> {
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::from_circuit_output(
vec![], // no public signers, no nonces
circuit_output,
);
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::from_circuit_output(
vec![], // no public signers, no nonces
circuit_output,
);
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, EncryptedAccountData, Nullifier, PrivateAction,
account::Account,
encryption::{Ciphertext, EphemeralPublicKey},
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_actions: vec![],
nonces: vec![],
private_actions: vec![PrivateAction {
nullifier: Nullifier::for_account_initialization(&account_id),
root: [0; 32],
commitment,
encrypted_post_state: EncryptedAccountData {
ciphertext: Ciphertext::from_inner(vec![]),
epk: EphemeralPublicKey(vec![]),
view_tag: 0,
},
}],
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"),
}
}