feat: assert non revertible phase when setting fee payer (#24479)

The fee payer should only be set in the non-revertible phase, but we
were not asserting this.

Author: nventuro

docs/docs-developers/docs/resources/migration_notes.md

@@ -9,6 +9,12 @@ Aztec is in active development. Each version may introduce breaking changes that
 
 ## TBD
 
+### [Aztec.nr] `set_as_fee_payer` now asserts it is called during the setup phase
+
+`PrivateContext::set_as_fee_payer` now asserts that execution is still in the setup (non-revertible) phase, i.e. that `end_setup` has not yet been called by any function in the transaction. Electing a fee payer in the revertible phase was never safe: compensation collected by the fee payer after `end_setup` can be discarded if a public call later reverts, while the protocol still debits the fee payer's fee-juice balance.
+
+**Impact**: A transaction in which `set_as_fee_payer` runs after the setup phase has ended now fails with `fee payer must be elected during the setup phase`. Standard fee payment flows, which call `set_as_fee_payer` before or together with `end_setup`, are unaffected.
+
 ### [PXE] Stores are now selected by `(l1ChainId, rollupAddress, schemaVersion)` instead of being wiped on mismatch
 
 Previously, connecting a PXE or embedded wallet to a different or redeployed rollup, or bumping the store schema version, wiped the existing on-disk store in place. That meant master account keys could be destroyed simply by pointing a wallet at a different network. PXE data stores now exist per `(l1ChainId, rollupAddress, schemaVersion)` triple, and switching networks (or upgrading) selects or creates the matching store instead of overwriting previous ones. The embedded wallet's `wallet_data` store is partitioned the same way, so accounts and aliases are per network: switching networks starts with an empty account list until accounts are re-imported, and switching back finds the originals intact.
@@ -140,8 +146,6 @@ Registering classes and instances are now separate, unvalidated operations. `reg
   The new class is used automatically once the upgrade takes effect on chain; no further PXE action is needed. Registering it beforehand is harmless: until the update activates, the node still resolves the contract's current class to the previous one, so it keeps running its old code.
 
 - `pxe.getContractInstance(address)` and `wallet.getContractMetadata(address).instance` now return the contract's **address preimage**, which no longer includes `currentContractClassId`.
-
-
 ### [Aztec.js] `AccountWithSecretKey` removed, read account keys from the `AccountManager` or PXE
 
 `AccountWithSecretKey` was a thin wrapper that bundled an account's transaction signer with its master secret key, used mainly to print or export the secret. It has been removed, and `AccountManager.getAccount()` now returns the plain `Account` signer. The wrapper's extra methods are no longer available on that value:

noir-projects/aztec-nr/aztec/src/context/private_context.nr

@@ -537,11 +537,22 @@ impl PrivateContext {
     /// Only one contract per transaction can declare itself as the fee payer, and it must have sufficient fee-juice
     /// balance (>= the gas limits specified in the TxContext) by the time we reach the public setup phase of the tx.
     ///
+    /// The fee payer must be elected during the setup (non-revertible) phase, i.e. before
+    /// [`end_setup`](PrivateContext::end_setup) is called - this function asserts so. This is because any compensation
+    /// collected by the fee payer during the revertible phase can be discarded if a public call later reverts, while
+    /// the protocol still debits the fee payer's fee-juice balance. Note that `end_setup` does not need to be called
+    /// by the electing function itself: it can be called later in the transaction (e.g. by the fee-juice contract when
+    /// claiming fee juice that pays for the very same transaction).
     pub fn set_as_fee_payer(&mut self) {
+        assert(!self.in_revertible_phase(), "fee payer must be elected during the setup phase");
         aztecnr_trace_log_format!("Setting {0} as fee payer")([self.this_address().to_field()]);
         self.is_fee_payer = true;
     }
 
+    /// Returns whether execution is currently in the revertible (app) phase of the transaction.
+    ///
+    /// A transaction is in the revertible phase if [`end_setup`](PrivateContext::end_setup) has already been called -
+    /// potentially by a different function of the same transaction.
     pub fn in_revertible_phase(&mut self) -> bool {
         let current_counter = self.side_effect_counter;
 

yarn-project/end-to-end/src/automine/phase_check.parallel.test.ts

@@ -1,4 +1,5 @@
 import { AztecAddress } from '@aztec/aztec.js/addresses';
+import { BatchCall } from '@aztec/aztec.js/contracts';
 import { SponsoredFeePaymentMethod } from '@aztec/aztec.js/fee';
 import { SPONSORED_FPC_SALT } from '@aztec/constants';
 import { Fr } from '@aztec/foundation/curves/bn254';
@@ -7,11 +8,23 @@ import { TestContract } from '@aztec/noir-test-contracts.js/Test';
 import { computeFeePayerBalanceLeafSlot } from '@aztec/protocol-contracts/fee-juice';
 import { getContractInstanceFromInstantiationParams } from '@aztec/stdlib/contract';
 import { PublicDataTreeLeaf } from '@aztec/stdlib/trees';
+import { ExecutionPayload } from '@aztec/stdlib/tx';
 import { defaultInitialAccountFeeJuice } from '@aztec/world-state/testing';
 
 import type { TestWallet } from '../test-wallet/test_wallet.js';
 import { AutomineTestContext } from './automine_test_context.js';
 
+/**
+ * Declares a contract as the tx fee payer without contributing any call to the fee payload, unlike
+ * SponsoredFeePaymentMethod which prepends the sponsor call (and thus makes the election run before any app call).
+ * This lets a test place the electing call anywhere in the app payload, e.g. after the setup phase has ended.
+ */
+class DeferredSponsoredFeePaymentMethod extends SponsoredFeePaymentMethod {
+  override async getExecutionPayload(): Promise<ExecutionPayload> {
+    return new ExecutionPayload([], [], [], [], await this.getFeePayer());
+  }
+}
+
 // Private functions should receive automatically a phase check that avoids any nested call changing the phase.
 // Functions that opt out of this phase check can be marked with #[allow_phase_change].
 //
@@ -78,4 +91,21 @@ describe('automine/phase_check', () => {
       },
     });
   });
+
+  it('should fail when the fee payer is elected after the setup phase has ended', async () => {
+    // BatchCall.simulate ignores the fee payment method, which would make the wallet fall back to
+    // PREEXISTING_FEE_JUICE and end setup in the account entrypoint before any app call runs. Build the payload via
+    // request() instead, which merges the payment method's payload (and thus its fee payer), and simulate it directly.
+    const lateElection = await new BatchCall(wallet, [
+      contract.methods.call_function_that_ends_setup_without_phase_check(),
+      sponsoredFPC.methods.sponsor_unconditionally(),
+    ]).request({
+      fee: {
+        paymentMethod: new DeferredSponsoredFeePaymentMethod(sponsoredFPC.address),
+      },
+    });
+    await expect(wallet.simulateTx(lateElection, { from: defaultAccountAddress })).rejects.toThrow(
+      'fee payer must be elected during the setup phase',
+    );
+  });
 });