apply pr suggestions for #273

nucypher · Sep 5, 2023 · cb258e4 · cb258e4
1 parent 67e43cd
commit cb258e4
Show file tree

Hide file tree

Showing 5 changed files with 112 additions and 116 deletions.
diff --git a/src/characters/cbd-recipient.ts b/src/characters/cbd-recipient.ts
@@ -1,26 +1,22 @@
 import {
-  AccessControlPolicy,
-  AuthenticatedData,
-  Ciphertext,
   combineDecryptionSharesSimple,
   Context,
   DecryptionShareSimple,
-  decryptWithSharedSecret,
   EncryptedThresholdDecryptionRequest,
   EncryptedThresholdDecryptionResponse,
   FerveoVariant,
   SessionSharedSecret,
   SessionStaticSecret,
   ThresholdDecryptionRequest,
+  ThresholdMessageKit,
 } from '@nucypher/nucypher-core';
 import { ethers } from 'ethers';
-import { keccak256 } from 'ethers/lib/utils';
 
 import { DkgCoordinatorAgent, DkgParticipant } from '../agents/coordinator';
 import { ConditionExpression } from '../conditions';
-import { DkgClient, DkgRitual } from '../dkg';
+import { DkgRitual } from '../dkg';
 import { PorterClient } from '../porter';
-import { fromJSON, toBytes, toJSON } from '../utils';
+import { fromJSON, objectEquals, toJSON } from '../utils';
 
 export type ThresholdDecrypterJSON = {
   porterUri: string;
@@ -46,55 +42,25 @@ export class ThresholdDecrypter {
   // Retrieve and decrypt ciphertext using provider and condition expression
   public async retrieveAndDecrypt(
     provider: ethers.providers.Provider,
-    signer: ethers.Signer,
     conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext
+    thresholdMessageKit: ThresholdMessageKit,
+    signer?: ethers.Signer
   ): Promise<Uint8Array> {
-    const acp = await this.makeAcp(provider, signer, conditionExpr, ciphertext);
-
     const decryptionShares = await this.retrieve(
       provider,
       conditionExpr,
-      ciphertext,
-      acp,
+      thresholdMessageKit,
       signer
     );
-
     const sharedSecret = combineDecryptionSharesSimple(decryptionShares);
-    return decryptWithSharedSecret(
-      ciphertext,
-      conditionExpr.asAad(),
-      sharedSecret
-    );
-  }
-
-  private async makeAcp(
-    provider: ethers.providers.Provider,
-    signer: ethers.Signer,
-    conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext
-  ) {
-    const dkgRitual = await DkgClient.getExistingRitual(
-      provider,
-      this.ritualId
-    );
-    const authData = new AuthenticatedData(
-      dkgRitual.dkgPublicKey,
-      conditionExpr.toWASMConditions()
-    );
-
-    const headerHash = keccak256(ciphertext.header.toBytes());
-    const authorization = await signer.signMessage(headerHash);
-
-    return new AccessControlPolicy(authData, toBytes(authorization));
+    return thresholdMessageKit.decryptWithSharedSecret(sharedSecret);
   }
 
   // Retrieve decryption shares
   public async retrieve(
     provider: ethers.providers.Provider,
     conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext,
-    acp: AccessControlPolicy,
+    thresholdMessageKit: ThresholdMessageKit,
     signer?: ethers.Signer
   ): Promise<DecryptionShareSimple[]> {
     const dkgParticipants = await DkgCoordinatorAgent.getParticipants(
@@ -106,10 +72,9 @@ export class ThresholdDecrypter {
       .toJson();
     const { sharedSecrets, encryptedRequests } = this.makeDecryptionRequests(
       this.ritualId,
-      ciphertext,
-      contextStr,
+      new Context(contextStr),
       dkgParticipants,
-      acp
+      thresholdMessageKit
     );
 
     const { encryptedResponses, errors } = await this.porter.cbdDecrypt(
@@ -154,20 +119,19 @@ export class ThresholdDecrypter {
 
   private makeDecryptionRequests(
     ritualId: number,
-    ciphertext: Ciphertext,
-    contextStr: string,
+    conditionContext: Context,
     dkgParticipants: Array<DkgParticipant>,
-    acp: AccessControlPolicy
+    thresholdMessageKit: ThresholdMessageKit
   ): {
     sharedSecrets: Record<string, SessionSharedSecret>;
     encryptedRequests: Record<string, EncryptedThresholdDecryptionRequest>;
   } {
     const decryptionRequest = new ThresholdDecryptionRequest(
       ritualId,
       FerveoVariant.simple,
-      ciphertext.header,
-      acp,
-      new Context(contextStr)
+      thresholdMessageKit.ciphertextHeader,
+      thresholdMessageKit.acp,
+      conditionContext
     );
 
     const ephemeralSessionKey = this.makeSessionKey();
@@ -234,8 +198,6 @@ export class ThresholdDecrypter {
   }
 
   public equals(other: ThresholdDecrypter): boolean {
-    return (
-      this.porter.porterUrl.toString() === other.porter.porterUrl.toString()
-    );
+    return objectEquals(this.toObj(), other.toObj());
   }
 }
diff --git a/src/characters/enrico.ts b/src/characters/enrico.ts
@@ -1,11 +1,13 @@
 import {
-  Ciphertext,
+  AccessControlPolicy,
   DkgPublicKey,
-  ferveoEncrypt,
+  encryptForDkg,
   MessageKit,
   PublicKey,
   SecretKey,
+  ThresholdMessageKit,
 } from '@nucypher/nucypher-core';
+import { arrayify, keccak256 } from 'ethers/lib/utils';
 
 import { ConditionExpression } from '../conditions';
 import { Keyring } from '../keyring';
@@ -51,26 +53,33 @@ export class Enrico {
 
   public encryptMessageCbd(
     plaintext: Uint8Array | string,
-    withConditions?: ConditionExpression
-  ): { ciphertext: Ciphertext; aad: Uint8Array } {
-    if (!withConditions) {
-      withConditions = this.conditions;
+    conditions?: ConditionExpression
+  ): ThresholdMessageKit {
+    if (!conditions) {
+      conditions = this.conditions;
     }
 
-    if (!withConditions) {
+    if (!conditions) {
       throw new Error('Conditions are required for CBD encryption.');
     }
 
     if (!(this.encryptingKey instanceof DkgPublicKey)) {
       throw new Error('Wrong key type. Use encryptMessagePre instead.');
     }
 
-    const aad = withConditions.asAad();
-    const ciphertext = ferveoEncrypt(
+    const [ciphertext, authenticatedData] = encryptForDkg(
       plaintext instanceof Uint8Array ? plaintext : toBytes(plaintext),
-      aad,
-      this.encryptingKey
+      this.encryptingKey,
+      conditions.toWASMConditions()
+    );
+
+    const headerHash = keccak256(ciphertext.header.toBytes());
+    const authorization = this.keyring.signer.sign(arrayify(headerHash));
+    const acp = new AccessControlPolicy(
+      authenticatedData,
+      authorization.toBEBytes()
     );
-    return { ciphertext, aad };
+
+    return new ThresholdMessageKit(ciphertext, acp);
   }
 }
diff --git a/test/integration/dkg-client.test.ts b/test/integration/dkg-client.test.ts
@@ -29,7 +29,7 @@ describe('DkgCoordinatorAgent', () => {
 
   it('fetches participants from the coordinator', async () => {
     const provider = fakeProvider(SecretKey.random().toBEBytes());
-    const fakeParticipants = fakeDkgParticipants(fakeRitualId);
+    const fakeParticipants = await fakeDkgParticipants(fakeRitualId);
     const getParticipantsSpy = mockGetParticipants(
       fakeParticipants.participants
     );

diff --git a/test/unit/cbd-strategy.test.ts b/test/unit/cbd-strategy.test.ts
@@ -30,8 +30,8 @@ const {
 
 // Shared test variables
 const aliceSecretKey = SecretKey.fromBEBytes(aliceSecretKeyBytes);
-const aliceSigner = fakeSigner(aliceSecretKey.toBEBytes());
 const aliceProvider = fakeProvider(aliceSecretKey.toBEBytes());
+const aliceSigner = fakeSigner(aliceSecretKey.toBEBytes());
 const ownsNFT = new ERC721Ownership({
   contractAddress: '0x1e988ba4692e52Bc50b375bcC8585b95c48AaD77',
   parameters: [3591],
@@ -54,10 +54,9 @@ async function makeDeployedCbdStrategy() {
 
   const mockedDkg = fakeDkgFlow(variant, 0, 4, 4);
   const mockedDkgRitual = fakeDkgRitual(mockedDkg);
-  const web3Provider = fakeProvider(aliceSecretKey.toBEBytes());
   const getUrsulasSpy = mockGetUrsulas(ursulas);
   const getExistingRitualSpy = mockGetExistingRitual(mockedDkgRitual);
-  const deployedStrategy = await strategy.deploy(web3Provider, ritualId);
+  const deployedStrategy = await strategy.deploy(aliceProvider, ritualId);
 
   expect(getUrsulasSpy).toHaveBeenCalled();
   expect(getExistingRitualSpy).toHaveBeenCalled();
@@ -104,20 +103,20 @@ describe('CbdDeployedStrategy', () => {
     const { mockedDkg, deployedStrategy } = await makeDeployedCbdStrategy();
 
     const message = 'this is a secret';
-    const { ciphertext, aad } = deployedStrategy
+    const thresholdMessageKit = deployedStrategy
       .makeEncrypter(conditionExpr)
       .encryptMessageCbd(message);
 
     // Setup mocks for `retrieveAndDecrypt`
-    const { decryptionShares } = fakeTDecFlow({
+    const { decryptionShares } = await fakeTDecFlow({
       ...mockedDkg,
       message: toBytes(message),
-      aad,
-      ciphertext,
+      conditionExpr,
+      dkgPublicKey: mockedDkg.dkg.publicKey(),
+      thresholdMessageKit,
     });
-    const { participantSecrets, participants } = fakeDkgParticipants(
-      mockedDkg.ritualId,
-      variant
+    const { participantSecrets, participants } = await fakeDkgParticipants(
+      mockedDkg.ritualId
     );
     const requesterSessionKey = SessionStaticSecret.random();
     const decryptSpy = mockCbdDecrypt(
@@ -133,9 +132,9 @@ describe('CbdDeployedStrategy', () => {
     const decryptedMessage =
       await deployedStrategy.decrypter.retrieveAndDecrypt(
         aliceProvider,
-        aliceSigner,
         conditionExpr,
-        ciphertext
+        thresholdMessageKit,
+        aliceSigner
       );
     expect(getUrsulasSpy).toHaveBeenCalled();
     expect(getParticipantsSpy).toHaveBeenCalled();