feat: make the System Program executable with a Clear instruction

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

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 },
 

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);
+}

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;

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()));