apply pr suggestions for #273

nucypher · Sep 3, 2023 · aad5942 · aad5942 · github-actions · Sep 3, 2023
1 parent b7161fd
commit aad5942
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 123 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;
@@ -47,51 +43,22 @@ export class ThresholdDecrypter {
   public async retrieveAndDecrypt(
     provider: ethers.providers.Web3Provider,
     conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext
+    thresholdMessageKit: ThresholdMessageKit
   ): Promise<Uint8Array> {
-    const acp = await this.makeAcp(provider, conditionExpr, ciphertext);
-
     const decryptionShares = await this.retrieve(
       provider,
       conditionExpr,
-      ciphertext,
-      acp
+      thresholdMessageKit
     );
-
     const sharedSecret = combineDecryptionSharesSimple(decryptionShares);
-    return decryptWithSharedSecret(
-      ciphertext,
-      conditionExpr.asAad(),
-      sharedSecret
-    );
-  }
-
-  private async makeAcp(
-    provider: ethers.providers.Web3Provider,
-    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 provider.getSigner().signMessage(headerHash);
-
-    return new AccessControlPolicy(authData, toBytes(authorization));
+    return thresholdMessageKit.decryptWithSharedSecret(sharedSecret);
   }
 
   // Retrieve decryption shares
   public async retrieve(
     provider: ethers.providers.Web3Provider,
     conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext,
-    acp: AccessControlPolicy
+    thresholdMessageKit: ThresholdMessageKit
   ): Promise<DecryptionShareSimple[]> {
     const dkgParticipants = await DkgCoordinatorAgent.getParticipants(
       provider,
@@ -100,10 +67,9 @@ export class ThresholdDecrypter {
     const contextStr = await conditionExpr.buildContext(provider).toJson();
     const { sharedSecrets, encryptedRequests } = this.makeDecryptionRequests(
       this.ritualId,
-      ciphertext,
-      contextStr,
+      new Context(contextStr),
       dkgParticipants,
-      acp
+      thresholdMessageKit
     );
 
     const { encryptedResponses, errors } = await this.porter.cbdDecrypt(
@@ -148,20 +114,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();
@@ -228,8 +193,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 = fakeWeb3Provider(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
@@ -39,6 +39,7 @@ const conditionExpr = new ConditionExpression(ownsNFT);
 const ursulas = fakeUrsulas();
 const variant = FerveoVariant.precomputed;
 const ritualId = 0;
+const web3Provider = fakeWeb3Provider(aliceSecretKey.toBEBytes());
 
 const makeCbdStrategy = async () => {
   const cohort = await makeCohort(ursulas);
@@ -52,7 +53,6 @@ async function makeDeployedCbdStrategy() {
 
   const mockedDkg = fakeDkgFlow(variant, 0, 4, 4);
   const mockedDkgRitual = fakeDkgRitual(mockedDkg);
-  const web3Provider = fakeWeb3Provider(aliceSecretKey.toBEBytes());
   const getUrsulasSpy = mockGetUrsulas(ursulas);
   const getExistingRitualSpy = mockGetExistingRitual(mockedDkgRitual);
   const deployedStrategy = await strategy.deploy(web3Provider, ritualId);
@@ -102,20 +102,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(
@@ -132,7 +132,7 @@ describe('CbdDeployedStrategy', () => {
       await deployedStrategy.decrypter.retrieveAndDecrypt(
         aliceProvider,
         conditionExpr,
-        ciphertext
+        thresholdMessageKit
       );
     expect(getUrsulasSpy).toHaveBeenCalled();
     expect(getParticipantsSpy).toHaveBeenCalled();