fix: one-shot ossify and ACL in ForkSelector

lidofinance · Feb 8, 2024 · 79d770d · 79d770d
1 parent 8dc7a63
commit 79d770d
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/src/ForkSelector.sol b/src/ForkSelector.sol
@@ -11,12 +11,30 @@ contract ForkSelector is IForkSelector {
     Slot[] public versionsLookup;
     Slot public terminalSlot = Slot.wrap(type(uint64).max);
 
+    address internal _admin;
+
     error NoSuitableForkVersion(Slot slot);
     error UnexpectedOrder();
+    error Unauthorized(address sender);
+    error ZeroAddress();
+    error Ossified();
+
+    function initialize(address admin) external {
+        if (admin == address(0)) {
+            revert ZeroAddress();
+        }
+
+        _admin = admin;
+    }
 
     /// @dev If any fork introduces a changed generalized index, we need to add it here.
     /// @dev The list of `versionsLookup` is expected to be sorted in ascending order.
-    function addForkAtSlot(ForkVersion fork, Slot slot) external onlyDao {
+    function addForkAtSlot(ForkVersion fork, Slot slot) external onlyAdmin {
+        // Allow adding forks at slots before the terminal slot.
+        if (Slot.unwrap(slot) > Slot.unwrap(terminalSlot)) {
+            revert Ossified();
+        }
+
         if (
             versionsLookup.length > 0 &&
             Slot.unwrap(versionsLookup[versionsLookup.length - 1]) >=
@@ -29,8 +47,7 @@ contract ForkSelector is IForkSelector {
         versionsLookup.push(slot);
     }
 
-    // TODO: Make it on-shot operation of leave it as is?
-    function ossifyAtSlot(Slot slot) external onlyDao {
+    function ossifyAtSlot(Slot slot) external onlyAdmin notOssified {
         terminalSlot = slot;
     }
 
@@ -50,8 +67,17 @@ contract ForkSelector is IForkSelector {
         revert NoSuitableForkVersion(slot);
     }
 
-    modifier onlyDao() {
-        // FIXME: implement.
+    modifier onlyAdmin() {
+        if (msg.sender != _admin) {
+            revert Unauthorized(msg.sender);
+        }
+        _;
+    }
+
+    modifier notOssified() {
+        if (Slot.unwrap(terminalSlot) != type(uint64).max) {
+            revert Ossified();
+        }
         _;
     }
 }
diff --git a/test/helpers/mocks/ForkSelectorMock.sol b/test/helpers/mocks/ForkSelectorMock.sol
@@ -6,8 +6,16 @@ import { ForkVersion, Slot } from "../../../src/lib/Types.sol";
 import { ForkSelector } from "../../../src/ForkSelector.sol";
 
 contract ForkSelectorMock is ForkSelector {
-    function usePrater() external {
+    function usePrater() external prankAdmin {
         this.addForkAtSlot(ForkVersion.wrap(0x03001020), Slot.wrap(5193728));
         this.addForkAtSlot(ForkVersion.wrap(0x04001020), Slot.wrap(7413760));
     }
+
+    // Used to bypass ACL while making external calls.
+    modifier prankAdmin() {
+        address oldAdmin = _admin;
+        _admin = address(this);
+        _;
+        _admin = oldAdmin;
+    }
 }