From 3fde3c686e8af3d5b684c50fad4cede31d6aa95c Mon Sep 17 00:00:00 2001 From: Artem Gureev Date: Fri, 14 Aug 2026 10:29:13 +0000 Subject: [PATCH] feat: make the System Program executable with a Clear instruction --- lee/state_machine/core/src/program/mod.rs | 41 +++++- lee/state_machine/src/error.rs | 3 + lee/state_machine/src/state/tests/clear.rs | 128 ++++++++++++++++++ lee/state_machine/src/state/tests/mod.rs | 1 + .../src/validated_state_diff/mod.rs | 49 ++++++- 5 files changed, 218 insertions(+), 4 deletions(-) create mode 100644 lee/state_machine/src/state/tests/clear.rs diff --git a/lee/state_machine/core/src/program/mod.rs b/lee/state_machine/core/src/program/mod.rs index 112948b62..25437efcf 100644 --- a/lee/state_machine/core/src/program/mod.rs +++ b/lee/state_machine/core/src/program/mod.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use crate::{ BlockId, Identifier, NullifierPublicKey, Timestamp, - account::{Account, AccountId, AccountWithMetadata}, + account::{Account, AccountId, AccountWithMetadata, Data}, encryption::ViewingPublicKey, }; @@ -269,6 +269,11 @@ pub enum Claim { Pda(PdaSeed), } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum SystemInstruction { + Clear, +} + impl AccountPostState { /// Creates a post state without a claim request. /// The executing program is not requesting ownership of the account. @@ -628,6 +633,15 @@ pub enum ExecutionValidationError { }, } +#[derive(thiserror::Error, Debug)] +pub enum ClearValidationError { + #[error("Unauthorized clear of account {account_id}")] + NotAuthorized { account_id: AccountId }, + + #[error("Invalid clear transition for account {account_id}")] + InvalidTransition { account_id: AccountId }, +} + /// Computes the set of public-PDA `AccountId`s the callee is authorized to mutate. /// /// Returns only public-form derivations, suitable for contexts where all accounts are public @@ -765,6 +779,31 @@ pub fn validate_execution( Ok(()) } +pub fn validate_clear( + pre: &AccountWithMetadata, + post: &Account, +) -> Result<(), ClearValidationError> { + if !pre.is_authorized { + return Err(ClearValidationError::NotAuthorized { + account_id: pre.account_id, + }); + } + + let expected = Account { + program_owner: DEFAULT_PROGRAM_ID, + balance: pre.account.balance, + data: Data::default(), + nonce: pre.account.nonce, + }; + if *post != expected { + return Err(ClearValidationError::InvalidTransition { + account_id: pre.account_id, + }); + } + + Ok(()) +} + fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> bool { let number_of_accounts = pre_states.len(); let number_of_account_ids = pre_states diff --git a/lee/state_machine/src/error.rs b/lee/state_machine/src/error.rs index 94e66e8b8..2b34ec61b 100644 --- a/lee/state_machine/src/error.rs +++ b/lee/state_machine/src/error.rs @@ -114,6 +114,9 @@ pub enum InvalidProgramBehaviorError { #[error(transparent)] ExecutionValidationFailed(#[from] lee_core::program::ExecutionValidationError), + #[error(transparent)] + ClearValidationFailed(#[from] lee_core::program::ClearValidationError), + #[error("Trying to claim account {account_id} which is not default")] ClaimedNonDefaultAccount { account_id: AccountId }, diff --git a/lee/state_machine/src/state/tests/clear.rs b/lee/state_machine/src/state/tests/clear.rs new file mode 100644 index 000000000..1d845a401 --- /dev/null +++ b/lee/state_machine/src/state/tests/clear.rs @@ -0,0 +1,128 @@ +use lee_core::program::{ClearValidationError, DEFAULT_PROGRAM_ID, SystemInstruction}; + +use super::*; + +const HOSTILE_OWNER: ProgramId = [1, 2, 3, 4, 5, 6, 7, 8]; + +#[test] +fn clear_reclaims_hostile_owned_account() { + let key = PrivateKey::try_new([1; 32]).unwrap(); + let id = AccountId::from(&PublicKey::new_from_private_key(&key)); + let mut state = V03State::new(); + state.force_insert_account( + id, + Account { + program_owner: HOSTILE_OWNER, + balance: 500, + data: vec![0xca, 0xfe].try_into().unwrap(), + nonce: Nonce(0), + }, + ); + + let message = public_transaction::Message::try_new( + DEFAULT_PROGRAM_ID, + vec![id], + vec![Nonce(0)], + SystemInstruction::Clear, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&key]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + assert_eq!( + state.get_account_by_id(id), + Account { + program_owner: DEFAULT_PROGRAM_ID, + balance: 500, + data: Data::default(), + nonce: Nonce(1), + } + ); +} + +#[test] +fn clear_by_non_signer_is_rejected() { + let victim_key = PrivateKey::try_new([1; 32]).unwrap(); + let victim = AccountId::from(&PublicKey::new_from_private_key(&victim_key)); + let attacker_key = PrivateKey::try_new([2; 32]).unwrap(); + + let victim_account = Account { + program_owner: HOSTILE_OWNER, + balance: 500, + data: vec![0xca, 0xfe].try_into().unwrap(), + nonce: Nonce(0), + }; + let mut state = V03State::new(); + state.force_insert_account(victim, victim_account.clone()); + + let message = public_transaction::Message::try_new( + DEFAULT_PROGRAM_ID, + vec![victim], + vec![Nonce(0)], + SystemInstruction::Clear, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&attacker_key]); + let tx = PublicTransaction::new(message, witness_set); + + let result = state.transition_from_public_transaction(&tx, 1, 0); + + assert!(matches!( + result, + Err(LeeError::InvalidProgramBehavior( + InvalidProgramBehaviorError::ClearValidationFailed( + ClearValidationError::NotAuthorized { account_id } + ) + )) if account_id == victim + )); + assert_eq!(state.get_account_by_id(victim), victim_account); +} + +#[test] +fn clear_touches_only_the_named_account() { + let key = PrivateKey::try_new([1; 32]).unwrap(); + let target = AccountId::from(&PublicKey::new_from_private_key(&key)); + let bystander = AccountId::new([9; 32]); + let bystander_account = Account { + program_owner: HOSTILE_OWNER, + balance: 777, + data: vec![0xaa].try_into().unwrap(), + nonce: Nonce(3), + }; + let mut state = V03State::new(); + state.force_insert_account( + target, + Account { + program_owner: HOSTILE_OWNER, + balance: 500, + data: vec![0xca, 0xfe].try_into().unwrap(), + nonce: Nonce(0), + }, + ); + state.force_insert_account(bystander, bystander_account.clone()); + + let message = public_transaction::Message::try_new( + DEFAULT_PROGRAM_ID, + vec![target], + vec![Nonce(0)], + SystemInstruction::Clear, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&key]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx, 1, 0).unwrap(); + + assert_eq!( + state.get_account_by_id(target), + Account { + program_owner: DEFAULT_PROGRAM_ID, + balance: 500, + data: Data::default(), + nonce: Nonce(1), + } + ); + assert_eq!(state.get_account_by_id(bystander), bystander_account); +} diff --git a/lee/state_machine/src/state/tests/mod.rs b/lee/state_machine/src/state/tests/mod.rs index 977496780..4dc0ad4da 100644 --- a/lee/state_machine/src/state/tests/mod.rs +++ b/lee/state_machine/src/state/tests/mod.rs @@ -35,6 +35,7 @@ mod authenticated_transfer; mod changer_claimer; mod circuit; mod claiming; +mod clear; mod flash_swap; mod genesis; mod privacy_preserving; diff --git a/lee/state_machine/src/validated_state_diff/mod.rs b/lee/state_machine/src/validated_state_diff/mod.rs index 3598b16ba..8c237d7fb 100644 --- a/lee/state_machine/src/validated_state_diff/mod.rs +++ b/lee/state_machine/src/validated_state_diff/mod.rs @@ -5,10 +5,10 @@ use std::{ use lee_core::{ BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, PublicAction, Timestamp, - account::{Account, AccountId, AccountWithMetadata}, + account::{Account, AccountId, AccountWithMetadata, Data}, program::{ - CallerData, ChainedCall, Claim, DEFAULT_PROGRAM_ID, compute_public_authorized_pdas, - validate_execution, + CallerData, ChainedCall, Claim, DEFAULT_PROGRAM_ID, SystemInstruction, + compute_public_authorized_pdas, validate_clear, validate_execution, }, }; use log::debug; @@ -111,6 +111,49 @@ impl ValidatedStateDiff { LeeError::MaxChainedCallsDepthExceeded ); + if chained_call.program_id == DEFAULT_PROGRAM_ID { + let instruction: SystemInstruction = + risc0_zkvm::serde::from_slice(&chained_call.instruction_data) + .map_err(|e| LeeError::InstructionSerializationError(e.to_string()))?; + + let authorized_pdas = + compute_public_authorized_pdas(caller_data.program_id, &chained_call.pda_seeds); + let is_authorized = |account_id: &AccountId| { + authorized_pdas.contains(account_id) + || caller_data.authorized_accounts.contains(account_id) + }; + + match instruction { + SystemInstruction::Clear => { + for pre_state in &chained_call.pre_states { + let account_id = pre_state.account_id; + let pre = AccountWithMetadata::new( + state_diff + .get(&account_id) + .cloned() + .unwrap_or_else(|| state.get_account_by_id(account_id)), + is_authorized(&account_id), + account_id, + ); + let post = Account { + program_owner: DEFAULT_PROGRAM_ID, + balance: pre.account.balance, + data: Data::default(), + nonce: pre.account.nonce, + }; + validate_clear(&pre, &post) + .map_err(InvalidProgramBehaviorError::ClearValidationFailed)?; + state_diff.insert(account_id, post); + } + } + } + + chain_calls_counter = chain_calls_counter + .checked_add(1) + .expect("we check the max depth at the beginning of the loop"); + continue; + } + // 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()));