Use SafeERC20 for reward flows and add non-compliant token tests

src/core/libraries/rollup/RewardLib.sol

@@ -87,7 +87,7 @@ library RewardLib {
     RollupStore storage rollupStore = STFLib.getStorage();
     uint256 amount = rewardStorage.sequencerRewards[_sequencer];
     rewardStorage.sequencerRewards[_sequencer] = 0;
-    rollupStore.config.feeAsset.transfer(_sequencer, amount);
+    rollupStore.config.feeAsset.safeTransfer(_sequencer, amount);
 
     return amount;
   }
@@ -119,7 +119,7 @@ library RewardLib {
       }
     }
 
-    rollupStore.config.feeAsset.transfer(_prover, accumulatedRewards);
+    rollupStore.config.feeAsset.safeTransfer(_prover, accumulatedRewards);
 
     return accumulatedRewards;
   }
@@ -215,7 +215,7 @@ library RewardLib {
       }
 
       if (t.totalBurn > 0) {
-        rollupStore.config.feeAsset.transfer(BURN_ADDRESS, t.totalBurn);
+        rollupStore.config.feeAsset.safeTransfer(BURN_ADDRESS, t.totalBurn);
       }
     }
   }

src/core/libraries/rollup/StakingLib.sol

@@ -181,7 +181,7 @@ library StakingLib {
     delete store.exits[_attester];
 
     store.gse.finalizeWithdraw(exit.withdrawalId);
-    store.stakingAsset.transfer(exit.recipientOrWithdrawer, exit.amount);
+    store.stakingAsset.safeTransfer(exit.recipientOrWithdrawer, exit.amount);
 
     emit IStakingCore.WithdrawFinalized(_attester, exit.recipientOrWithdrawer, exit.amount);
   }
@@ -282,7 +282,7 @@ library StakingLib {
     require(!store.exits[_attester].exists, Errors.Staking__AlreadyExiting(_attester));
     uint256 amount = store.gse.ACTIVATION_THRESHOLD();
 
-    store.stakingAsset.transferFrom(msg.sender, address(this), amount);
+    store.stakingAsset.safeTransferFrom(msg.sender, address(this), amount);
     store.entryQueue.enqueue(
       _attester, _withdrawer, _publicKeyInG1, _publicKeyInG2, _proofOfPossession, _moveWithLatestRollup
     );
@@ -327,8 +327,11 @@ library StakingLib {
 
     uint256 queueLength = store.entryQueue.length();
     uint256 numToDequeue = Math.min(maxAddableValidators, queueLength);
-    // Approve the GSE to spend the total stake amount needed for all deposits.
-    store.stakingAsset.approve(address(store.gse), amount * numToDequeue);
+    uint256 totalAmount = amount * numToDequeue;
+    if (totalAmount > 0) {
+      store.stakingAsset.safeApprove(address(store.gse), 0);
+      store.stakingAsset.safeApprove(address(store.gse), totalAmount);
+    }
     for (uint256 i = 0; i < numToDequeue; i++) {
       DepositArgs memory args = store.entryQueue.dequeue();
       (bool success, bytes memory data) = address(store.gse).call(
@@ -354,13 +357,15 @@ library StakingLib {
         //    where someone could drain the queue without making any deposits.
         //    We can safely assume data.length == 0 means out of gas since we only call trusted GSE contract.
         require(data.length > 0, Errors.Staking__DepositOutOfGas());
-        store.stakingAsset.transfer(args.withdrawer, amount);
+        store.stakingAsset.safeTransfer(args.withdrawer, amount);
         emit IStakingCore.FailedDeposit(
           args.attester, args.withdrawer, args.publicKeyInG1, args.publicKeyInG2, args.proofOfPossession
         );
       }
     }
-    store.stakingAsset.approve(address(store.gse), 0);
+    if (totalAmount > 0) {
+      store.stakingAsset.safeApprove(address(store.gse), 0);
+    }
   }
 
   /**

src/governance/GSE.sol

@@ -14,6 +14,7 @@ import {BN254Lib, G1Point, G2Point} from "@aztec/shared/libraries/BN254Lib.sol";
 import {Timestamp} from "@aztec/shared/libraries/TimeMath.sol";
 import {Ownable} from "@oz/access/Ownable.sol";
 import {IERC20} from "@oz/token/ERC20/IERC20.sol";
+import {SafeERC20} from "@oz/token/ERC20/utils/SafeERC20.sol";
 import {SafeCast} from "@oz/utils/math/SafeCast.sol";
 import {Checkpoints} from "@oz/utils/structs/Checkpoints.sol";
 
@@ -116,6 +117,7 @@ interface IGSE is IGSECore {
  * then have the deployer `setGovernance`, and then `transferOwnership` to Governance.
  */
 contract GSECore is IGSECore, Ownable {
+  using SafeERC20 for IERC20;
   using AddressSnapshotLib for SnapshottedAddressSet;
   using SafeCast for uint256;
   using SafeCast for uint224;
@@ -336,10 +338,11 @@ contract GSECore is IGSECore, Ownable {
     delegation.delegate(recipientInstance, _attester, recipientInstance);
     delegation.increaseBalance(recipientInstance, _attester, ACTIVATION_THRESHOLD);
 
-    ASSET.transferFrom(msg.sender, address(this), ACTIVATION_THRESHOLD);
+    ASSET.safeTransferFrom(msg.sender, address(this), ACTIVATION_THRESHOLD);
 
     Governance gov = getGovernance();
-    ASSET.approve(address(gov), ACTIVATION_THRESHOLD);
+    ASSET.safeApprove(address(gov), 0);
+    ASSET.safeApprove(address(gov), ACTIVATION_THRESHOLD);
     gov.deposit(address(this), ACTIVATION_THRESHOLD);
 
     emit Deposit(recipientInstance, _attester, _withdrawer);
@@ -469,8 +472,9 @@ contract GSECore is IGSECore, Ownable {
     Governance gov = getGovernance();
     uint256 amount = gov.getConfiguration().proposeConfig.lockAmount;
 
-    ASSET.transferFrom(msg.sender, address(this), amount);
-    ASSET.approve(address(gov), amount);
+    ASSET.safeTransferFrom(msg.sender, address(this), amount);
+    ASSET.safeApprove(address(gov), 0);
+    ASSET.safeApprove(address(gov), amount);
 
     gov.deposit(address(this), amount);
 

test/RollupNonCompliantAsset.t.sol

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.27;
+
+import {TestBase} from "@test/base/Base.sol";
+import {FailingERC20} from "@test/mocks/FailingERC20.sol";
+import {RollupBuilder} from "@test/builder/RollupBuilder.sol";
+import {IRollup} from "@aztec/core/interfaces/IRollup.sol";
+import {stdStorage, StdStorage} from "forge-std/Test.sol";
+
+contract RollupNonCompliantAssetTest is TestBase {
+  using stdStorage for StdStorage;
+
+  IRollup internal rollup;
+  FailingERC20 internal failingAsset;
+  StdStorage private stdstore;
+  address internal sequencer = address(0xBEEF);
+
+  function setUp() public {
+    failingAsset = new FailingERC20("fail", "FAIL", address(this));
+
+    RollupBuilder builder = new RollupBuilder(address(this)).setSlashingQuorum(1).setSlashingRoundSize(1);
+    builder.setTestERC20(failingAsset);
+    builder.deploy();
+
+    rollup = IRollup(address(builder.getConfig().rollup));
+    failingAsset = FailingERC20(address(builder.getConfig().testERC20));
+
+    address governance = address(builder.getConfig().governance);
+    vm.prank(governance);
+    rollup.setRewardsClaimable(true);
+  }
+
+  function test_ClaimSequencerRewardsRevertsWhenTransferFails() external {
+    uint256 amount = 1e18;
+
+    stdstore.target(address(rollup)).sig(IRollup.getSequencerRewards.selector).with_key(sequencer).checked_write(amount);
+
+    deal(address(failingAsset), address(rollup), amount);
+
+    failingAsset.setFailTransfer(true);
+
+    vm.expectRevert();
+    rollup.claimSequencerRewards(sequencer);
+  }
+}

test/builder/GseBuilder.sol

@@ -49,6 +49,11 @@ contract GSEBuilder is TestBase {
     config.flags.checkProofOfPossession = false;
   }
 
+  function setTestERC20(TestERC20 _testERC20) public returns (GSEBuilder) {
+    config.testERC20 = _testERC20;
+    return this;
+  }
+
   function setOpenFloodgates(bool _openFloodgates) public returns (GSEBuilder) {
     config.flags.openFloodgates = _openFloodgates;
     return this;

test/governance/gse/gse/nonCompliantAsset.t.sol

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.27;
+
+import {WithGSE} from "./base.sol";
+import {FailingERC20} from "@test/mocks/FailingERC20.sol";
+import {GSEBuilder} from "@test/builder/GseBuilder.sol";
+import {BN254Lib} from "@aztec/shared/libraries/BN254Lib.sol";
+
+contract GSENonCompliantAssetTest is WithGSE {
+  FailingERC20 internal failingAsset;
+  address internal rollup;
+
+  function setUp() public override {
+    failingAsset = new FailingERC20("fail", "FAIL", address(this));
+
+    GSEBuilder builder = new GSEBuilder();
+    builder.setTestERC20(failingAsset);
+    builder.deploy();
+
+    gse = builder.getConfig().gse;
+    stakingAsset = TestERC20(address(failingAsset));
+    governance = builder.getConfig().governance;
+
+    rollup = address(0x1234);
+    vm.prank(gse.owner());
+    gse.addRollup(rollup);
+  }
+
+  function test_DepositRevertsWhenTransferOrApproveFails() external {
+    uint256 activationThreshold = gse.ACTIVATION_THRESHOLD();
+
+    vm.prank(stakingAsset.owner());
+    stakingAsset.mint(rollup, activationThreshold);
+
+    vm.prank(rollup);
+    stakingAsset.approve(address(gse), activationThreshold);
+
+    failingAsset.setFailTransferFrom(true);
+
+    vm.prank(rollup);
+    vm.expectRevert();
+    gse.deposit(address(1), address(2), BN254Lib.g1Zero(), BN254Lib.g2Zero(), BN254Lib.g1Zero(), false);
+
+    failingAsset.setFailTransferFrom(false);
+    failingAsset.setFailApprove(true);
+
+    vm.prank(rollup);
+    vm.expectRevert();
+    gse.deposit(address(1), address(2), BN254Lib.g1Zero(), BN254Lib.g2Zero(), BN254Lib.g1Zero(), false);
+  }
+}

test/mocks/FailingERC20.sol

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity >=0.8.27;
+
+import {TestERC20} from "@aztec/mock/TestERC20.sol";
+
+contract FailingERC20 is TestERC20 {
+  bool public failTransfer;
+  bool public failTransferFrom;
+  bool public failApprove;
+
+  constructor(string memory _name, string memory _symbol, address _owner)
+    TestERC20(_name, _symbol, _owner)
+  {}
+
+  function setFailTransfer(bool _shouldFail) external {
+    failTransfer = _shouldFail;
+  }
+
+  function setFailTransferFrom(bool _shouldFail) external {
+    failTransferFrom = _shouldFail;
+  }
+
+  function setFailApprove(bool _shouldFail) external {
+    failApprove = _shouldFail;
+  }
+
+  function transfer(address _to, uint256 _value) public override returns (bool) {
+    bool success = super.transfer(_to, _value);
+    return failTransfer ? false : success;
+  }
+
+  function transferFrom(address _from, address _to, uint256 _value) public override returns (bool) {
+    bool success = super.transferFrom(_from, _to, _value);
+    return failTransferFrom ? false : success;
+  }
+
+  function approve(address _spender, uint256 _value) public override returns (bool) {
+    bool success = super.approve(_spender, _value);
+    return failApprove ? false : success;
+  }
+}

test/staking/nonCompliantAsset.t.sol

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.27;
+
+import {TestERC20} from "@aztec/mock/TestERC20.sol";
+import {StakingBase} from "./base.t.sol";
+import {FailingERC20} from "@test/mocks/FailingERC20.sol";
+import {RollupBuilder} from "@test/builder/RollupBuilder.sol";
+import {RollupConfigInput} from "@aztec/core/interfaces/IRollup.sol";
+import {BN254Lib} from "@aztec/shared/libraries/BN254Lib.sol";
+
+contract StakingNonCompliantAssetTest is StakingBase {
+  FailingERC20 internal failingAsset;
+
+  function setUp() public override {
+    failingAsset = new FailingERC20("fail", "FAIL", address(this));
+
+    RollupBuilder builder = new RollupBuilder(address(this)).setSlashingQuorum(1).setSlashingRoundSize(1);
+    builder.setTestERC20(failingAsset);
+    builder.deploy();
+
+    registry = builder.getConfig().registry;
+
+    RollupConfigInput memory rollupConfig = builder.getConfig().rollupConfigInput;
+    EPOCH_DURATION_SECONDS = rollupConfig.aztecEpochDuration * rollupConfig.aztecSlotDuration;
+
+    staking = IStaking(address(builder.getConfig().rollup));
+    stakingAsset = TestERC20(address(failingAsset));
+
+    ACTIVATION_THRESHOLD = staking.getActivationThreshold();
+    EJECTION_THRESHOLD = staking.getEjectionThreshold();
+    SLASHER = staking.getSlasher();
+  }
+
+  function test_DepositRevertsWhenTransferFromReturnsFalse() external {
+    mint(address(this), ACTIVATION_THRESHOLD);
+    stakingAsset.approve(address(staking), ACTIVATION_THRESHOLD);
+
+    failingAsset.setFailTransferFrom(true);
+
+    vm.expectRevert();
+    staking.deposit({
+      _attester: ATTESTER,
+      _withdrawer: WITHDRAWER,
+      _publicKeyInG1: BN254Lib.g1Zero(),
+      _publicKeyInG2: BN254Lib.g2Zero(),
+      _proofOfPossession: BN254Lib.g1Zero(),
+      _moveWithLatestRollup: true
+    });
+  }
+}