Merge pull request #479 from logos-blockchain/schouhy/fix-faucet-account-protection-mechanism

fix: Bug in faucet account protection mechanism

Cargo.lock

@@ -9211,6 +9211,7 @@ version = "0.1.0"
 dependencies = [
  "authenticated_transfer_core",
  "clock_core",
+ "faucet_core",
  "nssa_core",
  "risc0-zkvm",
  "serde",

common/src/transaction.rs

@@ -67,26 +67,17 @@ impl NSSATransaction {
     }
 
     /// Validates the transaction against the current state and returns the resulting diff
-    /// without applying it. Rejects transactions that modify clock system accounts and
-    /// rejects unsafe modifications of the system faucet account. Also rejects direct
-    /// invocation of the faucet program for user-submitted transactions.
+    /// without applying it. Rejects transactions that modify clock or faucet system accounts,
+    /// whether directly or indirectly via chain calls.
     ///
-    /// This check is required for all user transactions. Only sequencer transaction may bypass this
-    /// check.
+    /// This check is required for all user transactions. Only sequencer transactions may bypass
+    /// this check.
     pub fn validate_on_state(
         &self,
         state: &V03State,
         block_id: BlockId,
         timestamp: Timestamp,
     ) -> Result<ValidatedStateDiff, nssa::error::NssaError> {
-        if let Self::Public(tx) = self
-            && tx.message().program_id == nssa::program::Program::faucet().id()
-        {
-            return Err(nssa::error::NssaError::InvalidInput(
-                "Transaction invokes restricted faucet program".into(),
-            ));
-        }
-
         let diff = match self {
             Self::Public(tx) => {
                 ValidatedStateDiff::from_public_transaction(tx, state, block_id, timestamp)
@@ -111,6 +102,16 @@ impl NSSATransaction {
             ));
         }
 
+        let faucet_id = nssa::system_faucet_account_id();
+        if public_diff
+            .get(&faucet_id)
+            .is_some_and(|post| *post != state.get_account_by_id(faucet_id))
+        {
+            return Err(nssa::error::NssaError::InvalidInput(
+                "Transaction modifies system faucet account".into(),
+            ));
+        }
+
         Ok(diff)
     }
 

integration_tests/tests/auth_transfer/private.rs

@@ -1,6 +1,7 @@
 use std::time::Duration;
 
 use anyhow::{Context as _, Result};
+use common::transaction::NSSATransaction;
 use integration_tests::{
     TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx, private_mention,
     public_mention, verify_commitment_is_in_state,
@@ -623,3 +624,130 @@ async fn shielded_transfers_to_two_identifiers_same_npk() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+async fn ppt_that_chain_calls_faucet_is_dropped() -> Result<()> {
+    use nssa::{
+        EphemeralPublicKey, SharedSecretKey, execute_and_prove,
+        privacy_preserving_transaction::{self, circuit::ProgramWithDependencies},
+    };
+    use nssa_core::{InputAccountIdentity, account::AccountWithMetadata};
+
+    let ctx = TestContext::new().await?;
+
+    let binary = std::fs::read(
+        std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+            .join("../artifacts/test_program_methods/faucet_chain_caller.bin"),
+    )?;
+    let deploy_tx = NSSATransaction::ProgramDeployment(nssa::ProgramDeploymentTransaction::new(
+        nssa::program_deployment_transaction::Message::new(binary.clone()),
+    ));
+    ctx.sequencer_client().send_transaction(deploy_tx).await?;
+
+    info!("Waiting for deploy block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let faucet_account_id = nssa::system_faucet_account_id();
+    let attacker_id = ctx.existing_public_accounts()[0];
+    let faucet_program_id = Program::faucet().id();
+    let vault_program_id = Program::vault().id();
+    let auth_transfer_program_id = Program::authenticated_transfer_program().id();
+    let nsk: nssa_core::NullifierSecretKey = [3; 32];
+    let npk = NullifierPublicKey::from(&nsk);
+    let vpk = Secp256k1Point::from_scalar([4; 32]);
+    let ssk = SharedSecretKey::new([55; 32], &vpk);
+    let epk = EphemeralPublicKey::from_scalar([55; 32]);
+    let attacker_vault_id = {
+        let seed = vault_core::compute_vault_seed(attacker_id);
+        AccountId::for_private_pda(&vault_program_id, &seed, &npk, 1337)
+    };
+    let amount: u128 = 1;
+
+    let faucet_pre = AccountWithMetadata::new(
+        ctx.sequencer_client()
+            .get_account(faucet_account_id)
+            .await?,
+        false,
+        faucet_account_id,
+    );
+    let vault_pda_pre = AccountWithMetadata::new(
+        ctx.sequencer_client()
+            .get_account(attacker_vault_id)
+            .await?,
+        false,
+        attacker_vault_id,
+    );
+
+    let faucet_chain_caller = Program::new(binary)?;
+    let program_with_deps = ProgramWithDependencies::new(
+        faucet_chain_caller,
+        [
+            (faucet_program_id, Program::faucet()),
+            (vault_program_id, Program::vault()),
+            (
+                auth_transfer_program_id,
+                Program::authenticated_transfer_program(),
+            ),
+        ]
+        .into(),
+    );
+
+    let instruction =
+        Program::serialize_instruction((faucet_program_id, vault_program_id, attacker_id, amount))?;
+
+    let (output, proof) = execute_and_prove(
+        vec![faucet_pre, vault_pda_pre],
+        instruction,
+        vec![
+            InputAccountIdentity::Public,
+            InputAccountIdentity::PrivatePdaInit {
+                npk,
+                ssk,
+                identifier: 1337,
+            },
+        ],
+        &program_with_deps,
+    )?;
+
+    let message = privacy_preserving_transaction::Message::try_from_circuit_output(
+        vec![faucet_account_id],
+        vec![],
+        vec![(npk, vpk, epk)],
+        output,
+    )?;
+    let witness_set = privacy_preserving_transaction::WitnessSet::for_message(&message, proof, &[]);
+    let attack_ppt = NSSATransaction::PrivacyPreserving(nssa::PrivacyPreservingTransaction::new(
+        message,
+        witness_set,
+    ));
+
+    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 tx_hash = ctx.sequencer_client().send_transaction(attack_ppt).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 tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
+
+    assert_eq!(faucet_balance_after, faucet_balance_before);
+    assert_eq!(vault_balance_after, vault_balance_before);
+    assert!(tx_on_chain.is_none());
+
+    Ok(())
+}

integration_tests/tests/auth_transfer/public.rs

@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{path::PathBuf, time::Duration};
 
 use anyhow::Result;
 use common::transaction::NSSATransaction;
@@ -397,44 +397,6 @@ async fn cannot_transfer_funds_from_system_faucet_account() -> Result<()> {
     Ok(())
 }
 
-#[test]
-async fn can_transfer_funds_to_system_faucet_account() -> Result<()> {
-    let mut ctx = TestContext::new().await?;
-    let faucet_account_id = system_faucet_account_id();
-
-    let sender = ctx.existing_public_accounts()[0];
-    let sender_balance_before = ctx.sequencer_client().get_account_balance(sender).await?;
-    let faucet_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
-
-    let amount = 100_u128;
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(sender),
-        to: Some(public_mention(faucet_account_id)),
-        to_npk: None,
-        to_vpk: None,
-        to_identifier: Some(0),
-        amount,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
-    let sender_balance_after = ctx.sequencer_client().get_account_balance(sender).await?;
-    let faucet_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
-
-    assert_eq!(sender_balance_after, sender_balance_before - amount);
-    assert_eq!(faucet_balance_after, faucet_balance_before + amount);
-
-    Ok(())
-}
-
 #[test]
 async fn cannot_execute_faucet_program() -> Result<()> {
     let ctx = TestContext::new().await?;
@@ -492,3 +454,69 @@ async fn cannot_execute_faucet_program() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+async fn user_tx_that_chain_calls_faucet_is_dropped() -> Result<()> {
+    let ctx = TestContext::new().await?;
+
+    let binary = std::fs::read(
+        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+            .join("../artifacts/test_program_methods/faucet_chain_caller.bin"),
+    )?;
+    let faucet_chain_caller_id = Program::new(binary.clone())?.id();
+    let deploy_tx = NSSATransaction::ProgramDeployment(nssa::ProgramDeploymentTransaction::new(
+        nssa::program_deployment_transaction::Message::new(binary),
+    ));
+    ctx.sequencer_client().send_transaction(deploy_tx).await?;
+
+    info!("Waiting for deploy block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let faucet_account_id = system_faucet_account_id();
+    let attacker = ctx.existing_public_accounts()[0];
+    let faucet_program_id = Program::faucet().id();
+    let vault_program_id = Program::vault().id();
+    let attacker_vault_id = vault_core::compute_vault_account_id(vault_program_id, attacker);
+    let amount: u128 = 1;
+
+    let message = public_transaction::Message::try_new(
+        faucet_chain_caller_id,
+        vec![faucet_account_id, attacker_vault_id],
+        vec![],
+        (faucet_program_id, vault_program_id, attacker, amount),
+    )?;
+    let attack_tx = NSSATransaction::Public(nssa::PublicTransaction::new(
+        message,
+        nssa::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 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 tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
+
+    assert_eq!(faucet_balance_after, faucet_balance_before);
+    assert_eq!(vault_balance_after, vault_balance_before);
+    assert!(tx_on_chain.is_none());
+
+    Ok(())
+}

program_methods/guest/src/bin/privacy_preserving_circuit/execution_state.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::{HashMap, VecDeque, hash_map::Entry},
+    collections::{HashMap, HashSet, VecDeque, hash_map::Entry},
     convert::Infallible,
 };
 
@@ -49,6 +49,7 @@ pub struct ExecutionState {
     /// caller-seeds authorization paths to verify
     /// `AccountId::for_private_pda(program_id, seed, npk, identifier) == pre_state.account_id`.
     private_pda_npk_by_position: HashMap<usize, (NullifierPublicKey, Identifier)>,
+    authorized_accounts: HashSet<AccountId>,
 }
 
 impl ExecutionState {
@@ -107,6 +108,7 @@ impl ExecutionState {
             private_pda_bound_positions: HashMap::new(),
             pda_family_binding: HashMap::new(),
             private_pda_npk_by_position,
+            authorized_accounts: HashSet::new(),
         };
 
         let Some(first_output) = program_outputs.first() else {
@@ -246,10 +248,10 @@ impl ExecutionState {
         program_id: ProgramId,
         caller_program_id: Option<ProgramId>,
         caller_pda_seeds: &[PdaSeed],
-        pre_states: Vec<AccountWithMetadata>,
-        post_states: Vec<AccountPostState>,
+        output_pre_states: Vec<AccountWithMetadata>,
+        output_post_states: Vec<AccountPostState>,
     ) {
-        for (pre, mut post) in pre_states.into_iter().zip(post_states) {
+        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);
@@ -288,6 +290,7 @@ impl ExecutionState {
                         &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,
@@ -491,6 +494,7 @@ fn resolve_authorization_and_record_bindings(
     pda_family_binding: &mut HashMap<(ProgramId, PdaSeed), AccountId>,
     private_pda_bound_positions: &mut HashMap<usize, (ProgramId, PdaSeed)>,
     private_pda_npk_by_position: &HashMap<usize, (NullifierPublicKey, Identifier)>,
+    authorized_accounts: &mut HashSet<AccountId>,
     pre_account_id: AccountId,
     pre_state_position: usize,
     caller_program_id: Option<ProgramId>,
@@ -525,5 +529,13 @@ fn resolve_authorization_and_record_bindings(
         }
     }
 
-    previous_is_authorized || matched_caller_seed.is_some()
+    if authorized_accounts.contains(&pre_account_id) {
+        return true;
+    }
+
+    let authorized = previous_is_authorized || matched_caller_seed.is_some();
+    if authorized {
+        authorized_accounts.insert(pre_account_id);
+    }
+    authorized
 }

program_methods/guest/src/bin/vault.rs

@@ -40,7 +40,7 @@ fn main() {
         } => {
             let [sender, recipient_vault] = pre_states
                 .try_into()
-                .expect("Transfer requires exactly 3 accounts");
+                .expect("Transfer requires exactly 2 accounts");
 
             let seed = vault_core::compute_vault_seed(recipient_id);
 

test_program_methods/guest/Cargo.toml

@@ -11,6 +11,7 @@ workspace = true
 nssa_core.workspace = true
 authenticated_transfer_core.workspace = true
 clock_core.workspace = true
+faucet_core.workspace = true
 
 risc0-zkvm.workspace = true
 serde = { workspace = true, default-features = false }

test_program_methods/guest/src/bin/faucet_chain_caller.rs

@@ -0,0 +1,52 @@
+use nssa_core::{
+    account::AccountId,
+    program::{
+        AccountPostState, ChainedCall, ProgramId, ProgramInput, ProgramOutput, read_nssa_inputs,
+    },
+};
+use risc0_zkvm::serde::to_vec;
+
+type Instruction = (ProgramId, ProgramId, AccountId, u128);
+// (faucet_program_id, vault_program_id, recipient_id, amount)
+
+fn main() {
+    let (
+        ProgramInput {
+            self_program_id,
+            caller_program_id,
+            pre_states,
+            instruction: (faucet_program_id, vault_program_id, recipient_id, amount),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let post_states: Vec<_> = pre_states
+        .iter()
+        .map(|pre| AccountPostState::new(pre.account.clone()))
+        .collect();
+
+    assert_eq!(pre_states.len(), 2);
+    let [faucet_pre, vault_pda_pre] = [pre_states[0].clone(), pre_states[1].clone()];
+
+    let chained_calls = vec![ChainedCall {
+        program_id: faucet_program_id,
+        instruction_data: to_vec(&faucet_core::Instruction::Transfer {
+            vault_program_id,
+            recipient_id,
+            amount,
+        })
+        .unwrap(),
+        pre_states: vec![faucet_pre, vault_pda_pre],
+        pda_seeds: vec![],
+    }];
+
+    ProgramOutput::new(
+        self_program_id,
+        caller_program_id,
+        instruction_words,
+        pre_states,
+        post_states,
+    )
+    .with_chained_calls(chained_calls)
+    .write();
+}