Merge pull request #501 from zama-ai/feat/decryptSigning

feat: added kms signature in mocked mode

.env.example

@@ -3,6 +3,11 @@ export MNEMONIC="adapt mosquito move limb mobile illegal tree voyage juice mosqu
 export PRIVATE_KEY_GATEWAY_DEPLOYER="717fd99986df414889fd8b51069d4f90a50af72e542c58ee065f5883779099c6"
 export PRIVATE_KEY_GATEWAY_OWNER="717fd99986df414889fd8b51069d4f90a50af72e542c58ee065f5883779099c6"
 export PRIVATE_KEY_GATEWAY_RELAYER="7ec931411ad75a7c201469a385d6f18a325d4923f9f213bd882bbea87e160b67"
+export NUM_KMS_SIGNERS="1"
+export PRIVATE_KEY_KMS_SIGNER_0="26698d458a21b843aa1ddbd5c5b098821ddf4218bb52498c4aad3a84849275bb"
+export PRIVATE_KEY_KMS_SIGNER_1="e5b998ce1e664718772fa35e8d02b2b6a267a03a8ecadab15de5b125da7fa82b"
+export PRIVATE_KEY_KMS_SIGNER_2="dca817bfe824b12c92d61e56056b956617da156bcd730379cb9203c822c9ba8e"
+export PRIVATE_KEY_KMS_SIGNER_3="7ac1a2886ca07b3b7393ea5ff3613bb94d72129e2c7cbedc807eb55ff971394c"
 
 # Block explorer API keys
 export ARBISCAN_API_KEY="zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz"

examples/TestAsyncDecrypt.sol

@@ -345,4 +345,66 @@ contract TestAsyncDecrypt is GatewayCaller {
         yAddress = decAddress;
         yBytes256 = bytesRes;
     }
+
+    function requestEbytes256NonTrivialTrustless(einput inputHandle, bytes calldata inputProof) public {
+        ebytes256 inputNonTrivial = TFHE.asEbytes256(inputHandle, inputProof);
+        uint256[] memory cts = new uint256[](1);
+        cts[0] = Gateway.toUint256(inputNonTrivial);
+        uint256 requestID = Gateway.requestDecryption(
+            cts,
+            this.callbackBytes256Trustless.selector,
+            0,
+            block.timestamp + 100,
+            true
+        );
+        latestRequestID = requestID;
+        saveRequestedHandles(requestID, cts);
+    }
+
+    function callbackBytes256Trustless(
+        uint256 requestID,
+        bytes calldata decryptedInput,
+        bytes[] memory signatures
+    ) public onlyGateway returns (bytes memory) {
+        require(latestRequestID == requestID, "wrong requestID passed by Gateway");
+        uint256[] memory requestedHandles = loadRequestedHandles(latestRequestID);
+        bool isKMSVerified = Gateway.verifySignatures(requestedHandles, signatures);
+        require(isKMSVerified, "KMS did not verify this decryption result");
+        yBytes256 = decryptedInput;
+        return decryptedInput;
+    }
+
+    function requestMixedBytes256Trustless(einput inputHandle, bytes calldata inputProof) public {
+        ebytes256 xBytes256 = TFHE.asEbytes256(inputHandle, inputProof);
+        uint256[] memory cts = new uint256[](3);
+        cts[0] = Gateway.toUint256(xBool);
+        cts[1] = Gateway.toUint256(xBytes256);
+        cts[2] = Gateway.toUint256(xAddress);
+        Gateway.requestDecryption(cts, this.callbackMixedBytes256Trustless.selector, 0, block.timestamp + 100, true);
+        uint256 requestID = Gateway.requestDecryption(
+            cts,
+            this.callbackMixedBytes256Trustless.selector,
+            0,
+            block.timestamp + 100,
+            true
+        );
+        latestRequestID = requestID;
+        saveRequestedHandles(requestID, cts);
+    }
+
+    function callbackMixedBytes256Trustless(
+        uint256 requestID,
+        bool decBool,
+        bytes memory bytesRes,
+        address decAddress,
+        bytes[] memory signatures
+    ) public onlyGateway {
+        require(latestRequestID == requestID, "wrong requestID passed by Gateway");
+        uint256[] memory requestedHandles = loadRequestedHandles(latestRequestID);
+        bool isKMSVerified = Gateway.verifySignatures(requestedHandles, signatures);
+        require(isKMSVerified, "KMS did not verify this decryption result");
+        yBool = decBool;
+        yAddress = decAddress;
+        yBytes256 = bytesRes;
+    }
 }

gateway/GatewayContract.sol

@@ -174,28 +174,39 @@ contract GatewayContract is UUPSUpgradeable, Ownable2StepUpgradeable {
         bytes memory decryptedCts,
         bytes[] memory signatures
     ) external payable virtual onlyRelayer {
-        // TODO: this should be un-commented once KMS will have the signatures implemented
-        //require(
-        //    kmsVerifier.verifySignatures(decryptionRequests[requestID].cts, decryptedCts, signatures),
-        //    "KMS signature verification failed"
-        //);
         GatewayContractStorage storage $ = _getGatewayContractStorage();
+        require(
+            kmsVerifier.verifySignatures($.decryptionRequests[requestID].cts, decryptedCts, signatures),
+            "KMS signature verification failed"
+        );
         require(!$.isFulfilled[requestID], "Request is already fulfilled");
         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) {
-            callbackCalldata = abi.encodePacked(callbackCalldata, abi.encode(signatures));
+            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);
         }
+
         (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) {
+        uint256 newLength = input.length - 32;
+        bytes memory result = new bytes(newLength);
+        for (uint256 i = 0; i < newLength; i++) {
+            result[i] = input[i + 32];
+        }
+        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) {

gateway/lib/Gateway.sol

@@ -100,13 +100,36 @@ library Gateway {
     /// @notice this could be used only when signatures are made available to the callback, i.e when `passSignaturesToCaller` is set to true during request
     function verifySignatures(uint256[] memory handlesList, bytes[] memory signatures) internal returns (bool) {
         uint256 start = 4 + 32; // start position after skipping the selector (4 bytes) and the first argument (index, 32 bytes)
-        uint256 numArgs = handlesList.length; // Number of arguments before signatures
-        uint256 length = numArgs * 32; // TODO: fix the way we compute length in case the type of the handle is an ebytes256 (loop over all handles and add correct length corresponding to each type)
+        uint256 length = getSignedDataLength(handlesList);
         bytes memory decryptedResult = new bytes(length);
         assembly {
             calldatacopy(add(decryptedResult, 0x20), start, length) // Copy the relevant part of calldata to decryptedResult memory
         }
         FHEVMConfig.FHEVMConfigStruct storage $ = Impl.getFHEVMConfig();
         return IKMSVerifier($.KMSVerifierAddress).verifySignatures(handlesList, decryptedResult, signatures);
     }
+
+    function getSignedDataLength(uint256[] memory handlesList) private pure returns (uint256) {
+        uint256 handlesListlen = handlesList.length;
+        uint256 signedDataLength;
+        for (uint256 i = 0; i < handlesListlen; i++) {
+            uint8 typeCt = uint8(handlesList[i] >> 8);
+            if (typeCt < 9) {
+                signedDataLength += 32;
+            } else if (typeCt == 9) {
+                //ebytes64
+                signedDataLength += 128;
+            } else if (typeCt == 10) {
+                //ebytes128
+                signedDataLength += 192;
+            } else if (typeCt == 11) {
+                //ebytes256
+                signedDataLength += 320;
+            } else {
+                revert("Unsupported handle type");
+            }
+        }
+        signedDataLength += 32; // for the signatures offset
+        return signedDataLength;
+    }
 }

hardhat.config.ts

@@ -106,6 +106,7 @@ task('test', async (taskArgs, hre, runSuper) => {
     await hre.run('task:deployTFHEExecutor');
     await hre.run('task:deployKMSVerifier');
     await hre.run('task:deployFHEPayment');
+    await hre.run('task:addSigners', { numSigners: +process.env.NUM_KMS_SIGNERS! });
     await hre.run('task:launchFhevm', { skipGetCoin: false });
   }
   await hre.run('compile:specific', { contract: 'examples' });

lib/KMSVerifier.sol

@@ -8,10 +8,10 @@ import "@openzeppelin/contracts-upgradeable/utils/cryptography/EIP712Upgradeable
 import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import "@openzeppelin/contracts/utils/Strings.sol";
 
-/// @title KMS Verifier for signature verification and verifier management
+/// @title KMS Verifier for signature verification and signers management
 /// @author The developer
-/// @notice This contract allows for the management of verifiers and provides methods to verify signatures
-/// @dev The contract uses OpenZeppelin's SignatureChecker for cryptographic operations
+/// @notice This contract allows for the management of signers and provides methods to verify signatures
+/// @dev The contract uses OpenZeppelin's EIP712Upgradeable for cryptographic operations
 contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradeable {
     struct DecryptionResult {
         uint256[] handlesList;
@@ -31,9 +31,9 @@ contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradea
 
     /// @custom:storage-location erc7201:fhevm.storage.KMSVerifier
     struct KMSVerifierStorage {
-        mapping(address => bool) isVerifier; /// @notice Mapping to keep track of addresses that are verifiers
-        address[] verifiers; /// @notice Array to keep track of all verifiers
-        uint256 threshold; /// @notice The threshold for the number of verifiers required for a signature to be valid
+        mapping(address => bool) isSigner; /// @notice Mapping to keep track of addresses that are signers
+        address[] signers; /// @notice Array to keep track of all signers
+        uint256 threshold; /// @notice The threshold for the number of signers required for a signature to be valid
     }
 
     // keccak256(abi.encode(uint256(keccak256("fhevm.storage.KMSVerifier")) - 1)) & ~bytes32(uint256(0xff))
@@ -48,14 +48,14 @@ contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradea
 
     function _authorizeUpgrade(address _newImplementation) internal virtual override onlyOwner {}
 
-    function isVerifier(address account) public virtual returns (bool) {
+    function isSigner(address account) public virtual returns (bool) {
         KMSVerifierStorage storage $ = _getKMSVerifierStorage();
-        return $.isVerifier[account];
+        return $.isSigner[account];
     }
 
-    function getVerifiers() public view virtual returns (address[] memory) {
+    function getSigners() public view virtual returns (address[] memory) {
         KMSVerifierStorage storage $ = _getKMSVerifierStorage();
-        return $.verifiers;
+        return $.signers;
     }
 
     function getThreshold() public view virtual returns (uint256) {
@@ -67,13 +67,13 @@ contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradea
         return DECRYPTIONRESULT_TYPE;
     }
 
-    /// @notice Emitted when a verifier is added
-    /// @param verifier The address of the verifier that was added
-    event VerifierAdded(address indexed verifier);
+    /// @notice Emitted when a signer is added
+    /// @param signer The address of the signer that was added
+    event SignerAdded(address indexed signer);
 
-    /// @notice Emitted when a verifier is removed
-    /// @param verifier The address of the verifier that was removed
-    event VerifierRemoved(address indexed verifier);
+    /// @notice Emitted when a signer is removed
+    /// @param signer The address of the signer that was removed
+    event SignerRemoved(address indexed signer);
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
@@ -86,23 +86,23 @@ contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradea
         __EIP712_init(CONTRACT_NAME, "1");
     }
 
-    /// @notice Sets the threshold for the number of verifiers required for a signature to be valid
+    /// @notice Sets the threshold for the number of signers required for a signature to be valid
     function applyThreshold() internal virtual {
         KMSVerifierStorage storage $ = _getKMSVerifierStorage();
-        $.threshold = ($.verifiers.length - 1) / 3 + 1;
+        $.threshold = ($.signers.length - 1) / 3 + 1;
     }
 
-    /// @notice Adds a new verifier
-    /// @dev Only the owner can add a verifier
-    /// @param verifier The address to be added as a verifier
-    function addVerifier(address verifier) public virtual onlyOwner {
-        require(verifier != address(0), "KMSVerifier: Address is null");
+    /// @notice Adds a new signer
+    /// @dev Only the owner can add a signer
+    /// @param signer The address to be added as a signer
+    function addSigner(address signer) public virtual onlyOwner {
+        require(signer != address(0), "KMSVerifier: Address is null");
         KMSVerifierStorage storage $ = _getKMSVerifierStorage();
-        require($.isVerifier[verifier], "KMSVerifier: Address is already a verifier");
-        $.isVerifier[verifier] = true;
-        $.verifiers.push(verifier);
+        require(!$.isSigner[signer], "KMSVerifier: Address is already a signer");
+        $.isSigner[signer] = true;
+        $.signers.push(signer);
         applyThreshold();
-        emit VerifierAdded(verifier);
+        emit SignerAdded(signer);
     }
 
     function hashDecryptionResult(DecryptionResult memory decRes) internal view virtual returns (bytes32) {
@@ -118,29 +118,29 @@ contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradea
             );
     }
 
-    /// @notice Removes an existing verifier
-    /// @dev Only the owner can remove a verifier
-    /// @param verifier The address to be removed from verifiers
-    function removeVerifier(address verifier) public virtual onlyOwner {
+    /// @notice Removes an existing signer
+    /// @dev Only the owner can remove a signer
+    /// @param signer The address to be removed from signers
+    function removeSigner(address signer) public virtual onlyOwner {
         KMSVerifierStorage storage $ = _getKMSVerifierStorage();
-        require($.isVerifier[verifier], "KMSVerifier: Address is not a verifier");
+        require($.isSigner[signer], "KMSVerifier: Address is not a signer");
 
-        // Remove verifier from the mapping
-        $.isVerifier[verifier] = false;
+        // Remove signer from the mapping
+        $.isSigner[signer] = false;
 
-        // Find the index of the verifier and remove it from the array
-        for (uint i = 0; i < $.verifiers.length; i++) {
-            if ($.verifiers[i] == verifier) {
-                $.verifiers[i] = $.verifiers[$.verifiers.length - 1]; // Move the last element into the place to delete
-                $.verifiers.pop(); // Remove the last element
+        // Find the index of the signer and remove it from the array
+        for (uint i = 0; i < $.signers.length; i++) {
+            if ($.signers[i] == signer) {
+                $.signers[i] = $.signers[$.signers.length - 1]; // Move the last element into the place to delete
+                $.signers.pop(); // Remove the last element
                 applyThreshold();
-                emit VerifierRemoved(verifier);
+                emit SignerRemoved(signer);
                 return;
             }
         }
     }
 
-    /// @notice recovers the verifier's address from a `signature` and a `message` digest
+    /// @notice recovers the signer's address from a `signature` and a `message` digest
     /// @dev Utilizes ECDSA for actual address recovery
     /// @param message The hash of the message that was signed
     /// @param signature The signature to verify
@@ -178,21 +178,21 @@ contract KMSVerifier is UUPSUpgradeable, Ownable2StepUpgradeable, EIP712Upgradea
         require(numSignatures > 0, "KmsVerifier: no signatures provided");
         KMSVerifierStorage storage $ = _getKMSVerifierStorage();
         require(numSignatures >= $.threshold, "KmsVerifier: at least threshold number of signatures required");
-        address[] memory recoveredVerifiers = new address[](numSignatures);
+        address[] memory recoveredSigners = new address[](numSignatures);
         uint256 uniqueValidCount;
         for (uint256 i = 0; i < numSignatures; i++) {
             address signerRecovered = recoverSigner(message, signatures[i]);
-            if ($.isVerifier[signerRecovered]) {
+            if ($.isSigner[signerRecovered]) {
                 if (!tload(signerRecovered)) {
-                    recoveredVerifiers[uniqueValidCount] = signerRecovered;
+                    recoveredSigners[uniqueValidCount] = signerRecovered;
                     uniqueValidCount++;
                     tstore(signerRecovered, 1);
                 }
             }
             if (uniqueValidCount >= $.threshold) {
                 for (uint256 j = 0; i < uniqueValidCount; i++) {
                     /// @note : clearing transient storage for composability
-                    tstore(recoveredVerifiers[j], 0);
+                    tstore(recoveredSigners[j], 0);
                 }
                 return true;
             }

tasks/taskDeploy.ts

@@ -111,3 +111,18 @@ task('task:deployFHEPayment').setAction(async function (taskArguments: TaskArgum
   }
   console.log('FHEPayment was deployed at address:', address);
 });
+
+task('task:addSigners').setAction(async function (taskArguments: TaskArguments, { ethers }) {
+  const deployer = (await ethers.getSigners())[9];
+  const factory = await ethers.getContractFactory('KMSVerifier', deployer);
+  const kmsAdd = dotenv.parse(fs.readFileSync('lib/.env.kmsverifier')).KMS_VERIFIER_CONTRACT_ADDRESS;
+  const kmsVerifier = await factory.attach(kmsAdd);
+
+  for (let idx = 0; idx < taskArguments.numSigners; idx++) {
+    const privKeySigner = process.env[`PRIVATE_KEY_KMS_SIGNER_${idx}`];
+    const kmsSigner = new ethers.Wallet(privKeySigner).connect(ethers.provider);
+    const tx = await kmsVerifier.addSigner(kmsSigner.address);
+    await tx.wait();
+    console.log(`KMS signer no${idx} was added to KMSVerifier contract`);
+  }
+});