refactor: [hexsen-s3] update PoolManagerOwner to support 2 step poolM…

…anager ownership transfer
pancakeswap · Nov 7, 2024 · b4bc2a0 · b4bc2a0
1 parent 60db3ff
commit b4bc2a0
Show file tree

Hide file tree

Showing 6 changed files with 147 additions and 10 deletions.
diff --git a/src/base/PoolManagerOwnable2Step.sol b/src/base/PoolManagerOwnable2Step.sol
@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+// Copyright (C) 2024 PancakeSwap
+pragma solidity 0.8.26;
+
+import {IPoolManagerOwner} from "../interfaces/IPoolManagerOwner.sol";
+
+/**
+ * @notice dev This contract implements "Ownable2Step styled" poolManager ownership transfer
+ * functionality. Namely an extra acceptance step is added to the ownership transfer process.
+ * This is done to prevent accidental ownership transfer to a wrong address.
+ */
+abstract contract PoolManagerOwnable2Step is IPoolManagerOwner {
+    error NotPendingPoolManagerOwner();
+
+    address private _pendingPoolManagerOwner;
+
+    modifier isFromPendingPoolManagerOwner() {
+        if (_pendingPoolManagerOwner != msg.sender) {
+            revert NotPendingPoolManagerOwner();
+        }
+
+        _;
+    }
+
+    function _transferPoolManagerOwnership(address newPoolManagerOwner) internal {
+        _pendingPoolManagerOwner = newPoolManagerOwner;
+    }
+
+    /// @inheritdoc IPoolManagerOwner
+    function pendingPoolManagerOwner() public view override returns (address) {
+        return _pendingPoolManagerOwner;
+    }
+}
diff --git a/src/interfaces/IPoolManagerOwner.sol b/src/interfaces/IPoolManagerOwner.sol
@@ -18,6 +18,14 @@ interface IPoolManagerOwner {
 
     /// @notice transfer the ownership of pool manager to the new owner
     /// @dev used when a new PoolManagerOwner contract is created and we transfer pool manager owner to new contract
-    /// @param newOwner the address of the new owner
-    function transferPoolManagerOwnership(address newOwner) external;
+    /// @param newPoolManagerOwner the address of the new owner
+    function transferPoolManagerOwnership(address newPoolManagerOwner) external;
+
+    /// @notice accept the ownership of pool manager, only callable by the
+    /// pending pool manager owner set by latest transferPoolManagerOwnership
+    function acceptPoolManagerOwnership() external;
+
+    /// @notice get the current pending pool manager owner
+    /// @return the address of the pending pool manager owner
+    function pendingPoolManagerOwner() external view returns (address);
 }
diff --git a/src/pool-bin/BinPoolManagerOwner.sol b/src/pool-bin/BinPoolManagerOwner.sol
@@ -4,6 +4,7 @@ pragma solidity 0.8.26;
 
 import {IProtocolFeeController} from "../interfaces/IProtocolFeeController.sol";
 import {PausableRole} from "../base/PausableRole.sol";
+import {PoolManagerOwnable2Step} from "../base/PoolManagerOwnable2Step.sol";
 import {IBinPoolManager} from "./interfaces/IBinPoolManager.sol";
 import {IPoolManagerOwner} from "../interfaces/IPoolManagerOwner.sol";
 
@@ -20,7 +21,7 @@ interface IBinPoolManagerWithPauseOwnable is IBinPoolManager {
  * A seperate owner contract is used to handle some functionality so as to reduce the contract size
  * of PoolManager. This allow a higher optimizer run, reducing the gas cost for other poolManager functions
  */
-contract BinPoolManagerOwner is IPoolManagerOwner, PausableRole {
+contract BinPoolManagerOwner is IPoolManagerOwner, PoolManagerOwnable2Step, PausableRole {
     /// @notice Error thrown when owner set min share too small
     error MinShareTooSmall(uint256 minShare);
 
@@ -46,8 +47,13 @@ contract BinPoolManagerOwner is IPoolManagerOwner, PausableRole {
     }
 
     /// @inheritdoc IPoolManagerOwner
-    function transferPoolManagerOwnership(address newOwner) external override onlyOwner {
-        poolManager.transferOwnership(newOwner);
+    function transferPoolManagerOwnership(address newPoolManagerOwner) external override onlyOwner {
+        _transferPoolManagerOwnership(newPoolManagerOwner);
+    }
+
+    /// @inheritdoc IPoolManagerOwner
+    function acceptPoolManagerOwnership() external override isFromPendingPoolManagerOwner {
+        poolManager.transferOwnership(msg.sender);
     }
 
     /// @notice Set max bin steps for binPoolManager, see IBinPoolManager for more documentation about this function

diff --git a/src/pool-cl/CLPoolManagerOwner.sol b/src/pool-cl/CLPoolManagerOwner.sol
@@ -4,6 +4,7 @@ pragma solidity 0.8.26;
 
 import {IProtocolFeeController} from "../interfaces/IProtocolFeeController.sol";
 import {PausableRole} from "../base/PausableRole.sol";
+import {PoolManagerOwnable2Step} from "../base/PoolManagerOwnable2Step.sol";
 import {ICLPoolManager} from "./interfaces/ICLPoolManager.sol";
 import {IPoolManagerOwner} from "../interfaces/IPoolManagerOwner.sol";
 
@@ -20,7 +21,7 @@ interface ICLPoolManagerWithPauseOwnable is ICLPoolManager {
  * A seperate owner contract is used to handle some functionality so as to reduce the contract size
  * of PoolManager. This allow a higher optimizer run, reducing the gas cost for other poolManager functions.
  */
-contract CLPoolManagerOwner is IPoolManagerOwner, PausableRole {
+contract CLPoolManagerOwner is IPoolManagerOwner, PoolManagerOwnable2Step, PausableRole {
     ICLPoolManagerWithPauseOwnable public immutable poolManager;
 
     constructor(ICLPoolManagerWithPauseOwnable _poolManager) {
@@ -43,7 +44,12 @@ contract CLPoolManagerOwner is IPoolManagerOwner, PausableRole {
     }
 
     /// @inheritdoc IPoolManagerOwner
-    function transferPoolManagerOwnership(address newOwner) external override onlyOwner {
-        poolManager.transferOwnership(newOwner);
+    function transferPoolManagerOwnership(address newPoolManagerOwner) external override onlyOwner {
+        _transferPoolManagerOwnership(newPoolManagerOwner);
+    }
+
+    /// @inheritdoc IPoolManagerOwner
+    function acceptPoolManagerOwnership() external override isFromPendingPoolManagerOwner {
+        poolManager.transferOwnership(msg.sender);
     }
 }
diff --git a/test/pool-bin/BinPoolManagerOwner.t.sol b/test/pool-bin/BinPoolManagerOwner.t.sol
@@ -7,6 +7,7 @@ import {IVault} from "../../src/interfaces/IVault.sol";
 import {Vault} from "../../src/Vault.sol";
 import {BinPoolManager} from "../../src/pool-bin/BinPoolManager.sol";
 import {BinPoolManagerOwner, IBinPoolManagerWithPauseOwnable} from "../../src/pool-bin/BinPoolManagerOwner.sol";
+import {PoolManagerOwnable2Step} from "../../src/base/PoolManagerOwnable2Step.sol";
 import {Pausable} from "../../src/base/Pausable.sol";
 import {PausableRole} from "../../src/base/PausableRole.sol";
 
@@ -110,8 +111,15 @@ contract BinPoolManagerOwnerTest is Test {
         // before:
         assertEq(poolManager.owner(), address(binPoolManagerOwner));
 
-        // after:
+        // pending:
+        // it's still the original owner if new owner not accept yet
         binPoolManagerOwner.transferPoolManagerOwnership(alice);
+        assertEq(poolManager.owner(), address(binPoolManagerOwner));
+        assertEq(binPoolManagerOwner.pendingPoolManagerOwner(), alice);
+
+        // after:
+        vm.prank(alice);
+        binPoolManagerOwner.acceptPoolManagerOwnership();
         assertEq(poolManager.owner(), alice);
     }
 
@@ -122,6 +130,40 @@ contract BinPoolManagerOwnerTest is Test {
         binPoolManagerOwner.transferPoolManagerOwnership(alice);
     }
 
+    function test_TransferPoolManagerOwnership_NotPendingPoolManagerOwner() public {
+        binPoolManagerOwner.transferPoolManagerOwnership(alice);
+
+        // if it's not from alice then revert
+        vm.expectRevert(PoolManagerOwnable2Step.NotPendingPoolManagerOwner.selector);
+        binPoolManagerOwner.acceptPoolManagerOwnership();
+    }
+
+    function test_TransferPoolManagerOwnership_OverridePendingOwner() public {
+        // before:
+        assertEq(poolManager.owner(), address(binPoolManagerOwner));
+
+        // pending:
+        // it's still the original owner if new owner not accept yet
+        binPoolManagerOwner.transferPoolManagerOwnership(alice);
+        assertEq(poolManager.owner(), address(binPoolManagerOwner));
+        assertEq(binPoolManagerOwner.pendingPoolManagerOwner(), alice);
+
+        // override pending owner
+        binPoolManagerOwner.transferPoolManagerOwnership(makeAddr("bob"));
+        assertEq(poolManager.owner(), address(binPoolManagerOwner));
+        assertEq(binPoolManagerOwner.pendingPoolManagerOwner(), makeAddr("bob"));
+
+        // alice no longer the valid pending owner
+        vm.expectRevert(PoolManagerOwnable2Step.NotPendingPoolManagerOwner.selector);
+        vm.prank(alice);
+        binPoolManagerOwner.acceptPoolManagerOwnership();
+
+        // bob is the new owner
+        vm.prank(makeAddr("bob"));
+        binPoolManagerOwner.acceptPoolManagerOwnership();
+        assertEq(poolManager.owner(), makeAddr("bob"));
+    }
+
     function test_SetMaxBinStep_OnlyOwner() public {
         // before
         assertEq(poolManager.maxBinStep(), 100);

diff --git a/test/pool-cl/CLPoolManagerOwner.t.sol b/test/pool-cl/CLPoolManagerOwner.t.sol
@@ -7,6 +7,7 @@ import {Deployers} from "./helpers/Deployers.sol";
 import {IVault} from "../../src/interfaces/IVault.sol";
 import {CLPoolManager} from "../../src/pool-cl/CLPoolManager.sol";
 import {CLPoolManagerOwner, ICLPoolManagerWithPauseOwnable} from "../../src/pool-cl/CLPoolManagerOwner.sol";
+import {PoolManagerOwnable2Step} from "../../src/base/PoolManagerOwnable2Step.sol";
 import {Pausable} from "../../src/base/Pausable.sol";
 import {PausableRole} from "../../src/base/PausableRole.sol";
 
@@ -108,8 +109,15 @@ contract CLPoolManagerOwnerTest is Test, Deployers {
         // before:
         assertEq(poolManager.owner(), address(clPoolManagerOwner));
 
-        // after:
+        // pending:
+        // it's still the original owner if new owner not accept yet
         clPoolManagerOwner.transferPoolManagerOwnership(alice);
+        assertEq(poolManager.owner(), address(clPoolManagerOwner));
+        assertEq(clPoolManagerOwner.pendingPoolManagerOwner(), alice);
+
+        // after:
+        vm.prank(alice);
+        clPoolManagerOwner.acceptPoolManagerOwnership();
         assertEq(poolManager.owner(), alice);
     }
 
@@ -119,4 +127,38 @@ contract CLPoolManagerOwnerTest is Test, Deployers {
         vm.prank(alice);
         clPoolManagerOwner.transferPoolManagerOwnership(alice);
     }
+
+    function test_TransferPoolManagerOwnership_NotPendingPoolManagerOwner() public {
+        clPoolManagerOwner.transferPoolManagerOwnership(alice);
+
+        // if it's not from alice then revert
+        vm.expectRevert(PoolManagerOwnable2Step.NotPendingPoolManagerOwner.selector);
+        clPoolManagerOwner.acceptPoolManagerOwnership();
+    }
+
+    function test_TransferPoolManagerOwnership_OverridePendingOwner() public {
+        // before:
+        assertEq(poolManager.owner(), address(clPoolManagerOwner));
+
+        // pending:
+        // it's still the original owner if new owner not accept yet
+        clPoolManagerOwner.transferPoolManagerOwnership(alice);
+        assertEq(poolManager.owner(), address(clPoolManagerOwner));
+        assertEq(clPoolManagerOwner.pendingPoolManagerOwner(), alice);
+
+        // override pending owner
+        clPoolManagerOwner.transferPoolManagerOwnership(makeAddr("bob"));
+        assertEq(poolManager.owner(), address(clPoolManagerOwner));
+        assertEq(clPoolManagerOwner.pendingPoolManagerOwner(), makeAddr("bob"));
+
+        // alice no longer the valid pending owner
+        vm.expectRevert(PoolManagerOwnable2Step.NotPendingPoolManagerOwner.selector);
+        vm.prank(alice);
+        clPoolManagerOwner.acceptPoolManagerOwnership();
+
+        // bob is the new owner
+        vm.prank(makeAddr("bob"));
+        clPoolManagerOwner.acceptPoolManagerOwnership();
+        assertEq(poolManager.owner(), makeAddr("bob"));
+    }
 }