fix: Update Access Control to OwnableTwoSteps

LooksRare · Jul 4, 2024 · 9ca51b7 · 9ca51b7
1 parent d1c8d3f
commit 9ca51b7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 14 deletions.
diff --git a/contracts/BlastYield.sol b/contracts/BlastYield.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.20;
 
-import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
+import {OwnableTwoSteps} from "./OwnableTwoSteps.sol";
 import {IBlast, YieldMode as IBlast__YieldMode, GasMode as IBlast__GasMode} from "./interfaces/IBlast.sol";
 import {IBlastPoints} from "./interfaces/IBlastPoints.sol";
 import {IERC20Rebasing, YieldMode as IERC20Rebasing__YieldMode} from "./interfaces/IERC20Rebasing.sol";
@@ -12,10 +12,9 @@ import {IERC20Rebasing, YieldMode as IERC20Rebasing__YieldMode} from "./interfac
  * @notice This contract is a base contract for future contracts that wish to claim Blast WETH or USDB yield to inherit from.
  * @author LooksRare protocol team (👀,💎)
  */
-contract BlastYield is AccessControl {
+contract BlastYield is OwnableTwoSteps {
     address public immutable WETH;
     address public immutable USDB;
-    bytes32 private constant OPERATOR_ROLE = keccak256("OPERATOR_ROLE");
 
     /**
      * @param _blast Blast precompile
@@ -32,10 +31,7 @@ contract BlastYield is AccessControl {
         address _owner,
         address _usdb,
         address _weth
-    ) {
-        _grantRole(OPERATOR_ROLE, _owner);
-        _grantRole(OPERATOR_ROLE, _blastPointsOperator);
-
+    ) OwnableTwoSteps(_owner) {
         WETH = _weth;
         USDB = _usdb;
 
@@ -50,7 +46,7 @@ contract BlastYield is AccessControl {
      * @param wethReceiver The receiver of WETH.
      * @param usdbReceiver The receiver of USDB.
      */
-    function claim(address wethReceiver, address usdbReceiver) external virtual onlyRole(OPERATOR_ROLE) {
+    function claim(address wethReceiver, address usdbReceiver) external virtual onlyOwner {
         uint256 claimableWETH = IERC20Rebasing(WETH).getClaimableAmount(address(this));
         if (claimableWETH != 0) {
             IERC20Rebasing(WETH).claim(wethReceiver, claimableWETH);

diff --git a/test/foundry/BlastYield.t.sol b/test/foundry/BlastYield.t.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.20;
 
-import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
+import {IOwnableTwoSteps} from "../../contracts/interfaces/IOwnableTwoSteps.sol";
 import {BlastYield} from "../../contracts/BlastYield.sol";
 import {Test} from "../../lib/forge-std/src/Test.sol";
 import {YieldMode as IBlast__YieldMode, GasMode as IBlast__GasMode} from "../../contracts/interfaces/IBlast.sol";
@@ -38,8 +38,7 @@ contract BlastYield_Test is Test {
     function test_setUpState() public {
         assertEq(blastYield.WETH(), address(weth));
         assertEq(blastYield.USDB(), address(usdb));
-        assertTrue(blastYield.hasRole(OPERATOR_ROLE, owner));
-        assertTrue(blastYield.hasRole(OPERATOR_ROLE, operator));
+        assertEq(blastYield.owner(), owner);
 
         (IBlast__YieldMode yieldMode, IBlast__GasMode gasMode, address governor) = mockYield.config(
             address(blastYield)
@@ -65,9 +64,7 @@ contract BlastYield_Test is Test {
     }
 
     function test_claim_RevertIf_NotOwner() public asPrankedUser(user1) {
-        vm.expectRevert(
-            abi.encodeWithSelector(IAccessControl.AccessControlUnauthorizedAccount.selector, user1, OPERATOR_ROLE)
-        );
+        vm.expectRevert(IOwnableTwoSteps.NotOwner.selector);
         blastYield.claim(TREASURY, TREASURY);
     }