Forge cheatcodes to set confidential inputs (#31)

* Forge cheatcodes to set confidential inputs * Update README
flashbots · Jan 18, 2024 · b8e5c4d · b8e5c4d
1 parent f69be79
commit b8e5c4d
Show file tree

Hide file tree

Showing 6 changed files with 144 additions and 23 deletions.
diff --git a/README.md b/README.md
@@ -98,3 +98,28 @@ contract TestForge is Test, SuaveEnabled {
     }
 }
 ```
+
+### Confidential inputs
+
+Use the `setConfidentialInputs` function to set the confidential inputs during tests.
+
+```solidity
+// SPDX-License-Identifier: Unlicense
+pragma solidity ^0.8.13;
+
+import "forge-std/Test.sol";
+import "src/Test.sol";
+import "src/suavelib/Suave.sol";
+
+contract TestForge is Test, SuaveEnabled {
+    function testConfidentialInputs() public {
+        bytes memory input = hex"abcd";
+        setConfidentialInputs(input);
+
+        bytes memory found2 = Suave.confidentialInputs();
+        assertEq0(input, found2);
+    }
+}
+```
+
+The value for the confidential inputs gets reset for each test.
diff --git a/src/Test.sol b/src/Test.sol
@@ -2,9 +2,17 @@
 pragma solidity ^0.8.8;
 
 import "./forge/Registry.sol";
+import "./suavelib/Suave.sol";
 import "forge-std/Test.sol";
 
+interface ConfidentialInputsWrapperI {
+    function setConfidentialInputs(bytes memory) external;
+    function resetConfidentialInputs() external;
+}
+
 contract SuaveEnabled is Test {
+    ConfidentialInputsWrapperI constant confInputsWrapper = ConfidentialInputsWrapperI(Suave.CONFIDENTIAL_INPUTS);
+
     function setUp() public {
         string[] memory inputs = new string[](3);
         inputs[0] = "suave-geth";
@@ -22,6 +30,12 @@ contract SuaveEnabled is Test {
         }
 
         Registry.enable();
+
+        confInputsWrapper.resetConfidentialInputs();
+    }
+
+    function setConfidentialInputs(bytes memory data) internal {
+        confInputsWrapper.setConfidentialInputs(data);
     }
 
     function detectErrorMessage(bytes memory reason) internal pure returns (string memory) {

diff --git a/src/forge/ConfidentialInputs.sol b/src/forge/ConfidentialInputs.sol
@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.8;
+
+contract ConfidentialInputsWrapper {
+    bytes confidentialInputs;
+
+    function setConfidentialInputs(bytes memory _confidentialInputs) public {
+        confidentialInputs = _confidentialInputs;
+    }
+
+    function resetConfidentialInputs() public {
+        confidentialInputs = "";
+    }
+
+    fallback() external {
+        // copy bytes from storage to memory
+        bytes memory msgdata = new bytes(confidentialInputs.length);
+        for (uint256 i = 0; i < confidentialInputs.length; i++) {
+            msgdata[i] = confidentialInputs[i];
+        }
+
+        assembly {
+            let location := msgdata
+            let length := mload(msgdata)
+            return(add(location, 0x20), length)
+        }
+    }
+}
diff --git a/src/forge/Registry.sol b/src/forge/Registry.sol
@@ -16,12 +16,16 @@ library Registry {
         bytes memory code =
             hex"608060405234801561001057600080fd5b506004361061002b5760003560e01c8063671ff786146100a1575b6040516bffffffffffffffffffffffff193060601b1660208201526000906100959060340160408051808303601f19018152602036601f8101829004820285018201909352828452909291600091819084018382808284376000920191909152506100ca92505050565b90508081518060208301f35b6100b46100af366004610487565b61025e565b6040516100c19190610557565b60405180910390f35b606060006100d78461025e565b905060006100e48461025e565b60408051600480825260a0820190925291925060009190816020015b60608152602001906001900390816101005790505090506040518060400160405280600a8152602001690e6eac2ecca5acecae8d60b31b8152508160008151811061014d5761014d610571565b602002602001018190525060405180604001604052806005815260200164666f72676560d81b8152508160018151811061018957610189610571565b602002602001018190525082816002815181106101a8576101a8610571565b602002602001018190525081816003815181106101c7576101c7610571565b6020908102919091010152604051638916046760e01b8152600090737109709ecfa91a80626ff3989d68f67f5b1dd12d9063891604679061020c908590600401610587565b600060405180830381865afa158015610229573d6000803e3d6000fd5b505050506040513d6000823e601f3d908101601f1916820160405261025191908101906105eb565b9450505050505b92915050565b60606000825160026102709190610678565b67ffffffffffffffff81111561028857610288610418565b6040519080825280601f01601f1916602001820160405280156102b2576020820181803683370190505b5060408051808201909152601081526f181899199a1a9b1b9c1cb0b131b232b360811b602082015290915060005b84518110156103ee578182518683815181106102fe576102fe610571565b0160200151610310919060f81c6106a5565b8151811061032057610320610571565b01602001516001600160f81b0319168361033b836002610678565b8151811061034b5761034b610571565b60200101906001600160f81b031916908160001a90535081825186838151811061037757610377610571565b0160200151610389919060f81c6106b9565b8151811061039957610399610571565b01602001516001600160f81b031916836103b4836002610678565b6103bf9060016106cd565b815181106103cf576103cf610571565b60200101906001600160f81b031916908160001a9053506001016102e0565b508160405160200161040091906106e0565b60405160208183030381529060405292505050919050565b634e487b7160e01b600052604160045260246000fd5b604051601f8201601f1916810167ffffffffffffffff8111828210171561045757610457610418565b604052919050565b600067ffffffffffffffff82111561047957610479610418565b50601f01601f191660200190565b60006020828403121561049957600080fd5b813567ffffffffffffffff8111156104b057600080fd5b8201601f810184136104c157600080fd5b80356104d46104cf8261045f565b61042e565b8181528560208385010111156104e957600080fd5b81602084016020830137600091810160200191909152949350505050565b60005b8381101561052257818101518382015260200161050a565b50506000910152565b60008151808452610543816020860160208601610507565b601f01601f19169290920160200192915050565b60208152600061056a602083018461052b565b9392505050565b634e487b7160e01b600052603260045260246000fd5b600060208083016020845280855180835260408601915060408160051b87010192506020870160005b828110156105de57603f198886030184526105cc85835161052b565b945092850192908501906001016105b0565b5092979650505050505050565b6000602082840312156105fd57600080fd5b815167ffffffffffffffff81111561061457600080fd5b8201601f8101841361062557600080fd5b80516106336104cf8261045f565b81815285602083850101111561064857600080fd5b610659826020830160208601610507565b95945050505050565b634e487b7160e01b600052601160045260246000fd5b808202811582820484141761025857610258610662565b634e487b7160e01b600052601260045260246000fd5b6000826106b4576106b461068f565b500490565b6000826106c8576106c861068f565b500690565b8082018082111561025857610258610662565b61060f60f31b8152600082516106fd816002850160208701610507565b919091016002019291505056fea164736f6c6343000817000a";
         vm.etch(addr, code);
+
+        // enable is confidential wrapper
+        bytes memory confidentialCode =
+            hex"608060405234801561001057600080fd5b50600436106100365760003560e01c8063844f563414610125578063cb9b7f3f1461013a575b600080805461004490610170565b905067ffffffffffffffff81111561005e5761005e6101aa565b6040519080825280601f01601f191660200182016040528015610088576020820181803683370190505b50905060005b6000805461009b90610170565b905081101561011a5760008181546100b290610170565b81106100c0576100c06101c0565b8154600116156100df5790600052602060002090602091828204019190065b9054901a600160f81b028282815181106100fb576100fb6101c0565b60200101906001600160f81b031916908160001a90535060010161008e565b508081518060208301f35b6101386101333660046101d6565b610142565b005b610138610152565b600061014e82826102d8565b5050565b604080516020810190915260008082529061016d90826102d8565b50565b600181811c9082168061018457607f821691505b6020821081036101a457634e487b7160e01b600052602260045260246000fd5b50919050565b634e487b7160e01b600052604160045260246000fd5b634e487b7160e01b600052603260045260246000fd5b6000602082840312156101e857600080fd5b813567ffffffffffffffff8082111561020057600080fd5b818401915084601f83011261021457600080fd5b813581811115610226576102266101aa565b604051601f8201601f19908116603f0116810190838211818310171561024e5761024e6101aa565b8160405282815287602084870101111561026757600080fd5b826020860160208301376000928101602001929092525095945050505050565b601f8211156102d3576000816000526020600020601f850160051c810160208610156102b05750805b601f850160051c820191505b818110156102cf578281556001016102bc565b5050505b505050565b815167ffffffffffffffff8111156102f2576102f26101aa565b610306816103008454610170565b84610287565b602080601f83116001811461033b57600084156103235750858301515b600019600386901b1c1916600185901b1785556102cf565b600085815260208120601f198616915b8281101561036a5788860151825594840194600190910190840161034b565b50858210156103885787850151600019600388901b60f8161c191681555b5050505050600190811b0190555056fea164736f6c6343000817000a";
+        vm.etch(Suave.CONFIDENTIAL_INPUTS, confidentialCode);
     }
 
     function enable() public {
         enableLib(Suave.IS_CONFIDENTIAL_ADDR);
         enableLib(Suave.BUILD_ETH_BLOCK);
-        enableLib(Suave.CONFIDENTIAL_INPUTS);
         enableLib(Suave.CONFIDENTIAL_RETRIEVE);
         enableLib(Suave.CONFIDENTIAL_STORE);
         enableLib(Suave.DO_HTTPREQUEST);

diff --git a/test/Forge.t.sol b/test/Forge.t.sol
@@ -17,4 +17,16 @@ contract TestForge is Test, SuaveEnabled {
         bytes memory found = Suave.confidentialRetrieve(record.id, "key1");
         assertEq(keccak256(found), keccak256(value));
     }
+
+    function testConfidentialInputs() public {
+        // ensure that the confidential inputs are empty
+        bytes memory found = Suave.confidentialInputs();
+        assertEq(found.length, 0);
+
+        bytes memory input = hex"abcd";
+        setConfidentialInputs(input);
+
+        bytes memory found2 = Suave.confidentialInputs();
+        assertEq0(input, found2);
+    }
 }
diff --git a/tools/forge-gen/main.go b/tools/forge-gen/main.go
@@ -54,8 +54,13 @@ library Registry {
     function enableLib(address addr) public {
         // code for Wrapper
         bytes memory code =
-            hex"{{.Bytecode}}";
+            hex"{{.Bytecodes.Connector}}";
         vm.etch(addr, code);
+
+		// enable is confidential wrapper
+		bytes memory confidentialCode =
+			hex"{{.Bytecodes.Confidential}}";
+		vm.etch(Suave.CONFIDENTIAL_INPUTS, confidentialCode);
     }
 
     function enable() public {
@@ -65,14 +70,14 @@ library Registry {
     }
 }`
 
-func applyTemplate(bytecode string, precompileNames []string) error {
+func applyTemplate(bytecodes *forgeWrapperBytecodes, precompileNames []string) error {
 	t, err := template.New("template").Parse(templateFile)
 	if err != nil {
 		return err
 	}
 
 	input := map[string]interface{}{
-		"Bytecode":        bytecode,
+		"Bytecodes":       bytecodes,
 		"PrecompileNames": precompileNames,
 	}
 
@@ -96,37 +101,67 @@ func applyTemplate(bytecode string, precompileNames []string) error {
 	return nil
 }
 
-func getForgeConnectorBytecode() (string, error) {
+type forgeWrapperBytecodes struct {
+	Connector    string
+	Confidential string
+}
+
+func getForgeConnectorBytecode() (*forgeWrapperBytecodes, error) {
+	mirror := func(from, to string) error {
+		connectorSrc, err := os.ReadFile(resolvePath(from))
+		if err != nil {
+			return err
+		}
+		if err := writeFile(resolvePath(to), connectorSrc); err != nil {
+			return err
+		}
+		return nil
+	}
+
 	// mirror the Connector.sol contract to ./src
-	connectorSrc, err := os.ReadFile(resolvePath("../../src/forge/Connector.sol"))
-	if err != nil {
-		return "", err
+	if err := mirror("../../src/forge/Connector.sol", "./src-forge-test/Connector.sol"); err != nil {
+		return nil, err
 	}
-	if err := writeFile(resolvePath("./src-forge-test/Connector.sol"), connectorSrc); err != nil {
-		return "", err
+	// mirror the is confidential solver
+	if err := mirror("../../src/forge/ConfidentialInputs.sol", "./src-forge-test/ConfidentialInputs.sol"); err != nil {
+		return nil, err
 	}
 
 	// compile the Connector contract with forge and the local configuration
 	if _, err := execForgeCommand([]string{"build", "--config-path", resolvePath("./foundry.toml")}, ""); err != nil {
-		return "", err
+		return nil, err
 	}
 
-	abiContent, err := os.ReadFile(resolvePath("./out/Connector.sol/Connector.json"))
-	if err != nil {
-		return "", err
-	}
+	decodeBytecode := func(name string) (string, error) {
+		abiContent, err := os.ReadFile(resolvePath(name))
+		if err != nil {
+			return "", err
+		}
 
-	var abiArtifact struct {
-		DeployedBytecode struct {
-			Object string
+		var abiArtifact struct {
+			DeployedBytecode struct {
+				Object string
+			}
 		}
+		if err := json.Unmarshal(abiContent, &abiArtifact); err != nil {
+			return "", err
+		}
+
+		bytecode := abiArtifact.DeployedBytecode.Object[2:]
+		return bytecode, nil
 	}
-	if err := json.Unmarshal(abiContent, &abiArtifact); err != nil {
-		return "", err
+
+	res := &forgeWrapperBytecodes{}
+	var err error
+
+	if res.Connector, err = decodeBytecode("./out/Connector.sol/Connector.json"); err != nil {
+		return nil, err
+	}
+	if res.Confidential, err = decodeBytecode("./out/ConfidentialInputs.sol/ConfidentialInputsWrapper.json"); err != nil {
+		return nil, err
 	}
 
-	bytecode := abiArtifact.DeployedBytecode.Object[2:]
-	return bytecode, nil
+	return res, nil
 }
 
 func getPrecompileNames() ([]string, error) {
@@ -146,6 +181,9 @@ func getPrecompileNames() ([]string, error) {
 			if name == "ANYALLOWED" {
 				continue
 			}
+			if name == "CONFIDENTIAL_INPUTS" {
+				continue
+			}
 			names = append(names, name)
 		}
 	}
@@ -178,7 +216,7 @@ func execForgeCommand(args []string, stdin string) (string, error) {
 
 	// Run the command
 	if err := cmd.Run(); err != nil {
-		return "", fmt.Errorf("error running command: %v", err)
+		return "", fmt.Errorf("error running command: %v, %s", err, errBuf.String())
 	}
 
 	return outBuf.String(), nil