apply pr suggestions for #273

src/characters/cbd-recipient.ts

@@ -2,16 +2,17 @@ import {
   AccessControlPolicy,
   AuthenticatedData,
   Ciphertext,
+  CiphertextHeader,
   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';
@@ -49,27 +50,23 @@ export class ThresholdDecrypter {
     conditionExpr: ConditionExpression,
     ciphertext: Ciphertext
   ): Promise<Uint8Array> {
-    const acp = await this.makeAcp(provider, conditionExpr, ciphertext);
+    const acp = await this.makeAcp(provider, conditionExpr, ciphertext.header);
+    const messageKit = new ThresholdMessageKit(ciphertext, acp);
 
     const decryptionShares = await this.retrieve(
       provider,
       conditionExpr,
-      ciphertext,
-      acp
+      messageKit
     );
 
     const sharedSecret = combineDecryptionSharesSimple(decryptionShares);
-    return decryptWithSharedSecret(
-      ciphertext,
-      conditionExpr.asAad(),
-      sharedSecret
-    );
+    return messageKit.decryptWithSharedSecret(sharedSecret);
   }
 
   private async makeAcp(
     provider: ethers.providers.Web3Provider,
     conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext
+    ciphertextHeader: CiphertextHeader
   ) {
     const dkgRitual = await DkgClient.getExistingRitual(
       provider,
@@ -80,7 +77,7 @@ export class ThresholdDecrypter {
       conditionExpr.toWASMConditions()
     );
 
-    const headerHash = keccak256(ciphertext.header.toBytes());
+    const headerHash = keccak256(ciphertextHeader.toBytes());
     const authorization = await provider.getSigner().signMessage(headerHash);
 
     return new AccessControlPolicy(authData, toBytes(authorization));
@@ -90,8 +87,7 @@ export class ThresholdDecrypter {
   public async retrieve(
     provider: ethers.providers.Web3Provider,
     conditionExpr: ConditionExpression,
-    ciphertext: Ciphertext,
-    acp: AccessControlPolicy
+    messageKit: ThresholdMessageKit
   ): Promise<DecryptionShareSimple[]> {
     const dkgParticipants = await DkgCoordinatorAgent.getParticipants(
       provider,
@@ -100,10 +96,10 @@ export class ThresholdDecrypter {
     const contextStr = await conditionExpr.buildContext(provider).toJson();
     const { sharedSecrets, encryptedRequests } = this.makeDecryptionRequests(
       this.ritualId,
-      ciphertext,
+      messageKit.ciphertextHeader,
       contextStr,
       dkgParticipants,
-      acp
+      messageKit.acp
     );
 
     const { encryptedResponses, errors } = await this.porter.cbdDecrypt(
@@ -148,7 +144,7 @@ export class ThresholdDecrypter {
 
   private makeDecryptionRequests(
     ritualId: number,
-    ciphertext: Ciphertext,
+    ciphertextHeader: CiphertextHeader,
     contextStr: string,
     dkgParticipants: Array<DkgParticipant>,
     acp: AccessControlPolicy
@@ -159,7 +155,7 @@ export class ThresholdDecrypter {
     const decryptionRequest = new ThresholdDecryptionRequest(
       ritualId,
       FerveoVariant.simple,
-      ciphertext.header,
+      ciphertextHeader,
       acp,
       new Context(contextStr)
     );

src/characters/enrico.ts

@@ -1,7 +1,8 @@
 import {
+  AuthenticatedData,
   Ciphertext,
   DkgPublicKey,
-  ferveoEncrypt,
+  encryptForDkg,
   MessageKit,
   PublicKey,
   SecretKey,
@@ -52,7 +53,7 @@ export class Enrico {
   public encryptMessageCbd(
     plaintext: Uint8Array | string,
     withConditions?: ConditionExpression
-  ): { ciphertext: Ciphertext; aad: Uint8Array } {
+  ): { ciphertext: Ciphertext; authenticatedData: AuthenticatedData } {
     if (!withConditions) {
       withConditions = this.conditions;
     }
@@ -65,12 +66,11 @@ export class Enrico {
       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,
+      withConditions.toWASMConditions()
     );
-    return { ciphertext, aad };
+    return { ciphertext, authenticatedData };
   }
 }

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
     );

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 { ciphertext, authenticatedData } = deployedStrategy
       .makeEncrypter(conditionExpr)
       .encryptMessageCbd(message);
 
     // Setup mocks for `retrieveAndDecrypt`
-    const { decryptionShares } = fakeTDecFlow({
+    const { decryptionShares } = await fakeTDecFlow({
       ...mockedDkg,
       message: toBytes(message),
-      aad,
       ciphertext,
+      authenticatedData,
+      web3Provider,
     });
-    const { participantSecrets, participants } = fakeDkgParticipants(
-      mockedDkg.ritualId,
-      variant
+    const { participantSecrets, participants } = await fakeDkgParticipants(
+      mockedDkg.ritualId
     );
     const requesterSessionKey = SessionStaticSecret.random();
     const decryptSpy = mockCbdDecrypt(

test/utils.ts

@@ -4,19 +4,20 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 import { Block } from '@ethersproject/providers';
 import {
+  AccessControlPolicy,
   AggregatedTranscript,
+  AuthenticatedData,
   Capsule,
   CapsuleFrag,
   Ciphertext,
   combineDecryptionSharesSimple,
   DecryptionSharePrecomputed,
   DecryptionShareSimple,
-  decryptWithSharedSecret,
   Dkg,
   EncryptedThresholdDecryptionResponse,
   EncryptedTreasureMap,
+  encryptForDkg,
   EthereumAddress,
-  ferveoEncrypt,
   FerveoVariant,
   Keypair,
   PublicKey,
@@ -26,6 +27,7 @@ import {
   SessionStaticKey,
   SessionStaticSecret,
   ThresholdDecryptionResponse,
+  ThresholdMessageKit,
   Transcript,
   Validator,
   ValidatorMessage,
@@ -43,6 +45,8 @@ import {
   DkgRitualState,
 } from '../src/agents/coordinator';
 import { ThresholdDecrypter } from '../src/characters/cbd-recipient';
+import { ConditionExpression } from '../src/conditions';
+import { ERC721Balance } from '../src/conditions/predefined';
 import { DkgClient, DkgRitual } from '../src/dkg';
 import { BlockchainPolicy, PreEnactedPolicy } from '../src/policies/policy';
 import {
@@ -55,6 +59,8 @@ import {
 import { ChecksumAddress } from '../src/types';
 import { toBytes, toHexString, zip } from '../src/utils';
 
+import { TEST_CHAIN_ID, TEST_CONTRACT_ADDR } from './unit/testVariables';
+
 export const bytesEqual = (first: Uint8Array, second: Uint8Array): boolean =>
   first.length === second.length &&
   first.every((value, index) => value === second[index]);
@@ -288,21 +294,23 @@ interface FakeDkgRitualFlow {
   threshold: number;
   receivedMessages: ValidatorMessage[];
   ciphertext: Ciphertext;
-  aad: Uint8Array;
+  authenticatedData: AuthenticatedData;
   dkg: Dkg;
   message: Uint8Array;
+  web3Provider: ethers.providers.Web3Provider;
 }
 
-export const fakeTDecFlow = ({
+export const fakeTDecFlow = async ({
   validators,
   validatorKeypairs,
   ritualId,
   sharesNum,
   threshold,
   receivedMessages,
   ciphertext,
-  aad,
+  authenticatedData,
   message,
+  web3Provider,
 }: FakeDkgRitualFlow) => {
   // Having aggregated the transcripts, the validators can now create decryption shares
   const decryptionShares: (
@@ -320,7 +328,7 @@ export const fakeTDecFlow = ({
     const decryptionShare = aggregate.createDecryptionShareSimple(
       dkg,
       ciphertext.header,
-      aad,
+      authenticatedData.toBytes(),
       keypair
     );
     decryptionShares.push(decryptionShare);
@@ -331,44 +339,66 @@ export const fakeTDecFlow = ({
   const sharedSecret = combineDecryptionSharesSimple(decryptionShares);
 
   // The client should have access to the public parameters of the DKG
-  const plaintext = decryptWithSharedSecret(ciphertext, aad, sharedSecret);
+  const headerHash = keccak256(ciphertext.header.toBytes());
+  const authorization = await web3Provider.getSigner().signMessage(headerHash);
+  const acp = new AccessControlPolicy(
+    authenticatedData,
+    toBytes(authorization)
+  );
+  const messageKit = new ThresholdMessageKit(ciphertext, acp);
+
+  const plaintext = messageKit.decryptWithSharedSecret(sharedSecret);
   if (!bytesEqual(plaintext, message)) {
     throw new Error('Decryption failed');
   }
   return { decryptionShares, sharedSecret, plaintext };
 };
 
-export const fakeDkgTDecFlowE2e = (
-  variant: FerveoVariant,
-  message = toBytes('fake-message'),
-  aad = toBytes('fake-aad'),
+const fakeConditionExpr = () => {
+  const erc721Balance = new ERC721Balance({
+    chain: TEST_CHAIN_ID,
+    contractAddress: TEST_CONTRACT_ADDR,
+  });
+  return new ConditionExpression(erc721Balance);
+};
+
+export const fakeDkgTDecFlowE2e = async (
   ritualId = 0,
+  variant: FerveoVariant = FerveoVariant.precomputed,
+  conditions: ConditionExpression = fakeConditionExpr(),
+  message = toBytes('fake-message'),
+  web3Provider = fakeWeb3Provider(),
   sharesNum = 4,
   threshold = 4
 ) => {
   const ritual = fakeDkgFlow(variant, ritualId, sharesNum, threshold);
 
   // In the meantime, the client creates a ciphertext and decryption request
-  const ciphertext = ferveoEncrypt(message, aad, ritual.dkg.publicKey());
-  const { decryptionShares } = fakeTDecFlow({
+  const [ciphertext, authenticatedData] = encryptForDkg(
+    message,
+    ritual.dkg.publicKey(),
+    conditions.toWASMConditions()
+  );
+  const { decryptionShares } = await fakeTDecFlow({
     ...ritual,
     ciphertext,
-    aad,
+    authenticatedData,
     message,
+    web3Provider,
   });
 
   return {
     ...ritual,
     message,
-    aad,
+    authenticatedData,
     ciphertext,
     decryptionShares,
   };
 };
 
-export const fakeCoordinatorRitual = (
+export const fakeCoordinatorRitual = async (
   ritualId: number
-): {
+): Promise<{
   aggregationMismatch: boolean;
   initTimestamp: number;
   aggregatedTranscriptHash: string;
@@ -380,8 +410,8 @@ export const fakeCoordinatorRitual = (
   aggregatedTranscript: string;
   publicKeyHash: string;
   totalAggregations: number;
-} => {
-  const ritual = fakeDkgTDecFlowE2e(FerveoVariant.precomputed);
+}> => {
+  const ritual = await fakeDkgTDecFlowE2e();
   const dkgPkBytes = ritual.dkg.publicKey().toBytes();
   return {
     id: ritualId,
@@ -401,14 +431,13 @@ export const fakeCoordinatorRitual = (
   };
 };
 
-export const fakeDkgParticipants = (
-  ritualId: number,
-  variant = FerveoVariant.precomputed
-): {
+export const fakeDkgParticipants = async (
+  ritualId: number
+): Promise<{
   participants: DkgParticipant[];
   participantSecrets: Record<string, SessionStaticSecret>;
-} => {
-  const ritual = fakeDkgTDecFlowE2e(variant);
+}> => {
+  const ritual = await fakeDkgTDecFlowE2e(ritualId);
   const label = toBytes(`${ritualId}`);
 
   const participantSecrets: Record<string, SessionStaticSecret> =