test: added tests for kmsverifier

zama-ai · Sep 13, 2024 · c1a8df6 · c1a8df6
1 parent 4eeb0d2
commit c1a8df6
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 3 deletions.
diff --git a/.env.example b/.env.example
@@ -7,6 +7,7 @@ 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"

diff --git a/tasks/taskDeploy.ts b/tasks/taskDeploy.ts
@@ -119,7 +119,7 @@ task('task:addSigners').setAction(async function (taskArguments: TaskArguments,
   const kmsVerifier = await factory.attach(kmsAdd);
 
   for (let idx = 0; idx < taskArguments.numSigners; idx++) {
-    const privKeySigner = dotenv.parse(fs.readFileSync('.env'))[`PRIVATE_KEY_KMS_SIGNER_${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();

diff --git a/test/asyncDecrypt.ts b/test/asyncDecrypt.ts
@@ -166,9 +166,8 @@ async function computeDecryptSignatures(
 ): Promise<string[]> {
   const signatures: string[] = [];
 
-  const envConfig = dotenv.parse(fs.readFileSync('.env'));
   for (let idx = 0; idx < numSigners; idx++) {
-    const privKeySigner = envConfig[`PRIVATE_KEY_KMS_SIGNER_${idx}`];
+    const privKeySigner = process.env[`PRIVATE_KEY_KMS_SIGNER_${idx}`];
     if (privKeySigner) {
       const kmsSigner = new ethers.Wallet(privKeySigner).connect(ethers.provider);
       const signature = await kmsSign(handlesList, decryptedResult, kmsSigner);

diff --git a/test/kmsVerifier/kmsVerifier.ts b/test/kmsVerifier/kmsVerifier.ts
@@ -0,0 +1,95 @@
+import { expect } from 'chai';
+import dotenv from 'dotenv';
+import fs from 'fs';
+import { ethers } from 'hardhat';
+
+import { asyncDecrypt, awaitAllDecryptionResults } from '../asyncDecrypt';
+import { getSigners, initSigners } from '../signers';
+
+describe('KMSVerifier', function () {
+  before(async function () {
+    await initSigners(2);
+    this.signers = await getSigners();
+    this.kmsFactory = await ethers.getContractFactory('KMSVerifier');
+    await asyncDecrypt();
+  });
+
+  it('original owner adds one signer, then adds two more signers, then removes one signer', async function () {
+    if (process.env.HARDHAT_PARALLEL !== '1') {
+      // to avoid messing up other tests if used on the real node, in parallel testing
+
+      const origKMSAdd = dotenv.parse(fs.readFileSync('lib/.env.kmsverifier')).KMS_VERIFIER_CONTRACT_ADDRESS;
+      const deployer = (await ethers.getSigners())[9];
+      const kmsVerifier = await this.kmsFactory.attach(origKMSAdd);
+      expect(await kmsVerifier.getVersion()).to.equal('KMSVerifier v0.1.0');
+
+      const privKeySigner = process.env['PRIVATE_KEY_KMS_SIGNER_1']!;
+      const kmsSigner = new ethers.Wallet(privKeySigner).connect(ethers.provider);
+      const tx = await kmsVerifier.connect(deployer).addSigner(kmsSigner.address);
+      await tx.wait();
+
+      expect((await kmsVerifier.getSigners()).length).to.equal(2); // one signer has been added
+
+      const contractFactory = await ethers.getContractFactory('TestAsyncDecrypt');
+      const contract = await contractFactory.connect(this.signers.alice).deploy({
+        value: ethers.parseEther('0.001'),
+      });
+      const tx2 = await contract.requestBool({ gasLimit: 5_000_000 });
+      await tx2.wait();
+      await awaitAllDecryptionResults();
+      const y = await contract.yBool();
+      expect(y).to.equal(true); // in this case, one signature still suffices to pass the decrypt (threshold is still 1)
+
+      const kmsSignerDup = new ethers.Wallet(privKeySigner).connect(ethers.provider);
+      await expect(kmsVerifier.connect(deployer).addSigner(kmsSignerDup.address)).to.revertedWith(
+        'KMSVerifier: Address is already a signer',
+      ); // cannot add duplicated signer
+      expect((await kmsVerifier.getSigners()).length).to.equal(2);
+
+      const privKeySigner2 = process.env['PRIVATE_KEY_KMS_SIGNER_2']!;
+      const kmsSigner2 = new ethers.Wallet(privKeySigner2).connect(ethers.provider);
+      const tx3 = await kmsVerifier.connect(deployer).addSigner(kmsSigner2.address);
+      await tx3.wait();
+      const privKeySigner3 = process.env['PRIVATE_KEY_KMS_SIGNER_3']!;
+      const kmsSigner3 = new ethers.Wallet(privKeySigner3).connect(ethers.provider);
+      const tx4 = await kmsVerifier.connect(deployer).addSigner(kmsSigner3.address);
+      await tx4.wait();
+      expect((await kmsVerifier.getSigners()).length).to.equal(4); // 3rd and 4th signer has been added successfully
+
+      const tx5 = await contract.requestUint4({ gasLimit: 5_000_000 });
+      await tx5.wait();
+      await expect(awaitAllDecryptionResults()).to.revertedWith(
+        'KmsVerifier: at least threshold number of signatures required',
+      ); // should revert because now we are below the threshold! (we receive only 1 signature but threshold is 2)
+      const y2 = await contract.yUint4();
+      expect(y2).to.equal(0);
+
+      process.env.NUM_KMS_SIGNERS = '2';
+      await awaitAllDecryptionResults();
+      const y3 = await contract.yUint4();
+      expect(y3).to.equal(4); // with 2 signatures decryption should now succeed
+
+      process.env.NUM_KMS_SIGNERS = '4';
+      const tx6 = await contract.requestUint8({ gasLimit: 5_000_000 });
+      await tx6.wait();
+      await awaitAllDecryptionResults();
+      const y4 = await contract.yUint8();
+      expect(y4).to.equal(42); // even with more than 2 signatures decryption should still succeed
+
+      process.env.NUM_KMS_SIGNERS = '2';
+      process.env.PRIVATE_KEY_KMS_SIGNER_1 = process.env.PRIVATE_KEY_KMS_SIGNER_0;
+      const tx7 = await contract.requestUint16({ gasLimit: 5_000_000 });
+      await tx7.wait();
+      await expect(awaitAllDecryptionResults()).to.revertedWith('KMS signature verification failed'); // cannot use duplicated signatures if threshold is 2
+      const y5 = await contract.yUint16();
+      expect(y5).to.equal(0);
+
+      process.env.NUM_KMS_SIGNERS = '1';
+      const tx8 = await kmsVerifier.connect(deployer).removeSigner(kmsSigner2.address);
+      await tx8.wait();
+      await awaitAllDecryptionResults();
+      const y6 = await contract.yUint16();
+      expect(y6).to.equal(16); // after removing one of the 4 signers, one signature is enough for decryption
+    }
+  });
+});