FleetTreasuryPaymaster — whitelisted contract destinations

src/paymasters/FleetTreasuryPaymaster.sol → src/paymasters/BondTreasuryPaymaster.sol

@@ -7,63 +7,81 @@ import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol
 import {BasePaymaster} from "./BasePaymaster.sol";
 import {QuotaControl} from "../QuotaControl.sol";
 
-/// @notice Combined paymaster + bond treasury for FleetIdentity operations.
-/// @dev Holds ETH (to sponsor gas) and NODL (to sponsor bonds). Whitelisted
-///      users call FleetIdentityUpgradeable.claimUuidSponsored(), which calls
-///      this contract's `consumeSponsoredBond` to validate + consume quota,
-///      then pulls the NODL bond via `transferFrom`.
-///      The ZkSync paymaster validation ensures only whitelisted users calling
-///      FleetIdentity get gas-sponsored.
-contract FleetTreasuryPaymaster is BasePaymaster, QuotaControl {
+/// @notice ZkSync paymaster plus ERC-20 bond treasury for whitelisted contracts (e.g. `IBondTreasury` implementers).
+/// @dev Holds ETH (sponsored gas) and bond token balance. Whitelisted users interact with `isWhitelistedContract`
+///      destinations; bond consumers call `consumeSponsoredBond`, which checks caller + user whitelist, quota, balance,
+///      then `forceApprove(msg.sender, amount)` for `transferFrom`.
+///      Gas: `isWhitelistedContract[to] && isWhitelistedUser[from]`. Constructor seeds `address(this)` so management
+///      txs can be sponsored once those EOAs are user-whitelisted. On-chain mutators use `WHITELIST_ADMIN_ROLE` /
+///      `WITHDRAWER_ROLE` / `DEFAULT_ADMIN_ROLE`.
+contract BondTreasuryPaymaster is BasePaymaster, QuotaControl {
     using SafeERC20 for IERC20;
 
     bytes32 public constant WHITELIST_ADMIN_ROLE = keccak256("WHITELIST_ADMIN_ROLE");
 
-    address public immutable fleetIdentity;
     IERC20 public immutable bondToken;
 
     mapping(address => bool) public isWhitelistedUser;
+    /// @notice Sponsored tx destinations and allowed `consumeSponsoredBond` callers.
+    mapping(address => bool) public isWhitelistedContract;
 
     event WhitelistedUsersAdded(address[] users);
     event WhitelistedUsersRemoved(address[] users);
+    event WhitelistedContractsAdded(address[] contracts);
+    event WhitelistedContractsRemoved(address[] contracts);
     event TokensWithdrawn(address indexed token, address indexed to, uint256 amount);
 
     error UserIsNotWhitelisted();
-    error DestinationNotAllowed();
+    error DestIsNotWhitelisted();
     error PaymasterBalanceTooLow();
-    error NotFleetIdentity();
+    error CallerNotWhitelistedContract();
     error InsufficientBondBalance();
+    error ZeroAddress();
 
     constructor(
         address admin,
         address withdrawer,
-        address fleetIdentity_,
+        address[] memory initialWhitelistedContracts,
         address bondToken_,
         uint256 initialQuota,
         uint256 initialPeriod
     ) BasePaymaster(admin, withdrawer) QuotaControl(initialQuota, initialPeriod, admin) {
+        if (admin == address(0) || withdrawer == address(0)) revert ZeroAddress();
+        if (bondToken_ == address(0)) revert ZeroAddress();
+
         _grantRole(WHITELIST_ADMIN_ROLE, admin);
-        fleetIdentity = fleetIdentity_;
         bondToken = IERC20(bondToken_);
+        uint256 n = initialWhitelistedContracts.length;
+        for (uint256 i = 0; i < n; i++) {
+            if (initialWhitelistedContracts[i] == address(0)) revert ZeroAddress();
+            isWhitelistedContract[initialWhitelistedContracts[i]] = true;
+        }
+        if (n > 0) {
+            emit WhitelistedContractsAdded(initialWhitelistedContracts);
+        }
+        if (!isWhitelistedContract[address(this)]) {
+            isWhitelistedContract[address(this)] = true;
+            address[] memory selfDest = new address[](1);
+            selfDest[0] = address(this);
+            emit WhitelistedContractsAdded(selfDest);
+        }
     }
 
     // ──────────────────────────────────────────────
-    // Bond Treasury (called by FleetIdentity)
+    // Bond Treasury (whitelisted contracts)
     // ──────────────────────────────────────────────
 
     /// @notice Validate whitelist + consume quota for a sponsored bond.
-    /// @dev Only callable by the FleetIdentity contract during claimUuidSponsored.
-    ///      The actual NODL transfer is done separately by FleetIdentity via transferFrom.
+    /// @dev Callable only by `isWhitelistedContract`. Caller pulls via `transferFrom`.
     function consumeSponsoredBond(address user, uint256 amount) external {
-        if (msg.sender != fleetIdentity) revert NotFleetIdentity();
+        if (!isWhitelistedContract[msg.sender]) revert CallerNotWhitelistedContract();
         if (!isWhitelistedUser[user]) revert UserIsNotWhitelisted();
         if (bondToken.balanceOf(address(this)) < amount) revert InsufficientBondBalance();
 
         _checkedResetClaimed();
         _checkedUpdateClaimed(amount);
 
-        // Approve only the exact amount needed for this claim
-        bondToken.forceApprove(fleetIdentity, amount);
+        bondToken.forceApprove(msg.sender, amount);
     }
 
     // ──────────────────────────────────────────────
@@ -73,6 +91,7 @@ contract FleetTreasuryPaymaster is BasePaymaster, QuotaControl {
     function addWhitelistedUsers(address[] calldata users) external {
         _checkRole(WHITELIST_ADMIN_ROLE);
         for (uint256 i = 0; i < users.length; i++) {
+            if (users[i] == address(0)) revert ZeroAddress();
             isWhitelistedUser[users[i]] = true;
         }
         emit WhitelistedUsersAdded(users);
@@ -86,13 +105,31 @@ contract FleetTreasuryPaymaster is BasePaymaster, QuotaControl {
         emit WhitelistedUsersRemoved(users);
     }
 
+    function addWhitelistedContracts(address[] calldata contracts_) external {
+        _checkRole(WHITELIST_ADMIN_ROLE);
+        for (uint256 i = 0; i < contracts_.length; i++) {
+            if (contracts_[i] == address(0)) revert ZeroAddress();
+            isWhitelistedContract[contracts_[i]] = true;
+        }
+        emit WhitelistedContractsAdded(contracts_);
+    }
+
+    function removeWhitelistedContracts(address[] calldata contracts_) external {
+        _checkRole(WHITELIST_ADMIN_ROLE);
+        for (uint256 i = 0; i < contracts_.length; i++) {
+            isWhitelistedContract[contracts_[i]] = false;
+        }
+        emit WhitelistedContractsRemoved(contracts_);
+    }
+
     // ──────────────────────────────────────────────
     // ERC-20 Withdrawal
     // ──────────────────────────────────────────────
 
-    /// @notice Withdraw ERC-20 tokens (e.g. excess NODL) from this contract.
+    /// @notice Withdraw ERC-20 tokens (e.g. excess bond token) from this contract.
     function withdrawTokens(address token, address to, uint256 amount) external {
         _checkRole(WITHDRAWER_ROLE);
+        if (token == address(0) || to == address(0)) revert ZeroAddress();
         IERC20(token).safeTransfer(to, amount);
         emit TokensWithdrawn(token, to, amount);
     }
@@ -102,13 +139,8 @@ contract FleetTreasuryPaymaster is BasePaymaster, QuotaControl {
     // ──────────────────────────────────────────────
 
     function _validateAndPayGeneralFlow(address from, address to, uint256 requiredETH) internal view override {
-        if (to == fleetIdentity) {
-            if (!isWhitelistedUser[from]) revert UserIsNotWhitelisted();
-        } else if (to == address(this)) {
-            if (!hasRole(WHITELIST_ADMIN_ROLE, from)) revert DestinationNotAllowed();
-        } else {
-            revert DestinationNotAllowed();
-        }
+        if (!isWhitelistedContract[to]) revert DestIsNotWhitelisted();
+        if (!isWhitelistedUser[from]) revert UserIsNotWhitelisted();
         if (address(this).balance < requiredETH) revert PaymasterBalanceTooLow();
     }