chore: custom errors for DecryptionOracleCaller

zama-ai · Dec 26, 2024 · 0dc4304 · 0dc4304
1 parent 29800ab
commit 0dc4304
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/contracts/decryptionOracleLib/DecryptionOracleCaller.sol b/contracts/decryptionOracleLib/DecryptionOracleCaller.sol
@@ -23,6 +23,11 @@ struct DecryptionOracleConfigStruct {
 }
 
 abstract contract DecryptionOracleCaller {
+    error HandlesAlreadySavedForRequestID();
+    error NoHandleFoundForRequestID();
+    error InvalidKMSSignatures();
+    error UnsupportedHandleType();
+
     mapping(uint256 => ebool[]) private paramsEBool;
     mapping(uint256 => euint4[]) private paramsEUint4;
     mapping(uint256 => euint8[]) private paramsEUint8;
@@ -73,12 +78,16 @@ abstract contract DecryptionOracleCaller {
     }
 
     function saveRequestedHandles(uint256 requestID, uint256[] memory handlesList) internal {
-        require(requestedHandles[requestID].length == 0, "requested handles already saved");
+        if (requestedHandles[requestID].length != 0) {
+            revert HandlesAlreadySavedForRequestID();
+        }
         requestedHandles[requestID] = handlesList;
     }
 
     function loadRequestedHandles(uint256 requestID) internal view returns (uint256[] memory) {
-        require(requestedHandles[requestID].length != 0, "requested handles were not saved for this requestID");
+        if (requestedHandles[requestID].length == 0) {
+            revert NoHandleFoundForRequestID();
+        }
         return requestedHandles[requestID];
     }
 
@@ -232,7 +241,7 @@ abstract contract DecryptionOracleCaller {
                 //ebytes256
                 signedDataLength += 320;
             } else {
-                revert("Unsupported handle type");
+                revert UnsupportedHandleType();
             }
         }
         signedDataLength += 32; // add offset of signatures
@@ -242,7 +251,9 @@ abstract contract DecryptionOracleCaller {
     modifier checkSignatures(uint256 requestID, bytes[] memory signatures) {
         uint256[] memory handlesList = loadRequestedHandles(requestID);
         bool isVerified = verifySignatures(handlesList, signatures);
-        require(isVerified, "KMS signature verification failed");
+        if (!isVerified) {
+            revert InvalidKMSSignatures();
+        }
         _;
     }
 }
diff --git a/contracts/test/kmsVerifier/kmsVerifier.ts b/contracts/test/kmsVerifier/kmsVerifier.ts
@@ -117,7 +117,7 @@ describe('KMSVerifier', function () {
       process.env.PRIVATE_KEY_KMS_SIGNER_1 = process.env.PRIVATE_KEY_KMS_SIGNER_0;
       const tx7 = await contract.requestUint16();
       await tx7.wait();
-      await expect(awaitAllDecryptionResults()).to.revertedWith('KMS signature verification failed'); // cannot use duplicated signatures if threshold is 2
+      await expect(awaitAllDecryptionResults()).to.revertedWithCustomError(contract, 'InvalidKMSSignatures'); // cannot use duplicated signatures if threshold is 2
       const y5 = await contract.yUint16();
       expect(y5).to.equal(0);