Merge pull request #181 from zama-ai/fixGatewayContract

fix: make GatewayContract compatible with current KMS in trustless case

contracts/gateway/GatewayContract.sol

@@ -24,7 +24,7 @@ contract GatewayContract is UUPSUpgradeable, Ownable2StepUpgradeable {
     /// @notice Version of the contract
     uint256 private constant MAJOR_VERSION = 0;
     uint256 private constant MINOR_VERSION = 1;
-    uint256 private constant PATCH_VERSION = 0;
+    uint256 private constant PATCH_VERSION = 1;
 
     IKMSVerifier private constant kmsVerifier = IKMSVerifier(kmsVerifierAdd);
     address private constant aclAddress = aclAdd;
@@ -184,35 +184,35 @@ contract GatewayContract is UUPSUpgradeable, Ownable2StepUpgradeable {
         bytes[] memory signatures
     ) external payable virtual onlyRelayer {
         GatewayContractStorage storage $ = _getGatewayContractStorage();
-        require(
-            kmsVerifier.verifyDecryptionEIP712KMSSignatures(
-                aclAddress,
-                $.decryptionRequests[requestID].cts,
-                decryptedCts,
-                signatures
-            ),
-            "KMS signature verification failed"
-        );
         require(!$.isFulfilled[requestID], "Request is already fulfilled");
+        $.isFulfilled[requestID] = true;
         DecryptionRequest memory decryptionReq = $.decryptionRequests[requestID];
         require(block.timestamp <= decryptionReq.maxTimestamp, "Too late");
         bytes memory callbackCalldata = abi.encodeWithSelector(decryptionReq.callbackSelector, requestID);
         bool passSignatures = decryptionReq.passSignaturesToCaller;
         callbackCalldata = abi.encodePacked(callbackCalldata, decryptedCts); // decryptedCts MUST be correctly abi-encoded by the relayer, according to the requested types of `ctsHandles`
         if (passSignatures) {
             bytes memory packedSignatures = abi.encode(signatures);
-            bytes memory packedSignaturesNoOffset = removeOffset(packedSignatures); // remove the offset (the first 32 bytes) before concatenating with the first part of calldata
-            callbackCalldata = abi.encodePacked(callbackCalldata, packedSignaturesNoOffset);
+            packedSignatures = removeOffset(packedSignatures);
+            callbackCalldata = replaceOffsets(callbackCalldata, decryptionReq.cts, packedSignatures);
+        } else {
+            require(
+                kmsVerifier.verifyDecryptionEIP712KMSSignatures(
+                    aclAddress,
+                    $.decryptionRequests[requestID].cts,
+                    decryptedCts,
+                    signatures
+                ),
+                "KMS signature verification failed"
+            );
         }
-
         (bool success, bytes memory result) = (decryptionReq.contractCaller).call{value: decryptionReq.msgValue}(
             callbackCalldata
         );
         emit ResultCallback(requestID, success, result);
-        $.isFulfilled[requestID] = true;
     }
 
-    function removeOffset(bytes memory input) public pure virtual returns (bytes memory) {
+    function removeOffset(bytes memory input) internal pure virtual returns (bytes memory) {
         uint256 newLength = input.length - 32;
         bytes memory result = new bytes(newLength);
         for (uint256 i = 0; i < newLength; i++) {
@@ -221,6 +221,74 @@ contract GatewayContract is UUPSUpgradeable, Ownable2StepUpgradeable {
         return result;
     }
 
+    function replaceOffsets(
+        bytes memory input,
+        uint256[] memory handlesList,
+        bytes memory packedSignatures
+    ) internal pure virtual returns (bytes memory) {
+        uint256 numArgs = handlesList.length;
+        uint256 signaturesOffset = 64; // requestID is always first argument of callback + own size of offset = 32+32
+        for (uint256 i = 0; i < numArgs; i++) {
+            uint8 typeCt = uint8(handlesList[i] >> 8);
+            if (typeCt >= 9) {
+                input = addToBytes32Slice(input, 32 * (i + 1) + 4); // because we append the signatures, all bytes offsets are shifted by 0x20
+                if (typeCt == 9) {
+                    //ebytes64
+                    signaturesOffset += 128;
+                } else if (typeCt == 10) {
+                    //ebytes128
+                    signaturesOffset += 192;
+                } else if (typeCt == 11) {
+                    //ebytes256
+                    signaturesOffset += 320;
+                }
+            } else {
+                signaturesOffset += 32;
+            }
+        }
+        input = interleaveBytes(input, 4 + 32 * (numArgs + 1), signaturesOffset); // we add the offset of the signatures at correct location
+        input = abi.encodePacked(input, packedSignatures);
+        return input;
+    }
+
+    function addToBytes32Slice(bytes memory data, uint256 offset) internal pure virtual returns (bytes memory) {
+        // @note: data is assumed to be more than 32+offset bytes long
+        assembly {
+            let ptr := add(add(data, 0x20), offset)
+            let val := mload(ptr)
+            val := add(val, 0x20)
+            mstore(ptr, val)
+        }
+        return data;
+    }
+
+    function interleaveBytes(
+        bytes memory input,
+        uint256 position,
+        uint256 valueToPutINBetween
+    ) internal pure virtual returns (bytes memory) {
+        // @note: we assume position <= input.length
+        uint256 originalLength = input.length;
+        uint256 newLength = originalLength + 32;
+        bytes memory result = new bytes(newLength);
+        for (uint256 i = 0; i < position; i++) {
+            result[i] = input[i];
+        }
+
+        // Insert 32 bytes with the specified value
+        bytes32 valueToInsert = bytes32(valueToPutINBetween);
+        for (uint256 i = 0; i < 32; i++) {
+            result[position + i] = valueToInsert[i];
+        }
+
+        // Copy remainder of input (after insertion point)
+        for (uint256 i = position; i < originalLength; i++) {
+            result[i + 32] = input[i];
+        }
+
+        return result;
+    }
+
     /// @notice Getter for the name and version of the contract
     /// @return string representing the name and the version of the contract
     function getVersion() external pure virtual returns (string memory) {

contracts/gatewayLib/lib/Gateway.sol

@@ -2,9 +2,8 @@
 
 pragma solidity ^0.8.24;
 
-import "../../addresses/GatewayContractAddress.sol";
 import "../../lib/Impl.sol";
-import "../../addresses/ACLAddress.sol";
+import "../../addresses/GatewayContractAddress.sol";
 
 interface IKMSVerifier {
     function verifyDecryptionEIP712KMSSignatures(
@@ -25,24 +24,24 @@ interface IGatewayContract {
     ) external returns (uint256);
 }
 
-library Gateway {
-    struct GatewayConfigStruct {
-        address GatewayContractAddress;
-    }
+struct GatewayConfigStruct {
+    address GatewayContractAddress;
+}
 
+library Gateway {
     // keccak256(abi.encode(uint256(keccak256("fhevm.storage.GatewayConfig")) - 1)) & ~bytes32(uint256(0xff))
     bytes32 private constant GatewayLocation = 0x93ab6e17f2c461cce6ea5d4ec117e51dda77a64affc2b2c05f8cd440def0e700;
 
-    function defaultGatewayAddress() internal pure returns (address) {
-        return GATEWAY_CONTRACT_PREDEPLOY_ADDRESS;
-    }
-
     function getGetwayConfig() internal pure returns (GatewayConfigStruct storage $) {
         assembly {
             $.slot := GatewayLocation
         }
     }
 
+    function defaultGatewayAddress() internal pure returns (address) {
+        return GATEWAY_CONTRACT_PREDEPLOY_ADDRESS;
+    }
+
     function setGateway(address gatewayAddress) internal {
         GatewayConfigStruct storage $ = getGetwayConfig();
         $.GatewayContractAddress = gatewayAddress;
@@ -130,10 +129,12 @@ library Gateway {
         assembly {
             calldatacopy(add(decryptedResult, 0x20), start, length) // Copy the relevant part of calldata to decryptedResult memory
         }
+
+        decryptedResult = shiftOffsets(decryptedResult, handlesList);
         FHEVMConfig.FHEVMConfigStruct storage $ = Impl.getFHEVMConfig();
         return
             IKMSVerifier($.KMSVerifierAddress).verifyDecryptionEIP712KMSSignatures(
-                aclAdd,
+                $.ACLAddress,
                 handlesList,
                 decryptedResult,
                 signatures
@@ -160,7 +161,45 @@ library Gateway {
                 revert("Unsupported handle type");
             }
         }
-        signedDataLength += 32; // for the signatures offset
+        signedDataLength += 32; // add offset of signatures
         return signedDataLength;
     }
+
+    function shiftOffsets(bytes memory input, uint256[] memory handlesList) private pure returns (bytes memory) {
+        uint256 numArgs = handlesList.length;
+        for (uint256 i = 0; i < numArgs; i++) {
+            uint8 typeCt = uint8(handlesList[i] >> 8);
+            if (typeCt >= 9) {
+                input = subToBytes32Slice(input, 32 * i); // because we append the signatures, all bytes offsets are shifted by 0x20
+            }
+        }
+        input = remove32Slice(input, 32 * numArgs);
+        return input;
+    }
+
+    function subToBytes32Slice(bytes memory data, uint256 offset) private pure returns (bytes memory) {
+        // @note: data is assumed to be more than 32+offset bytes long
+        assembly {
+            let ptr := add(add(data, 0x20), offset)
+            let val := mload(ptr)
+            val := sub(val, 0x20)
+            mstore(ptr, val)
+        }
+        return data;
+    }
+
+    function remove32Slice(bytes memory input, uint256 start) public pure returns (bytes memory) {
+        // @note we assume start+32 is less than input.length
+        bytes memory result = new bytes(input.length - 32);
+
+        for (uint256 i = 0; i < start; i++) {
+            result[i] = input[i];
+        }
+
+        for (uint256 i = start + 32; i < input.length; i++) {
+            result[i - 32] = input[i];
+        }
+
+        return result;
+    }
 }

contracts/package.json

@@ -1,6 +1,6 @@
 {
   "name": "fhevm-core-contracts",
-  "version": "0.6.0",
+  "version": "0.6.1-0",
   "description": "fhEVM contracts",
   "repository": {
     "type": "git",

contracts/test/asyncDecrypt.ts

@@ -122,7 +122,6 @@ const fulfillAllPastRequestsIds = async (mocked: boolean) => {
     const handles = event.args[1];
     const typesList = handles.map((handle) => parseInt(handle.toString(16).slice(-4, -2), 16));
     const msgValue = event.args[4];
-    const passSignaturesToCaller = event.args[6];
     if (!results.includes(requestID)) {
       // if request is not already fulfilled
       if (mocked) {
@@ -154,13 +153,8 @@ const fulfillAllPastRequestsIds = async (mocked: boolean) => {
         const abiCoder = new ethers.AbiCoder();
         let encodedData;
         let calldata;
-        if (!passSignaturesToCaller) {
-          encodedData = abiCoder.encode(['uint256', ...types], [31, ...valuesFormatted4]); // 31 is just a dummy uint256 requestID to get correct abi encoding for the remaining arguments (i.e everything except the requestID)
-          calldata = '0x' + encodedData.slice(66); // we just pop the dummy requestID to get the correct value to pass for `decryptedCts`
-        } else {
-          encodedData = abiCoder.encode(['uint256', ...types, 'bytes[]'], [31, ...valuesFormatted4, []]); // adding also a dummy empty array of bytes for correct abi-encoding when used with signatures
-          calldata = '0x' + encodedData.slice(66).slice(0, -64); // we also pop the last 32 bytes (empty bytes[])
-        }
+        encodedData = abiCoder.encode(['uint256', ...types], [31, ...valuesFormatted4]); // 31 is just a dummy uint256 requestID to get correct abi encoding for the remaining arguments (i.e everything except the requestID)
+        calldata = '0x' + encodedData.slice(66); // we just pop the dummy requestID to get the correct value to pass for `decryptedCts`
 
         const numSigners = +process.env.NUM_KMS_SIGNERS!;
         const decryptResultsEIP712signatures = await computeDecryptSignatures(handles, calldata, numSigners);

contracts/test/upgrades/upgrades.ts

@@ -84,7 +84,7 @@ describe('Upgrades', function () {
       kind: 'uups',
     });
     await gateway.waitForDeployment();
-    expect(await gateway.getVersion()).to.equal('GatewayContract v0.1.0');
+    expect(await gateway.getVersion()).to.equal('GatewayContract v0.1.1');
     const gateway2 = await upgrades.upgradeProxy(gateway, this.gatewayFactoryUpgraded);
     await gateway2.waitForDeployment();
     expect(await gateway2.getVersion()).to.equal('GatewayContract v0.2.0');

fhevm-engine/coprocessor/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "coprocessor"
-version = "0.6.0"
+version = "0.6.1-alpha.0"
 default-run = "coprocessor"
 authors.workspace = true
 edition.workspace = true

fhevm-engine/executor/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "executor"
-version = "0.6.0"
+version = "0.6.1-alpha.0"
 authors.workspace = true
 edition.workspace = true
 license.workspace = true

fhevm-engine/fhevm-engine-common/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "fhevm-engine-common"
-version = "0.6.0"
+version = "0.6.1-alpha.0"
 authors.workspace = true
 edition.workspace = true
 license.workspace = true

fhevm-engine/scheduler/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "scheduler"
-version = "0.6.0"
+version = "0.6.1-alpha.0"
 edition = "2021"
 
 [dependencies]