fix: set condOsp with postUpgradeInit

OffchainLabs · Apr 16, 2024 · e5ba177 · e5ba177
1 parent f05cd14
commit e5ba177
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 18 deletions.
diff --git a/src/challenge/ChallengeManager.sol b/src/challenge/ChallengeManager.sol
@@ -111,20 +111,19 @@ contract ChallengeManager is DelegateCallAware, IChallengeManager {
         osp = osp_;
     }
 
-    function postUpgradeInit(IOneStepProofEntry osp_) external onlyDelegated onlyProxyOwner {
-        // when updating to 4844 we need to create new osp contracts and set them here
-        // on the challenge manager
+    function postUpgradeInit(
+        IOneStepProofEntry osp_,
+        bytes32 condRoot,
+        IOneStepProofEntry condOsp
+    ) external onlyDelegated onlyProxyOwner {
+        emit ConditonalOSPSet(condRoot, condOsp);
+        ospCond[condRoot] = condOsp;
         osp = osp_;
     }
 
-    function setConditionalOsp(bytes32 wasmModuleRoot, IOneStepProofEntry osp_) external onlyDelegated onlyProxyOwner {
-        emit ConditonalOSPSet(wasmModuleRoot, osp_);
-        ospCond[wasmModuleRoot] = osp_;
-    }
-
-    function getOSP(bytes32 wasmModuleRoot) public view returns (IOneStepProofEntry) {
+    function getOsp(bytes32 wasmModuleRoot) public view returns (IOneStepProofEntry) {
         IOneStepProofEntry t = ospCond[wasmModuleRoot];
-        if (address(t) == address(0)){
+        if (address(t) == address(0)) {
             return osp;
         } else {
             return t;
@@ -274,7 +273,7 @@ contract ChallengeManager is DelegateCallAware, IChallengeManager {
             require(challengeLength == 1, "TOO_LONG");
         }
 
-        bytes32 afterHash = getOSP(challenge.wasmModuleRoot).proveOneStep(
+        bytes32 afterHash = getOsp(challenge.wasmModuleRoot).proveOneStep(
             ExecutionContext({maxInboxMessagesRead: challenge.maxInboxMessages, bridge: bridge}),
             challengeStart,
             selection.oldSegments[selection.challengePosition],

diff --git a/src/challenge/IChallengeManager.sol b/src/challenge/IChallengeManager.sol
@@ -48,9 +48,18 @@ interface IChallengeManager {
         IOneStepProofEntry osp_
     ) external;
 
-    function setConditionalOsp(bytes32 wasmModuleRoot, IOneStepProofEntry osp_) external;
+    function postUpgradeInit(
+        IOneStepProofEntry osp_,
+        bytes32 condRoot,
+        IOneStepProofEntry condOsp
+    ) external;
+
+    /// @notice Get the default osp, which is used for all wasm module roots that don't have a conditional OSP set
+    ///         Use getOsp(wasmModuleRoot) to get the OSP for a specific wasm module root
+    function osp() external view returns (IOneStepProofEntry);
 
-    function getOSP(bytes32 wasmModuleRoot) external view returns (IOneStepProofEntry);
+    /// @notice Get the OSP for a given wasm module root
+    function getOsp(bytes32 wasmModuleRoot) external view returns (IOneStepProofEntry);
 
     function createChallenge(
         bytes32 wasmModuleRoot_,

diff --git a/test/foundry/ChallengeManager.t.sol b/test/foundry/ChallengeManager.t.sol
@@ -11,9 +11,12 @@ contract ChallengeManagerTest is Test {
     IBridge bridge = IBridge(address(139));
     IOneStepProofEntry osp = IOneStepProofEntry(address(140));
     IOneStepProofEntry newOsp = IOneStepProofEntry(address(141));
+    IOneStepProofEntry condOsp = IOneStepProofEntry(address(142));
     address proxyAdmin = address(141);
     ChallengeManager chalmanImpl = new ChallengeManager();
 
+    bytes32 randomRoot = keccak256(abi.encodePacked("randomRoot"));
+
     function deploy() public returns (ChallengeManager) {
         ChallengeManager chalman = ChallengeManager(
             address(new TransparentUpgradeableProxy(address(chalmanImpl), proxyAdmin, ""))
@@ -40,23 +43,29 @@ contract ChallengeManagerTest is Test {
         vm.prank(proxyAdmin);
         TransparentUpgradeableProxy(payable(address(chalman))).upgradeToAndCall(
             address(chalmanImpl),
-            abi.encodeWithSelector(ChallengeManager.postUpgradeInit.selector, newOsp)
+            abi.encodeWithSelector(
+                ChallengeManager.postUpgradeInit.selector,
+                newOsp,
+                randomRoot,
+                condOsp
+            )
         );
 
-        assertEq(address(chalman.osp()), address(newOsp), "New osp not set");
+        assertEq(address(chalman.getOsp(bytes32(0))), address(newOsp), "New osp not set");
+        assertEq(address(chalman.getOsp(randomRoot)), address(condOsp), "Cond osp not set");
     }
 
     function testPostUpgradeInitFailsNotAdmin() public {
         ChallengeManager chalman = deploy();
 
         vm.expectRevert(abi.encodeWithSelector(NotOwner.selector, address(151), proxyAdmin));
         vm.prank(address(151));
-        chalman.postUpgradeInit(osp);
+        chalman.postUpgradeInit(newOsp, randomRoot, condOsp);
     }
 
     function testPostUpgradeInitFailsNotDelCall() public {
         vm.expectRevert(bytes("Function must be called through delegatecall"));
         vm.prank(proxyAdmin);
-        chalmanImpl.postUpgradeInit(osp);
+        chalmanImpl.postUpgradeInit(newOsp, randomRoot, condOsp);
     }
 }
diff --git a/test/signatures/ChallengeManager b/test/signatures/ChallengeManager
@@ -7,11 +7,13 @@
   "clearChallenge(uint64)": "56e9df97",
   "createChallenge(bytes32,uint8[2],(bytes32[2],uint64[2])[2],uint64,address,address,uint256,uint256)": "14eab5e7",
   "currentResponder(uint64)": "23a9ef23",
+  "getOsp(bytes32)": "3690b011",
   "initialize(address,address,address,address)": "f8c8765e",
   "isTimedOut(uint64)": "9ede42b9",
   "oneStepProveExecution(uint64,(uint256,uint256,bytes32[],uint256),bytes)": "d248d124",
   "osp()": "f26a62c6",
-  "postUpgradeInit(address)": "c474d2c5",
+  "ospCond(bytes32)": "dc74bf8b",
+  "postUpgradeInit(address,bytes32,address)": "5038934d",
   "resultReceiver()": "3504f1d7",
   "sequencerInbox()": "ee35f327",
   "timeout(uint64)": "1b45c86a",

diff --git a/test/storage/ChallengeManager b/test/storage/ChallengeManager
@@ -6,3 +6,4 @@
 | sequencerInbox         | contract ISequencerInbox                          | 3    | 0      | 20    | src/challenge/ChallengeManager.sol:ChallengeManager |
 | bridge                 | contract IBridge                                  | 4    | 0      | 20    | src/challenge/ChallengeManager.sol:ChallengeManager |
 | osp                    | contract IOneStepProofEntry                       | 5    | 0      | 20    | src/challenge/ChallengeManager.sol:ChallengeManager |
+| ospCond                | mapping(bytes32 => contract IOneStepProofEntry)   | 6    | 0      | 32    | src/challenge/ChallengeManager.sol:ChallengeManager |