feat: construct extraData and add context signer fetching

src/index.ts

@@ -51,6 +51,7 @@ import { RelayerZKProofBuilder } from './relayer-provider/RelayerZKProofBuilder'
 import { CoprocessorSignersVerifier } from './sdk/coprocessor/CoprocessorSignersVerifier';
 import { InputProof } from './sdk/coprocessor/InputProof';
 import { KmsEIP712, TKMSPkeKeypair } from './sdk';
+import { buildRequestExtraData } from './sdk/kms/extraData';
 
 export { getErrorCauseStatus, getErrorCauseCode } from './relayer/error';
 export type { EncryptionBits } from './base/types/primitives';
@@ -97,18 +98,25 @@ export interface FhevmInstance {
     options?: RelayerInputProofOptionsType,
   ): Promise<InputProofBytesType>;
   generateKeypair(): KeypairType<BytesHexNo0x>;
+  /**
+   * Returns the current KMS context extraData for user/delegated user decrypt.
+   * Pass the returned value to both `createEIP712`, `createDelegatedUserDecryptEIP712`, `userDecrypt`, and `delegatedUserDecrypt`.
+   */
+  getExtraData(): Promise<BytesHex>;
   createEIP712(
     publicKey: string,
     contractAddresses: string[],
     startTimestamp: number,
     durationDays: number,
+    extraData: BytesHex,
   ): KmsUserDecryptEIP712Type;
   createDelegatedUserDecryptEIP712(
     publicKey: string,
     contractAddresses: string[],
     delegatorAddress: string,
     startTimestamp: number,
     durationDays: number,
+    extraData: BytesHex,
   ): KmsDelegatedUserDecryptEIP712Type;
   publicDecrypt(
     handles: (string | Uint8Array)[],
@@ -123,6 +131,7 @@ export interface FhevmInstance {
     userAddress: string,
     startTimestamp: number,
     durationDays: number,
+    extraData: BytesHex,
     options?: RelayerUserDecryptOptionsType,
   ): Promise<UserDecryptResults>;
   delegatedUserDecrypt(
@@ -135,6 +144,7 @@ export interface FhevmInstance {
     delegateAddress: string,
     startTimestamp: number,
     durationDays: number,
+    extraData: BytesHex,
     options?: RelayerUserDecryptOptionsType,
   ): Promise<UserDecryptResults>;
   getPublicKey(): { publicKeyId: string; publicKey: Uint8Array } | null;
@@ -176,6 +186,7 @@ export const createInstance = async (
   const thresholdCoprocessorSigners =
     relayerFhevm.fhevmHostChain.coprocessorSignerThreshold;
   const provider = relayerFhevm.fhevmHostChain.ethersProvider;
+  const kmsContextCache = relayerFhevm.fhevmHostChain.kmsContextCache;
 
   return {
     config: relayerFhevm.fhevmHostChain,
@@ -221,11 +232,16 @@ export const createInstance = async (
     generateKeypair: () => {
       return TKMSPkeKeypair.generate().toBytesHexNo0x();
     },
+    getExtraData: async (): Promise<BytesHex> => {
+      const contextId = await kmsContextCache.getCurrentContextId();
+      return buildRequestExtraData(contextId);
+    },
     createEIP712: (
       publicKey: string,
       contractAddresses: string[],
       startTimestamp: number,
       durationDays: number,
+      extraData: BytesHex = '0x00',
     ): KmsUserDecryptEIP712Type => {
       const kmsEIP712 = new KmsEIP712({
         chainId: BigInt(chainId),
@@ -236,7 +252,7 @@ export const createInstance = async (
         contractAddresses,
         startTimestamp,
         durationDays,
-        extraData: '0x00',
+        extraData,
       });
     },
     createDelegatedUserDecryptEIP712: (
@@ -245,6 +261,7 @@ export const createInstance = async (
       delegatorAddress: string,
       startTimestamp: number,
       durationDays: number,
+      extraData: BytesHex = '0x00',
     ): KmsDelegatedUserDecryptEIP712Type => {
       const kmsEIP712 = new KmsEIP712({
         chainId: BigInt(chainId),
@@ -256,7 +273,7 @@ export const createInstance = async (
         delegatorAddress,
         startTimestamp,
         durationDays,
-        extraData: '0x00',
+        extraData,
       });
     },
     publicDecrypt: publicDecryptRequest({
@@ -268,6 +285,7 @@ export const createInstance = async (
       relayerProvider: relayerFhevm.relayerProvider,
       provider,
       defaultOptions,
+      kmsContextCache,
     }),
     userDecrypt: userDecryptRequest({
       kmsSigners,
@@ -278,6 +296,7 @@ export const createInstance = async (
       relayerProvider: relayerFhevm.relayerProvider,
       provider,
       defaultOptions,
+      kmsContextCache,
     }),
     delegatedUserDecrypt: delegatedUserDecryptRequest({
       kmsSigners,
@@ -288,6 +307,7 @@ export const createInstance = async (
       relayerProvider: relayerFhevm.relayerProvider,
       provider,
       defaultOptions,
+      kmsContextCache,
     }),
     getPublicKey: () => {
       const pk = relayerFhevm.getPublicKeyBytes();

src/relayer-provider/AbstractRelayerProvider.ts

@@ -337,9 +337,9 @@ export function assertIsRelayerUserDecryptResult(
     });
   }
   for (let i = 0; i < value.length; ++i) {
-    // Missing extraData
     assertRecordBytesHexNo0xProperty(value[i], 'payload', `${name}[i]`);
     assertRecordBytesHexNo0xProperty(value[i], 'signature', `${name}[i]`);
+    assertRecordBytesHexProperty(value[i], 'extraData', `${name}[i]`);
   }
 }
 

src/relayer-provider/types/public-api.d.ts

@@ -283,7 +283,7 @@ export type RelayerPublicDecryptResult = {
 export type RelayerUserDecryptResult = Array<{
   payload: BytesHexNo0x;
   signature: BytesHexNo0x;
-  //extraData: BytesHex;
+  extraData: BytesHex;
 }>;
 
 export type RelayerInputProofResult = {
@@ -323,7 +323,7 @@ export type RelayerApiError500Type = {
 
 /*
   protocol_paused : that means the user should not retry
-  readiness_check_timed_out: configurable array of interval 
+  readiness_check_timed_out: configurable array of interval
 */
 export type RelayerApiError503Type = {
   label:

src/relayer-provider/v1/RelayerV1Provider.ts

@@ -16,7 +16,6 @@ import {
   assertIsRelayerPublicDecryptResult,
   assertIsRelayerUserDecryptResult,
 } from '../AbstractRelayerProvider';
-
 export class RelayerV1Provider extends AbstractRelayerProvider {
   public override get version(): number {
     return 1;
@@ -86,6 +85,7 @@ export class RelayerV1Provider extends AbstractRelayerProvider {
       payload,
       options,
     );
+
     assertIsRelayerUserDecryptResult(
       json.response,
       'RelayerUserDecryptResult()',

src/relayer-provider/v2/RelayerV2Provider_delegated-user-decrypt.test.ts

@@ -25,6 +25,8 @@ import type {
 } from '../../base/types/primitives';
 import type { RelayerDelegatedUserDecryptPayload } from '../types/public-api';
 import type { FhevmInstanceConfig } from '../../types/relayer';
+import { buildRequestExtraData } from '../../sdk/kms/extraData';
+import { TEST_KMS_CONTEXT_ID } from '../../test/config';
 
 ////////////////////////////////////////////////////////////////////////////////
 //
@@ -259,6 +261,7 @@ describeIfFetchMock('RelayerV2Provider - Delegated User Decrypt', () => {
             {
               payload: 'deadbeef',
               signature: 'deadbeef',
+              extraData: '0x00',
             },
           ],
         },
@@ -269,6 +272,49 @@ describeIfFetchMock('RelayerV2Provider - Delegated User Decrypt', () => {
     const result = await relayerProvider.fetchPostDelegatedUserDecrypt(payload);
     console.log(JSON.stringify(result, null, 2));
   });
+
+  it('v2:delegated-user-decrypt: context-bearing extraData round-trips through provider', async () => {
+    const contextExtraData = buildRequestExtraData(TEST_KMS_CONTEXT_ID);
+    const contextPayload: RelayerDelegatedUserDecryptPayload = {
+      ...payload,
+      extraData: contextExtraData,
+    };
+
+    const bodyJson = {
+      status: 'queued',
+      requestId: 'ctx-req',
+      result: { jobId: '789' },
+    };
+
+    fetchMock.post(DELEGATED_USER_DECRYPT_URL, {
+      status: 202,
+      body: bodyJson,
+      headers: { 'Content-Type': 'application/json' },
+    });
+
+    fetchMock.get(`${DELEGATED_USER_DECRYPT_URL}/789`, {
+      status: 200,
+      body: {
+        status: 'succeeded',
+        requestId: 'ctx-req',
+        result: {
+          result: [
+            {
+              payload: 'deadbeef',
+              signature: 'deadbeef',
+              extraData: contextExtraData,
+            },
+          ],
+        },
+      },
+      headers: { 'Content-Type': 'application/json' },
+    });
+
+    const result =
+      await relayerProvider.fetchPostDelegatedUserDecrypt(contextPayload);
+    expect(result).toHaveLength(1);
+    expect(result[0].extraData).toBe(contextExtraData);
+  });
 });
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -325,6 +371,7 @@ describeIfFetch('FhevmInstance.delegatedUserDecrypt:sepolia:', () => {
       delegatorAddress,
       startTimestamp,
       durationDays,
+      extraData,
     );
 
     const kmsEIP712 = new KmsEIP712({
@@ -365,6 +412,7 @@ describeIfFetch('FhevmInstance.delegatedUserDecrypt:sepolia:', () => {
       delegatorAddress,
       startTimestamp,
       durationDays,
+      extraData,
     );
 
     expect(eCount in res).toBe(true);
@@ -402,7 +450,10 @@ defaultDescribe('FhevmInstance.createDelegatedUserDecryptEIP712', () => {
     consoleLogSpy.mockRestore();
   });
 
-  async function testDelegateCreateEIP712(config: FhevmInstanceConfig) {
+  async function testDelegateCreateEIP712(
+    config: FhevmInstanceConfig,
+    extraData: BytesHex = '0x00',
+  ) {
     const mnemonic = TEST_CONFIG.mnemonic;
     if (typeof mnemonic !== 'string' || mnemonic.length === 0) {
       throw new Error(`Missing MNEMONIC env variable in .env file!`);
@@ -419,7 +470,6 @@ defaultDescribe('FhevmInstance.createDelegatedUserDecryptEIP712', () => {
 
     const startTimestamp = timestampNow();
     const durationDays = 365;
-    const extraData = '0x00';
     const contractAddress = TEST_CONFIG.testContracts.FHETestAddress;
     const contractAddresses = [contractAddress];
     const delegatorAddress = '0xDeaDbeefdEAdbeefdEadbEEFdeadbeEFdEaDbeeF';
@@ -446,6 +496,7 @@ defaultDescribe('FhevmInstance.createDelegatedUserDecryptEIP712', () => {
       delegatorAddress,
       startTimestamp,
       durationDays,
+      extraData,
     );
 
     const kmsEIP712 = new KmsEIP712({
@@ -511,4 +562,14 @@ defaultDescribe('FhevmInstance.createDelegatedUserDecryptEIP712', () => {
     });
     await testDelegateCreateEIP712(TEST_CONFIG.v2.fhevmInstanceConfig);
   }, 600000);
+
+  it('v2: createDelegatedUserDecryptEIP712 with context-bearing extraData', async () => {
+    setupAllFetchMockRoutes({
+      enableInputProofRoutes: false,
+    });
+    await testDelegateCreateEIP712(
+      TEST_CONFIG.v2.fhevmInstanceConfig,
+      buildRequestExtraData(TEST_KMS_CONTEXT_ID),
+    );
+  }, 600000);
 });