chore: custom errors for DecryptionOracleCaller

chore: fix acl test
zama-ai · Dec 26, 2024 · ea67017 · ea67017
1 parent 29800ab
commit ea67017
Show file tree

Hide file tree

Showing 3 changed files with 73 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/acl/acl.ts b/contracts/test/acl/acl.ts
@@ -0,0 +1,57 @@
+import { expect } from 'chai';
+import { ethers } from 'hardhat';
+
+import { initDecryptionOracle } from '../asyncDecrypt';
+import { createInstances } from '../instance';
+import { getSigners, initSigners } from '../signers';
+
+describe.only('ACL', function () {
+  before(async function () {
+    await initSigners(2);
+    this.signers = await getSigners();
+    this.instances = await createInstances(this.signers);
+    const aclFactory = await ethers.getContractFactory('ACL');
+    await initDecryptionOracle();
+    const acl = await aclFactory.deploy();
+    await acl.waitForDeployment();
+    this.acl = acl;
+    this.tfheAddress = await acl.getTFHEExecutorAddress();
+
+    const amountToDistribute = BigInt(100 * 1e18);
+    await ethers.provider.send('hardhat_impersonateAccount', [this.tfheAddress]);
+    await ethers.provider.send('hardhat_setBalance', [this.tfheAddress, '0x' + amountToDistribute.toString(16)]);
+    this.tfheExecutor = await ethers.getSigner(this.tfheAddress);
+  });
+
+  it('allowTransient() is not persistent', async function () {
+    const randomHandle = 3290232n;
+    const randomAccount = this.signers.bob.address;
+    await this.acl.connect(this.tfheExecutor).allowTransient(randomHandle, randomAccount);
+
+    /// @dev The isAllowed returns false since it is transient.
+    expect(await this.acl.isAllowed(randomHandle, randomAccount)).to.be.eq(false);
+
+    /// @dev The isAllowed returns false since it is transient.
+    expect(await this.acl.allowedTransient(randomHandle, randomAccount)).to.be.eq(false);
+  });
+
+  it('allowTransient() reverts if sender is not allowed', async function () {
+    const randomHandle = 3290232n;
+    const randomAccount = this.signers.alice.address;
+    const sender = this.signers.alice;
+
+    await expect(this.acl.connect(sender).allowTransient(randomHandle, randomAccount))
+      .to.be.revertedWithCustomError(this.acl, 'SenderNotAllowed')
+      .withArgs(sender);
+  });
+
+  it('allow() reverts if sender is not allowed', async function () {
+    const randomHandle = 3290232n;
+    const randomAccount = this.signers.alice.address;
+    const sender = this.signers.alice;
+
+    await expect(this.acl.connect(sender).allow(randomHandle, randomAccount))
+      .to.be.revertedWithCustomError(this.acl, 'SenderNotAllowed')
+      .withArgs(sender);
+  });
+});
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);