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,35 @@ 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()` and `userDecrypt(options.extraData)`.
+   */
+  getExtraData(): Promise<BytesHex>;
+  /**
+   * NOTE: Calling without extraData uses legacy '0x00'. Use:
+   *   const extraData = await instance.getExtraData();
+   *   const eip712 = instance.createEIP712(pubKey, contracts, start, days, extraData);
+   */
   createEIP712(
     publicKey: string,
     contractAddresses: string[],
     startTimestamp: number,
     durationDays: number,
+    extraData?: BytesHex,
   ): KmsUserDecryptEIP712Type;
+  /**
+   * NOTE: Calling without extraData uses legacy '0x00'. Use:
+   *   const extraData = await instance.getExtraData();
+   *   const eip712 = instance.createDelegatedUserDecryptEIP712(pubKey, contracts, delegator, start, days, extraData);
+   */
   createDelegatedUserDecryptEIP712(
     publicKey: string,
     contractAddresses: string[],
     delegatorAddress: string,
     startTimestamp: number,
     durationDays: number,
+    extraData?: BytesHex,
   ): KmsDelegatedUserDecryptEIP712Type;
   publicDecrypt(
     handles: (string | Uint8Array)[],
@@ -176,6 +194,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 +240,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 +260,7 @@ export const createInstance = async (
         contractAddresses,
         startTimestamp,
         durationDays,
-        extraData: '0x00',
+        extraData,
       });
     },
     createDelegatedUserDecryptEIP712: (
@@ -245,6 +269,7 @@ export const createInstance = async (
       delegatorAddress: string,
       startTimestamp: number,
       durationDays: number,
+      extraData: BytesHex = '0x00',
     ): KmsDelegatedUserDecryptEIP712Type => {
       const kmsEIP712 = new KmsEIP712({
         chainId: BigInt(chainId),
@@ -256,7 +281,7 @@ export const createInstance = async (
         delegatorAddress,
         startTimestamp,
         durationDays,
-        extraData: '0x00',
+        extraData,
       });
     },
     publicDecrypt: publicDecryptRequest({
@@ -268,6 +293,7 @@ export const createInstance = async (
       relayerProvider: relayerFhevm.relayerProvider,
       provider,
       defaultOptions,
+      kmsContextCache,
     }),
     userDecrypt: userDecryptRequest({
       kmsSigners,
@@ -278,6 +304,7 @@ export const createInstance = async (
       relayerProvider: relayerFhevm.relayerProvider,
       provider,
       defaultOptions,
+      kmsContextCache,
     }),
     delegatedUserDecrypt: delegatedUserDecryptRequest({
       kmsSigners,
@@ -288,6 +315,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

@@ -66,6 +66,8 @@ export type RelayerUserDecryptOptionsType = Prettify<
     signal?: AbortSignal;
     timeout?: number;
     onProgress?: (args: RelayerUserDecryptProgressArgs) => void;
+    /** Context-aware extraData from getExtraData(). Optional — defaults to legacy '0x00' when omitted. */
+    extraData?: BytesHex;
   }
 >;
 
@@ -283,7 +285,7 @@ export type RelayerPublicDecryptResult = {
 export type RelayerUserDecryptResult = Array<{
   payload: BytesHexNo0x;
   signature: BytesHexNo0x;
-  //extraData: BytesHex;
+  extraData: BytesHex;
 }>;
 
 export type RelayerInputProofResult = {

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

@@ -259,6 +259,7 @@ describeIfFetchMock('RelayerV2Provider - Delegated User Decrypt', () => {
             {
               payload: 'deadbeef',
               signature: 'deadbeef',
+              extraData: '0x00',
             },
           ],
         },

src/relayer-provider/v2/RelayerV2Provider_user-decrypt.test.ts

@@ -247,7 +247,7 @@ describeIfFetchMock('RelayerV2Provider', () => {
             {
               payload: 'deadbeef',
               signature: 'deadbeef',
-              //extraData: '0x00',
+              extraData: '0x00',
             },
           ],
         },

src/relayer-provider/v2/guards/RelayerV2ResultUserDecrypt.test.ts

@@ -3,20 +3,19 @@ import { assertIsRelayerV2ResultUserDecrypt } from './RelayerV2ResultUserDecrypt
 
 // Jest Command line
 // =================
-// npx jest --colors --passWithNoTests ./src/relayer-provider/v2/types/RelayerV2ResultUserDecrypt.test.ts
-// npx jest --colors --passWithNoTests --coverage ./src/relayer-provider/v2/types/RelayerV2ResultUserDecrypt.test.ts --collectCoverageFrom=./src/relayer-provider/v2/types/RelayerV2ResultUserDecrypt.ts
+// npx jest --colors --passWithNoTests ./src/relayer-provider/v2/guards/RelayerV2ResultUserDecrypt.test.ts
 
 describe('RelayerV2ResultUserDecrypt', () => {
   it('assertIsRelayerV2ResultUserDecrypt', () => {
-    // True
+    // True — valid response with extraData
     expect(() =>
       assertIsRelayerV2ResultUserDecrypt(
         {
           result: [
             {
               payload: 'deadbeef',
               signature: 'deadbeef',
-              //extraData: '0x00',
+              extraData: '0x00',
             },
           ],
         },
@@ -110,5 +109,44 @@ describe('RelayerV2ResultUserDecrypt', () => {
         type: 'string',
       }),
     );
+
+    // Missing extraData on item throws
+    expect(() =>
+      assertIsRelayerV2ResultUserDecrypt(
+        {
+          result: [{ payload: 'deadbeef', signature: 'deadbeef' }],
+        },
+        'Foo',
+      ),
+    ).toThrow(
+      InvalidPropertyError.missingProperty({
+        objName: 'Foo.result[0]',
+        property: 'extraData',
+        expectedType: 'BytesHex',
+      }),
+    );
+
+    // Invalid (non-hex) extraData throws
+    expect(() =>
+      assertIsRelayerV2ResultUserDecrypt(
+        {
+          result: [
+            {
+              payload: 'deadbeef',
+              signature: 'deadbeef',
+              extraData: 'not-hex',
+            },
+          ],
+        },
+        'Foo',
+      ),
+    ).toThrow(
+      new InvalidPropertyError({
+        objName: 'Foo.result[0]',
+        property: 'extraData',
+        expectedType: 'BytesHex',
+        type: 'string',
+      }),
+    );
   });
 });

src/relayer-provider/v2/guards/RelayerV2ResultUserDecrypt.ts

@@ -1,5 +1,8 @@
 import type { RelayerV2ResultUserDecrypt } from '../types';
-import { assertRecordBytesHexNo0xProperty } from '@base/bytes';
+import {
+  assertRecordBytesHexNo0xProperty,
+  assertRecordBytesHexProperty,
+} from '@base/bytes';
 import { assertRecordArrayProperty } from '@base/record';
 
 /**
@@ -21,7 +24,6 @@ export function assertIsRelayerV2ResultUserDecrypt(
 
   assertRecordArrayProperty(value, 'result' satisfies keyof T, name);
   for (let i = 0; i < value.result.length; ++i) {
-    // Missing extraData
     assertRecordBytesHexNo0xProperty(
       value.result[i],
       'payload' satisfies keyof ResultItem,
@@ -32,5 +34,10 @@ export function assertIsRelayerV2ResultUserDecrypt(
       'signature' satisfies keyof ResultItem,
       `${name}.result[${i}]`,
     );
+    assertRecordBytesHexProperty(
+      value.result[i],
+      'extraData' satisfies keyof ResultItem,
+      `${name}.result[${i}]`,
+    );
   }
 }

src/relayer/publicDecrypt.test.ts

@@ -1,4 +1,5 @@
 import type { RelayerPublicDecryptPayload } from '../relayer-provider/types/public-api';
+import type { KmsContextCache } from '../sdk/kms/KmsContextCache';
 import { publicDecryptRequest } from './publicDecrypt';
 import fetchMock from 'fetch-mock';
 import { ethers } from 'ethers';
@@ -8,6 +9,12 @@ import { TEST_CONFIG, TEST_KMS } from '../test/config';
 import { fetchRelayerV1Post } from '../relayer-provider/v1/fetchRelayerV1';
 import { ChecksummedAddress } from 'src/node';
 
+// Mock KmsContextCache that returns init-time signers (legacy behavior)
+const mockKmsContextCache = {
+  getCurrentContextId: jest.fn().mockResolvedValue(0n),
+  getSignersForContext: jest.fn().mockResolvedValue(TEST_KMS.addresses),
+} as unknown as KmsContextCache;
+
 // Jest Command line
 // =================
 // npx jest --colors --passWithNoTests ./src/relayer/publicDecrypt.test.ts
@@ -58,6 +65,7 @@ describeIfFetchMock('publicDecrypt', () => {
         .aclContractAddress as ChecksummedAddress,
       relayerProvider,
       provider: new ethers.JsonRpcProvider('https://devnet.zama.ai'),
+      kmsContextCache: mockKmsContextCache,
     });
   });
 });

src/relayer/publicDecrypt.ts

@@ -23,6 +23,12 @@ import { solidityPrimitiveTypeNameFromFheTypeId } from '@sdk/FheType';
 import { FhevmHandle } from '@sdk/FhevmHandle';
 import { fhevmHandleCheck2048EncryptedBits } from './decryptUtils';
 import { ACL } from '@sdk/ACL';
+import type { KmsContextCache } from '@sdk/kms/KmsContextCache';
+import {
+  isLegacyExtraData,
+  parseExtraData,
+  buildRequestExtraData,
+} from '@sdk/kms/extraData';
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -208,6 +214,7 @@ export const publicDecryptRequest =
     aclContractAddress,
     relayerProvider,
     provider,
+    kmsContextCache,
     defaultOptions,
   }: {
     kmsSigners: ChecksummedAddress[];
@@ -217,13 +224,16 @@ export const publicDecryptRequest =
     aclContractAddress: ChecksummedAddress;
     relayerProvider: AbstractRelayerProvider;
     provider: EthersProviderType;
+    kmsContextCache: KmsContextCache;
     defaultOptions?: FhevmInstanceOptions;
   }) =>
   async (
     _handles: (Uint8Array | string)[],
     options?: RelayerPublicDecryptOptionsType,
   ): Promise<PublicDecryptResults> => {
-    const extraData: `0x${string}` = '0x00';
+    // Request side: build dynamic extraData from current context ID
+    const currentContextId = await kmsContextCache.getCurrentContextId();
+    const extraData = buildRequestExtraData(currentContextId);
 
     const orderedFhevmHandles: FhevmHandle[] = _handles.map(FhevmHandle.from);
     const orderedHandlesBytes32Hex: Bytes32Hex[] = orderedFhevmHandles.map(
@@ -256,12 +266,10 @@ export const publicDecryptRequest =
     const decryptedResult: `0x${string}` = ensure0x(json.decryptedValue);
     const kmsSignatures: `0x${string}`[] = json.signatures.map(ensure0x);
 
-    ////////////////////////////////////////////////////////////////////////////
-    //
-    // Warning!!!! Do not use '0x00' here!! Only '0x' is permitted!
-    //
-    ////////////////////////////////////////////////////////////////////////////
-    const signedExtraData = '0x';
+    // Always use the raw response extraData for EIP-712 signature verification.
+    // The KMS signs whatever extraData bytes it receives — the SDK must verify
+    // against the same bytes, whether legacy or context-bearing.
+    const signedExtraData: `0x${string}` = json.extraData;
 
     ////////////////////////////////////////////////////////////////////////////
     // Compute the PublicDecryptionProof
@@ -287,6 +295,20 @@ export const publicDecryptRequest =
     ////////////////////////////////////////////////////////////////////////////
 
     // verify signatures on decryption:
+
+    // Response side: resolve signers based on response extraData
+    let effectiveSigners: string[];
+    if (isLegacyExtraData(signedExtraData)) {
+      // Legacy path: use init-time signers
+      effectiveSigners = [...kmsSigners];
+    } else {
+      // Context path: parse contextId, fetch context-specific signers
+      // Fail closed: RPC errors propagate — no silent fallback to init-time signers
+      const { contextId } = parseExtraData(signedExtraData);
+      effectiveSigners = await kmsContextCache.getSignersForContext(contextId);
+    }
+
+    // Verify signatures on decryption
     const domain = {
       name: 'Decryption',
       version: '1',
@@ -318,7 +340,7 @@ export const publicDecryptRequest =
     );
 
     const thresholdReached = isThresholdReached(
-      kmsSigners,
+      effectiveSigners,
       recoveredAddresses,
       thresholdSigners,
     );