Fix to allow address(this) in peekers for forge confstore (#73)

flashbots · Mar 27, 2024 · 1c7c09d · 1c7c09d
1 parent 04dfd3a
commit 1c7c09d
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 12 deletions.
diff --git a/src/forge/ConfidentialStore.sol b/src/forge/ConfidentialStore.sol
@@ -53,9 +53,13 @@ contract ConfidentialStore is Test {
     }
 
     function confidentialStore(Suave.DataId dataId, string memory key, bytes memory value) public {
+        confidentialStore(dataId, key, value, msg.sender);
+    }
+
+    function confidentialStore(Suave.DataId dataId, string memory key, bytes memory value, address sender) public {
         address[] memory allowedStores = dataRecords[dataId].allowedStores;
         for (uint256 i = 0; i < allowedStores.length; i++) {
-            if (allowedStores[i] == msg.sender || allowedStores[i] == Suave.ANYALLOWED) {
+            if (allowedStores[i] == sender || allowedStores[i] == Suave.ANYALLOWED) {
                 dataRecordsContent[dataId][key] = value;
                 return;
             }
@@ -65,9 +69,17 @@ contract ConfidentialStore is Test {
     }
 
     function confidentialRetrieve(Suave.DataId dataId, string memory key) public view returns (bytes memory) {
+        return confidentialRetrieve(dataId, key, msg.sender);
+    }
+
+    function confidentialRetrieve(Suave.DataId dataId, string memory key, address sender)
+        public
+        view
+        returns (bytes memory)
+    {
         address[] memory allowedPeekers = dataRecords[dataId].allowedPeekers;
         for (uint256 i = 0; i < allowedPeekers.length; i++) {
-            if (allowedPeekers[i] == msg.sender || allowedPeekers[i] == Suave.ANYALLOWED) {
+            if (allowedPeekers[i] == sender || allowedPeekers[i] == Suave.ANYALLOWED) {
                 return dataRecordsContent[dataId][key];
             }
         }

diff --git a/src/forge/ConfidentialStoreConnector.sol b/src/forge/ConfidentialStoreConnector.sol
@@ -9,24 +9,27 @@ contract ConfidentialStoreConnector {
         address confidentialStoreAddr = 0x0101010101010101010101010101010101010101;
 
         address addr = address(this);
-        bytes4 sig;
+        bytes memory input;
 
         if (addr == Suave.CONFIDENTIAL_STORE) {
-            sig = ConfidentialStore.confidentialStore.selector;
+            (Suave.DataId dataId, string memory key, bytes memory val) =
+                abi.decode(msg.data, (Suave.DataId, string, bytes));
+
+            input =
+                abi.encodeWithSignature("confidentialStore(bytes16,string,bytes,address)", dataId, key, val, msg.sender);
         } else if (addr == Suave.CONFIDENTIAL_RETRIEVE) {
-            sig = ConfidentialStore.confidentialRetrieve.selector;
+            (Suave.DataId dataId, string memory key) = abi.decode(msg.data, (Suave.DataId, string));
+
+            input = abi.encodeWithSignature("confidentialRetrieve(bytes16,string,address)", dataId, key, msg.sender);
         } else if (addr == Suave.FETCH_DATA_RECORDS) {
-            sig = ConfidentialStore.fetchDataRecords.selector;
+            input = abi.encodePacked(ConfidentialStore.fetchDataRecords.selector, msg.data);
         } else if (addr == Suave.NEW_DATA_RECORD) {
-            sig = ConfidentialStore.newDataRecord.selector;
+            input = abi.encodePacked(ConfidentialStore.newDataRecord.selector, msg.data);
         } else {
             revert("function signature not found in the confidential store");
         }
 
-        bytes memory input = msg.data;
-
-        // call 'confidentialStore' with the selector and the input data.
-        (bool success, bytes memory output) = confidentialStoreAddr.call(abi.encodePacked(sig, input));
+        (bool success, bytes memory output) = confidentialStoreAddr.call(input);
         if (!success) {
             revert("Call to confidentialStore failed");
         }

diff --git a/test/Forge.t.sol b/test/Forge.t.sol
@@ -28,7 +28,18 @@ contract TestForge is Test, SuaveEnabled {
     }
 
     function testForgeConfidentialStoreRecordStore() public {
-        Suave.DataRecord memory record = Suave.newDataRecord(0, addressList, addressList, "namespace");
+        // test with the wildcard
+        _testForgeConfidentialStoreRecordStore(addressList);
+
+        // test with address(this) as the allowed address
+        address[] memory addrList = new address[](1);
+        addrList[0] = address(this);
+
+        _testForgeConfidentialStoreRecordStore(addrList);
+    }
+
+    function _testForgeConfidentialStoreRecordStore(address[] memory addrList) public {
+        Suave.DataRecord memory record = Suave.newDataRecord(0, addrList, addrList, "namespace");
 
         bytes memory value = abi.encode("suave works with forge!");
         Suave.confidentialStore(record.id, "key1", value);