restrict returnExcessETHToVault to ops

PufferFinance · Sep 24, 2024 · 4e850e1 · 4e850e1
1 parent f8a4bd4
commit 4e850e1
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 8 deletions.
diff --git a/mainnet-contracts/script/AccessManagerMigrations/04_Generate2StepWithdrawalsCalldata.s.sol b/mainnet-contracts/script/AccessManagerMigrations/04_Generate2StepWithdrawalsCalldata.s.sol
@@ -9,7 +9,8 @@ import {
     PUBLIC_ROLE,
     ROLE_ID_DAO,
     ROLE_ID_VAULT_WITHDRAWER,
-    ROLE_ID_PUFETH_BURNER
+    ROLE_ID_PUFETH_BURNER,
+    ROLE_ID_OPERATIONS_MULTISIG
 } from "../../script/Roles.sol";
 import { PufferWithdrawalManager } from "../../src/PufferWithdrawalManager.sol";
 import { PufferVaultV2 } from "../../src/PufferVaultV2.sol";
@@ -22,7 +23,7 @@ contract Generate2StepWithdrawalsCalldata is Script {
         address paymaster,
         address withdrawalFinalizer
     ) public pure returns (bytes memory) {
-        bytes[] memory calldatas = new bytes[](14);
+        bytes[] memory calldatas = new bytes[](15);
 
         bytes4[] memory paymasterSelectors = new bytes4[](1);
         paymasterSelectors[0] = PufferWithdrawalManager.finalizeWithdrawals.selector;
@@ -34,11 +35,10 @@ contract Generate2StepWithdrawalsCalldata is Script {
         );
 
         // Everybody can complete queued withdrawals
-        bytes4[] memory publicSelectors = new bytes4[](4);
+        bytes4[] memory publicSelectors = new bytes4[](3);
         publicSelectors[0] = PufferWithdrawalManager.completeQueuedWithdrawal.selector;
         publicSelectors[1] = PufferWithdrawalManager.requestWithdrawal.selector;
         publicSelectors[2] = PufferWithdrawalManager.requestWithdrawalWithPermit.selector;
-        publicSelectors[3] = PufferWithdrawalManager.returnExcessETHToVault.selector;
         calldatas[1] = abi.encodeWithSelector(
             AccessManager.setTargetFunctionRole.selector, withdrawalManagerProxy, publicSelectors, PUBLIC_ROLE
         );
@@ -101,6 +101,15 @@ contract Generate2StepWithdrawalsCalldata is Script {
         calldatas[13] =
             abi.encodeWithSelector(AccessManager.grantRole.selector, ROLE_ID_PUFETH_BURNER, pufferProtocolProxy, 0);
 
+        bytes4[] memory opsSelectors = new bytes4[](1);
+        opsSelectors[0] = PufferWithdrawalManager.returnExcessETHToVault.selector;
+        calldatas[14] = abi.encodeWithSelector(
+            AccessManager.setTargetFunctionRole.selector,
+            withdrawalManagerProxy,
+            opsSelectors,
+            ROLE_ID_OPERATIONS_MULTISIG
+        );
+
         bytes memory encodedMulticall = abi.encodeCall(Multicall.multicall, (calldatas));
 
         return encodedMulticall;

diff --git a/mainnet-contracts/script/DeployPufferWithdrawalManager.s.sol b/mainnet-contracts/script/DeployPufferWithdrawalManager.s.sol
@@ -16,7 +16,7 @@ contract DeployPufferWithdrawalManager is DeployerHelper {
     PufferWithdrawalManager public withdrawalManager;
     bytes public encodedCalldata;
 
-    uint256 public BATCH_SIZE = 10; // @todo figure out a good batch size
+    uint256 public BATCH_SIZE = 10;
 
     function run() public {
         Generate2StepWithdrawalsCalldata calldataGenerator = new Generate2StepWithdrawalsCalldata();

diff --git a/mainnet-contracts/src/PufferWithdrawalManager.sol b/mainnet-contracts/src/PufferWithdrawalManager.sol
@@ -183,7 +183,7 @@ contract PufferWithdrawalManager is
 
     /**
      * @inheritdoc IPufferWithdrawalManager
-     * @dev Restricted in this context is like `whenNotPaused` modifier from Pausable.sol
+     * @dev Restricted access to ROLE_ID_WITHDRAWAL_FINALIZER
      */
     function completeQueuedWithdrawal(uint256 withdrawalIdx) external restricted {
         WithdrawalManagerStorage storage $ = _getWithdrawalManagerStorage();
@@ -226,7 +226,7 @@ contract PufferWithdrawalManager is
 
     /**
      * @inheritdoc IPufferWithdrawalManager
-     * @dev Restricted in this context is like the `whenNotPaused` modifier from Pausable.sol
+     * @dev Restricted access to ROLE_ID_OPERATIONS_MULTISIG
      */
     function returnExcessETHToVault(uint256[] calldata batchIndices) external restricted {
         WithdrawalManagerStorage storage $ = _getWithdrawalManagerStorage();

diff --git a/mainnet-contracts/test/unit/PufferWithdrawalManager.t.sol b/mainnet-contracts/test/unit/PufferWithdrawalManager.t.sol
@@ -153,6 +153,7 @@ contract PufferWithdrawalManagerTest is UnitTestHelper {
         uint256[] memory batches = new uint256[](1);
         batches[0] = 1;
 
+        vm.startPrank(OPERATIONS_MULTISIG);
         // If there is no excess ETH, this should revert
         vm.expectRevert(IPufferWithdrawalManager.AlreadyReturned.selector);
         withdrawalManager.returnExcessETHToVault(batches);
@@ -587,10 +588,10 @@ contract PufferWithdrawalManagerTest is UnitTestHelper {
         vm.stopPrank();
 
         // Return the dust from the batch
-        vm.startPrank(PAYMASTER);
         uint256[] memory batches = new uint256[](1);
         batches[0] = 1;
 
+        vm.startPrank(OPERATIONS_MULTISIG);
         vm.expectRevert(abi.encodeWithSelector(IPufferWithdrawalManager.NotAllWithdrawalsClaimed.selector));
         withdrawalManager.returnExcessETHToVault(batches);