@fhevm/sdk 1.1.0-alpha.4 → 1.1.0-alpha.5

package/core/kms/TransportKeyPair-p.ts

@@ -2,6 +2,9 @@ import type { ErrorMetadataParams } from '../base/errors/ErrorBase.js';
 import type { FhevmRuntime, WithDecrypt } from '../types/coreFhevmRuntime.js';
 import type { Bytes, BytesHex } from '../types/primitives.js';
 import type { TkmsPrivateKey } from '../types/tkms-p.js';
+import type { TkmsVersion } from '../../wasm/tkms/KmsLibApi.js';
+import type { NativeClient, OptionalNativeClient } from '../types/coreFhevmClient.js';
+import type { FhevmChain } from '../types/fhevmChain.js';
 import { assertIsBytesOrBytesHex, bytesToHexLarge, hexToBytesFaster } from '../base/bytes.js';
 import { InvalidTypeError } from '../base/errors/InvalidTypeError.js';
 import { assertRecordNonNullableProperty } from '../base/record.js';
@@ -18,6 +21,7 @@ declare const TransportKeyPairBrand: unique symbol;
 
 export type TransportKeyPair = {
   readonly publicKey: BytesHex;
+  readonly tkmsVersion: TkmsVersion;
   readonly [TransportKeyPairBrand]: true;
 };
 
@@ -40,6 +44,7 @@ class TransportKeyPairImpl implements TransportKeyPair {
 
   readonly #publicKeyBytesHex: BytesHex;
   readonly #privateKeyBytes: Bytes;
+  readonly #tkmsVersion: TkmsVersion;
 
   readonly #runtime: WeakRef<FhevmRuntime>;
   #tkmsPrivateKey: TkmsPrivateKey | undefined;
@@ -51,6 +56,7 @@ class TransportKeyPairImpl implements TransportKeyPair {
       readonly publicKeyBytesHex: BytesHex;
       readonly privateKeyBytes: Bytes;
       readonly tkmsPrivateKey?: TkmsPrivateKey | undefined;
+      readonly tkmsVersion: TkmsVersion;
     },
   ) {
     if (privateToken !== PRIVATE_TOKEN) {
@@ -60,12 +66,17 @@ class TransportKeyPairImpl implements TransportKeyPair {
     this.#publicKeyBytesHex = parameters.publicKeyBytesHex;
     this.#privateKeyBytes = parameters.privateKeyBytes;
     this.#tkmsPrivateKey = parameters.tkmsPrivateKey;
+    this.#tkmsVersion = parameters.tkmsVersion;
   }
 
   public get publicKey(): BytesHex {
     return this.#publicKeyBytesHex;
   }
 
+  public get tkmsVersion(): TkmsVersion {
+    return this.#tkmsVersion;
+  }
+
   /**
    * Hidden method (Symbol key) that returns the deserialized TKMS private key.
    * Lazily deserializes and verifies on first call.
@@ -94,11 +105,12 @@ class TransportKeyPairImpl implements TransportKeyPair {
 
     this.#tkmsPrivateKey = await runtimeWithDecrypt.decrypt.deserializeTkmsPrivateKey({
       tkmsPrivateKeyBytes: this.#privateKeyBytes,
+      tkmsVersion: this.#tkmsVersion,
     });
 
     // Verify the key is valid
     await verifyTkmsPublicKey(
-      { runtime: runtimeWithDecrypt },
+      { runtime: runtimeWithDecrypt, tkmsVersion: this.#tkmsVersion },
       {
         tkmsPrivateKey: this.#tkmsPrivateKey,
         tkmsPublicKeyBytesHex: this.#publicKeyBytesHex,
@@ -114,8 +126,9 @@ class TransportKeyPairImpl implements TransportKeyPair {
    * Access is protected by the Symbol key (not exported) and the private token.
    */
   [SerializeFn](privateToken: symbol): {
-    publicKey: BytesHex;
-    privateKey: BytesHex;
+    readonly publicKey: BytesHex;
+    readonly privateKey: BytesHex;
+    readonly tkmsVersion: TkmsVersion;
   } {
     if (privateToken !== PRIVATE_TOKEN) {
       throw new Error('Unauthorized');
@@ -123,6 +136,7 @@ class TransportKeyPairImpl implements TransportKeyPair {
     return {
       publicKey: this.#publicKeyBytesHex,
       privateKey: bytesToHexLarge(this.#privateKeyBytes, false /* no0x */),
+      tkmsVersion: this.#tkmsVersion,
     };
   }
 
@@ -130,8 +144,8 @@ class TransportKeyPairImpl implements TransportKeyPair {
    * Prevents accidental private key exposure via `JSON.stringify`.
    * Only the public key is included in the output.
    */
-  public toJSON(): { publicKey: string } {
-    return { publicKey: this.#publicKeyBytesHex };
+  public toJSON(): { publicKey: string; tkmsVersion: string } {
+    return { publicKey: this.#publicKeyBytesHex, tkmsVersion: this.#tkmsVersion };
   }
 }
 
@@ -167,16 +181,25 @@ export function assertIsTransportKeyPair(
 ////////////////////////////////////////////////////////////////////////////////
 
 /** Generates a fresh {@link TransportKeyPair}. */
-export async function generateTransportKeyPair(context: { readonly runtime: WithDecrypt }): Promise<TransportKeyPair> {
-  const tkmsPrivateKey = await context.runtime.decrypt.generateTkmsPrivateKey();
-  const tkmsPrivateKeyBytes = await context.runtime.decrypt.serializeTkmsPrivateKey({ tkmsPrivateKey });
+export async function generateTransportKeyPair(context: {
+  readonly runtime: WithDecrypt;
+  readonly chain: FhevmChain;
+  readonly client: NativeClient;
+  readonly tkmsVersion: TkmsVersion;
+}): Promise<TransportKeyPair> {
+  const tkmsVersion = context.tkmsVersion;
+
+  const tkmsPrivateKey = await context.runtime.decrypt.generateTkmsPrivateKey({ tkmsVersion });
+  const tkmsPrivateKeyBytes = await context.runtime.decrypt.serializeTkmsPrivateKey({ tkmsPrivateKey, tkmsVersion });
   const tkmsPublicKeyBytesHex = await context.runtime.decrypt.getTkmsPublicKeyHex({
     tkmsPrivateKey,
+    tkmsVersion,
   });
   return new TransportKeyPairImpl(PRIVATE_TOKEN, context.runtime, {
     privateKeyBytes: tkmsPrivateKeyBytes,
     publicKeyBytesHex: tkmsPublicKeyBytesHex,
     tkmsPrivateKey,
+    tkmsVersion,
   });
 }
 
@@ -190,7 +213,11 @@ export async function generateTransportKeyPair(context: { readonly runtime: With
  * @throws {InvalidTypeError} If `value` is not a recognized key pair shape.
  */
 export async function toTransportKeyPair(
-  context: { readonly runtime: FhevmRuntime },
+  context: {
+    readonly runtime: FhevmRuntime;
+    readonly chain: FhevmChain;
+    readonly client?: OptionalNativeClient;
+  },
   value: unknown,
 ): Promise<TransportKeyPair> {
   if (isTransportKeyPair(value)) {
@@ -210,9 +237,20 @@ export async function toTransportKeyPair(
     expectedType: 'Bytes | BytesHex',
     ...options,
   });
+  assertRecordNonNullableProperty(value, 'tkmsVersion', name, {
+    expectedType: 'string',
+    ...options,
+  });
 
   const rawPublicKey = value.publicKey;
   const rawPrivateKey = value.privateKey;
+  const rawTkmsVersion = value.tkmsVersion;
+
+  // TODO: this is not correct. We must add a resolution function
+  let resolvedTkmsVersion: TkmsVersion = '0.13.20-0';
+  if (rawTkmsVersion === '0.13.10') {
+    resolvedTkmsVersion = rawTkmsVersion;
+  }
 
   assertIsBytesOrBytesHex(rawPublicKey, { subject: `${name}.publicKey` });
   assertIsBytesOrBytesHex(rawPrivateKey, { subject: `${name}.privateKey` });
@@ -236,14 +274,22 @@ export async function toTransportKeyPair(
   if (runtimeWithDecrypt !== undefined) {
     tkmsPrivateKey = await runtimeWithDecrypt.decrypt.deserializeTkmsPrivateKey({
       tkmsPrivateKeyBytes,
+      tkmsVersion: resolvedTkmsVersion,
     });
-    await verifyTkmsPublicKey({ runtime: runtimeWithDecrypt }, { tkmsPrivateKey, tkmsPublicKeyBytesHex });
+    await verifyTkmsPublicKey(
+      { runtime: runtimeWithDecrypt, tkmsVersion: resolvedTkmsVersion },
+      {
+        tkmsPrivateKey,
+        tkmsPublicKeyBytesHex,
+      },
+    );
   }
 
   return new TransportKeyPairImpl(PRIVATE_TOKEN, context.runtime, {
     privateKeyBytes: tkmsPrivateKeyBytes,
     publicKeyBytesHex: tkmsPublicKeyBytesHex,
     tkmsPrivateKey,
+    tkmsVersion: resolvedTkmsVersion,
   });
 }
 

package/core/kms/createKmsUserDecryptEip712.ts

@@ -15,6 +15,7 @@ import { assertRecordStringProperty, ensure0x } from '../base/string.js';
 import { assertIsUint64, assertIsUintNumber } from '../base/uint.js';
 import { assertIsKmsExtraData } from './kmsExtraData.js';
 import { kmsUserDecryptEip712Types } from './kmsUserDecryptEip712Types.js';
+import { isDeepEqual } from '../base/object.js';
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -120,6 +121,9 @@ export function _assertIsKmsUserDecryptEip712Base(
   assertIsKmsEip712Domain((value as Record<string, unknown>).domain, `${name}.domain`, options);
 
   assertRecordNonNullableProperty(value, 'types', name, options);
+  if (!isDeepEqual(value.types, kmsUserDecryptEip712Types)) {
+    throw new Error('Unexpected KmsUserDecryptEip712Types');
+  }
 
   assertRecordStringProperty(value, 'primaryType', name, {
     expectedValue: primaryType,

package/core/kms/createKmsUserDecryptEip712V2.ts

@@ -0,0 +1,159 @@
+import type { ErrorMetadataParams } from '../base/errors/ErrorBase.js';
+import type { KmsUserDecryptEip712V2 } from '../types/kms.js';
+import type { BytesHex } from '../types/primitives.js';
+import type { KmsUserDecryptEip712V2Message } from '../types/kms.js';
+import {
+  addressToChecksummedAddress,
+  assertIsAddress,
+  assertIsAddressArray,
+  assertRecordAddressProperty,
+  assertRecordChecksummedAddressArrayProperty,
+} from '../base/address.js';
+import { assertIsKmsEip712Domain, createKmsEip712Domain } from './createKmsEip712Domain.js';
+import { asBytesHex, assertIsBytesHex, assertRecordBytesHexProperty, bytesToHexLarge } from '../base/bytes.js';
+import { isDeepEqual } from '../base/object.js';
+import { assertRecordNonNullableProperty } from '../base/record.js';
+import { assertRecordStringProperty, ensure0x } from '../base/string.js';
+import { assertIsUint256, assertIsUint64, assertIsUintString, MAX_UINT256 } from '../base/uint.js';
+import { assertIsKmsExtraData } from './kmsExtraData.js';
+import { kmsUserDecryptEip712V2Types } from './kmsUserDecryptEip712V2Types.js';
+
+////////////////////////////////////////////////////////////////////////////////
+
+export type CreateKmsUserDecryptEip712V2Parameters = {
+  readonly verifyingContractAddressDecryption: string;
+  readonly chainId: number | bigint;
+  readonly userAddress: string;
+  readonly publicKey: string | Uint8Array;
+  readonly allowedContracts: readonly string[];
+  readonly startTimestamp: number | bigint;
+  readonly durationSeconds: number | bigint;
+  readonly extraData: string;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// createKmsUserDecryptEip712V2 (RFC-016)
+////////////////////////////////////////////////////////////////////////////////
+
+export function createKmsUserDecryptEip712V2(
+  parameters: CreateKmsUserDecryptEip712V2Parameters,
+): KmsUserDecryptEip712V2 {
+  const {
+    verifyingContractAddressDecryption,
+    chainId,
+    userAddress,
+    publicKey,
+    allowedContracts,
+    startTimestamp,
+    durationSeconds,
+    extraData,
+  } = parameters;
+  const publicKeyBytesHex = _verifyPublicKeyArg(publicKey);
+
+  assertIsUint64(chainId, {});
+  assertIsAddress(verifyingContractAddressDecryption, {});
+  assertIsAddress(userAddress, {});
+  assertIsAddressArray(allowedContracts, {});
+  assertIsUint256(startTimestamp, {});
+  assertIsUint256(durationSeconds, {});
+  assertIsBytesHex(extraData, {});
+  assertIsKmsExtraData(extraData, {});
+
+  const checksummedUserAddress = addressToChecksummedAddress(userAddress);
+  const checksummedContractAddresses = allowedContracts.map(addressToChecksummedAddress);
+
+  const primaryType: KmsUserDecryptEip712V2['primaryType'] = 'UserDecryptRequestVerification';
+
+  const domain = createKmsEip712Domain({
+    chainId,
+    verifyingContractAddressDecryption,
+  });
+
+  const eip712 = {
+    domain,
+    types: kmsUserDecryptEip712V2Types,
+    primaryType,
+    message: {
+      userAddress: checksummedUserAddress,
+      publicKey: publicKeyBytesHex,
+      allowedContracts: checksummedContractAddresses,
+      startTimestamp: startTimestamp.toString(),
+      durationSeconds: durationSeconds.toString(),
+      extraData,
+    } satisfies KmsUserDecryptEip712V2Message,
+  };
+
+  Object.freeze(eip712);
+  Object.freeze(eip712.domain);
+  Object.freeze(eip712.types);
+  Object.freeze(eip712.types.EIP712Domain);
+  Object.freeze(eip712.types.UserDecryptRequestVerification);
+  Object.freeze(eip712.message);
+  Object.freeze(eip712.message.allowedContracts);
+
+  return eip712;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+export function assertIsKmsUserDecryptEip712V2(
+  value: unknown,
+  name: string,
+  options: ErrorMetadataParams,
+): asserts value is KmsUserDecryptEip712V2 {
+  assertRecordNonNullableProperty(value, 'domain', name, options);
+  assertIsKmsEip712Domain((value as Record<string, unknown>).domain, `${name}.domain`, options);
+
+  assertRecordNonNullableProperty(value, 'types', name, options);
+  if (!isDeepEqual(value.types, kmsUserDecryptEip712V2Types)) {
+    throw new Error('Unexpected KmsUserDecryptEip712V2Types');
+  }
+
+  assertRecordStringProperty(value, 'primaryType', name, {
+    expectedValue: 'UserDecryptRequestVerification',
+    ...options,
+  });
+
+  assertRecordNonNullableProperty(value, 'message', name, options);
+  const msg = (value as Record<string, unknown>).message as Record<string, unknown>;
+  _assertIsKmsUserDecryptEip712V2Message(msg, `${name}.message`, options);
+}
+
+function _assertIsKmsUserDecryptEip712V2Message(
+  msg: unknown,
+  msgName: string,
+  options: ErrorMetadataParams,
+): asserts msg is KmsUserDecryptEip712V2Message {
+  type MessageType = KmsUserDecryptEip712V2['message'];
+  assertRecordAddressProperty(msg, 'userAddress' satisfies keyof MessageType, msgName, options);
+  assertRecordBytesHexProperty(msg, 'publicKey' satisfies keyof MessageType, msgName, options);
+  assertRecordChecksummedAddressArrayProperty(msg, 'allowedContracts' satisfies keyof MessageType, msgName, options);
+  assertRecordStringProperty(msg, 'startTimestamp' satisfies keyof MessageType, msgName, options);
+  assertRecordStringProperty(msg, 'durationSeconds' satisfies keyof MessageType, msgName, options);
+  assertRecordBytesHexProperty(msg, 'extraData' satisfies keyof MessageType, msgName, options);
+  assertIsKmsExtraData(msg.extraData, options);
+  assertIsUintString(msg.startTimestamp, { max: MAX_UINT256 });
+  assertIsUintString(msg.durationSeconds, { max: MAX_UINT256 });
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+function _verifyPublicKeyArg(value: unknown): BytesHex {
+  if (value === null || value === undefined) {
+    throw new Error(`Missing publicKey argument.`);
+  }
+
+  let publicKeyBytesHex: BytesHex;
+
+  const pk = value;
+
+  if (typeof pk === 'string') {
+    publicKeyBytesHex = asBytesHex(ensure0x(pk));
+  } else if (pk instanceof Uint8Array) {
+    publicKeyBytesHex = bytesToHexLarge(pk);
+  } else {
+    throw new Error(`Invalid publicKey argument.`);
+  }
+
+  return publicKeyBytesHex;
+}

package/core/kms/{decryptKmsSignedcryptedShares-p.ts → decryptKmsSigncryptedShares-p.ts}

@@ -2,14 +2,16 @@ import type { WithDecrypt } from '../types/coreFhevmRuntime.js';
 import type { KmsSigncryptedShares } from '../types/kms.js';
 import type { ClearValue } from '../types/encryptedTypes-p.js';
 import type { TransportKeyPair } from './TransportKeyPair-p.js';
+import type { TkmsVersion } from '../../wasm/tkms/KmsLibApi.js';
 import { transportKeyPairToTkmsPrivateKey } from './TransportKeyPair-p.js';
 
 ////////////////////////////////////////////////////////////////////////////////
-// decryptKmsSignedcryptedShares (with privateKey)
+// decryptKmsSigncryptedShares (with privateKey)
 ////////////////////////////////////////////////////////////////////////////////
 
 type Context = {
   readonly runtime: WithDecrypt;
+  readonly tkmsVersion: TkmsVersion;
 };
 
 type Parameters = {
@@ -21,9 +23,16 @@ type ReturnType = readonly ClearValue[];
 
 ////////////////////////////////////////////////////////////////////////////////
 
-export async function decryptKmsSignedcryptedShares(context: Context, parameters: Parameters): Promise<ReturnType> {
+export async function decryptKmsSigncryptedShares(context: Context, parameters: Parameters): Promise<ReturnType> {
   const { transportKeyPair: transportKeyPair, kmsSigncryptedShares } = parameters;
 
+  if (context.tkmsVersion !== kmsSigncryptedShares.tkmsVersion) {
+    throw new Error('TkmsVersion mismatch');
+  }
+  if (context.tkmsVersion !== transportKeyPair.tkmsVersion) {
+    throw new Error('TkmsVersion mismatch');
+  }
+
   // also validates `transportKeyPair`
   const tkmsPrivateKey = await transportKeyPairToTkmsPrivateKey(context, transportKeyPair);
 
@@ -31,6 +40,7 @@ export async function decryptKmsSignedcryptedShares(context: Context, parameters
   const orderedDecryptedHandles: readonly ClearValue[] = await context.runtime.decrypt.decryptAndReconstruct({
     shares: kmsSigncryptedShares,
     tkmsPrivateKey,
+    tkmsVersion: context.tkmsVersion,
   });
 
   return orderedDecryptedHandles;

package/core/kms/decryptValuesFromPairs.ts

@@ -1,5 +1,6 @@
 ////////////////////////////////////////////////////////////////////////////////
 
+import type { TkmsVersion } from '../../wasm/tkms/KmsLibApi.js';
 import type { WithDecrypt } from '../types/coreFhevmRuntime.js';
 import type { Handle } from '../types/encryptedTypes-p.js';
 import type { FhevmChain } from '../types/fhevmChain.js';
@@ -8,8 +9,8 @@ import type { ChecksummedAddress, TypedValue } from '../types/primitives.js';
 import type { RelayerDelegatedUserDecryptOptions, RelayerUserDecryptOptions } from '../types/relayer.js';
 import type { SignedDelegatedDecryptionPermit, SignedSelfDecryptionPermit } from '../types/signedDecryptionPermit.js';
 import type { TransportKeyPair } from './TransportKeyPair-p.js';
-import { decryptKmsSignedcryptedShares } from './decryptKmsSignedcryptedShares-p.js';
-import { fetchKmsSignedcryptedShares } from './fetchKmsSignedcryptedShares-p.js';
+import { decryptKmsSigncryptedShares } from './decryptKmsSigncryptedShares-p.js';
+import { fetchKmsSigncryptedShares } from './fetchKmsSigncryptedShares-p.js';
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -17,6 +18,7 @@ type Context = {
   readonly chain: FhevmChain;
   readonly runtime: WithDecrypt;
   readonly client: NonNullable<object>;
+  readonly tkmsVersion: TkmsVersion;
   readonly options: { readonly batchRpcCalls: boolean };
 };
 
@@ -47,10 +49,10 @@ export type ReturnType = readonly TypedValue[];
 export async function decryptValuesFromPairs(fhevm: Context, parameters: Parameters): Promise<ReturnType> {
   const { transportKeyPair: transportKeyPair } = parameters;
 
-  const kmsSigncryptedShares: KmsSigncryptedShares = await fetchKmsSignedcryptedShares(fhevm, parameters);
+  const kmsSigncryptedShares: KmsSigncryptedShares = await fetchKmsSigncryptedShares(fhevm, parameters);
 
   // Using the `KmsSigncryptedShares` decrypt and reconstruct clear values
-  return decryptKmsSignedcryptedShares(fhevm, {
+  return decryptKmsSigncryptedShares(fhevm, {
     kmsSigncryptedShares,
     transportKeyPair: transportKeyPair,
   });

package/core/kms/{fetchKmsSignedcryptedShares-p.ts → fetchKmsSigncryptedShares-p.ts}

@@ -16,6 +16,7 @@ import { checkPersistAllowed } from '../host-contracts/checkPersistAllowed.js';
 import { assertExtraDataMatchesKmsSingersContext } from '../host-contracts/KmsSignersContext-p.js';
 import { createKmsEip712Domain } from './createKmsEip712Domain.js';
 import { checkDelegation } from '../host-contracts/checkDelegation.js';
+import { resolveFhevmTkmsVersion } from '../runtime/resolveFhevmVersions-p.js';
 
 /*
     See: in KMS (eip712Domain)
@@ -64,12 +65,16 @@ type ReturnType = KmsSigncryptedShares;
 const MAX_USER_DECRYPT_CONTRACT_ADDRESSES = 10;
 
 ////////////////////////////////////////////////////////////////////////////////
-// fetchKmsSignedcryptedShares
+// fetchKmsSigncryptedShares
 ////////////////////////////////////////////////////////////////////////////////
 
-export async function fetchKmsSignedcryptedShares(context: Context, parameters: Parameters): Promise<ReturnType> {
+export async function fetchKmsSigncryptedShares(context: Context, parameters: Parameters): Promise<ReturnType> {
   const { signedPermit, options, pairs } = parameters;
 
+  // This helper must support base clients, where TKMS is not mandatory
+  // and tkmsVersion may not be initialized in the CoreFhevm instance yet.
+  const tkmsVersion = await resolveFhevmTkmsVersion(context);
+
   // Check: every requested contractAddress is listed in the permit
   assertPermitIncludesContractAddresses(
     signedPermit,
@@ -121,14 +126,14 @@ export async function fetchKmsSignedcryptedShares(context: Context, parameters:
   // 7. Check: ACL permissions (user is signer or delegatorAddress)
   if (signedPermit.isDelegated) {
     await checkDelegation(context, {
-      address: context.chain.fhevm.contracts.acl.address as ChecksummedAddress,
+      aclAddress: context.chain.fhevm.contracts.acl.address as ChecksummedAddress,
       delegate: signerAddress,
       delegator: encryptedDataOwnerAddress,
       handleContractPairs,
     });
   } else {
     await checkPersistAllowed(context, {
-      address: context.chain.fhevm.contracts.acl.address as ChecksummedAddress,
+      aclAddress: context.chain.fhevm.contracts.acl.address as ChecksummedAddress,
       userAddress: encryptedDataOwnerAddress,
       handleContractPairs,
     });
@@ -152,7 +157,8 @@ export async function fetchKmsSignedcryptedShares(context: Context, parameters:
   // though the context ID matches. Consider comparing the decoded `kmsContextId`
   // instead of the raw `extraData` bytes.
   const requestedKmsSignersContext: KmsSignersContext = await readKmsSignersContext(context, {
-    address: context.chain.fhevm.contracts.kmsVerifier.address as ChecksummedAddress,
+    kmsVerifierAddress: context.chain.fhevm.contracts.kmsVerifier.address as ChecksummedAddress,
+    protocolConfigAddress: context.chain.fhevm.contracts.protocolConfig?.address as ChecksummedAddress | undefined,
   });
 
   assertExtraDataMatchesKmsSingersContext(
@@ -201,6 +207,7 @@ export async function fetchKmsSignedcryptedShares(context: Context, parameters:
     eip712Signature: signature,
     eip712SignerAddress: signerAddress,
     handles,
+    tkmsVersion,
   };
 
   // 12. The returned KmsSigncryptedShares is guaranteed to be fully verified:

package/core/kms/kmsExtraData.ts

@@ -1,49 +1,115 @@
-import type { ByteLength, Bytes32Hex, BytesHex, Uint256BigInt, Uint8Number } from '../types/primitives.js';
+import type { Bytes32Hex, BytesHex, Uint256BigInt, Uint8Number } from '../types/primitives.js';
 import type { ErrorMetadataParams } from '../base/errors/ErrorBase.js';
 import { asUint256BigInt, asUint8Number } from '../base/uint.js';
 import { asBytes32Hex, assertIsBytesHex } from '../base/bytes.js';
 
-const EXTRA_DATA_V1 = 0x01;
+export const EXTRA_DATA_V0: Uint8Number = 0x00 as Uint8Number; // 0x00
+export const EXTRA_DATA_V1: Uint8Number = 0x01 as Uint8Number; // 1 version byte + 32-byte big-endian context ID = 33 bytes
+export const EXTRA_DATA_V2: Uint8Number = 0x02 as Uint8Number; // 1 version byte + 32-byte big-endian context ID + 32-byte big-endian epoch ID = 65 bytes
 
 export function fromKmsExtraData(extraData: BytesHex): {
   version: Uint8Number;
   kmsContextId: Uint256BigInt;
+  kmsEpochId: Uint256BigInt;
 } {
   if ((extraData as string) === '0x00' || (extraData as string) === '0x') {
     return {
-      version: 0 as Uint8Number,
+      version: EXTRA_DATA_V0,
       kmsContextId: 0n as Uint256BigInt,
+      kmsEpochId: 0n as Uint256BigInt,
     };
   }
 
+  if (extraData.length <= 4) {
+    throw new Error(`Unsupported extraData length ${extraData.length}: must be more than 4 bytes`);
+  }
+
   // First byte = version (characters 2-3 after '0x')
-  const version = asUint8Number(parseInt(extraData.slice(2, 4), 16));
+  const version = asUint8Number(Number('0x' + extraData.slice(2, 4)));
 
   if (version === EXTRA_DATA_V1) {
-    // 1 version byte + 32 contextId bytes = 33 bytes = 66 hex chars + 2 for '0x' = 68
-    if (extraData.length < 68) {
-      throw new Error(`extraData too short for v1: expected at least 33 bytes, got ${(extraData.length - 2) / 2}`);
+    // ExtraData v1 format: 1 version byte + 32-byte big-endian context ID = 33 bytes = 66 hex chars + 2 for '0x' = 68
+    // 0x01aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+    // ↑   ↑                                                               ↑
+    // 0   4                                                               68
+    // 0         1         2         3         4         5         6
+    // 01234567890123456789012345678901234567890123456789012345678901234567
+
+    if (extraData.length != 68) {
+      throw new Error(`Invalid extraData length for v1: expected 68, got ${extraData.length}`);
     }
     // 32-byte contextId starts at byte 1 (hex chars 4..67)
     const contextIdBytes32Hex: Bytes32Hex = asBytes32Hex(`0x${extraData.slice(4, 68)}`);
     const kmsContextId = asUint256BigInt(BigInt(contextIdBytes32Hex));
-    return { version, kmsContextId };
+    return { version, kmsContextId, kmsEpochId: 0n as Uint256BigInt };
+  }
+
+  if (version === EXTRA_DATA_V2) {
+    // ExtraData v2 format: 1 version byte + 32-byte big-endian context ID + 32-byte big-endian epoch ID = 65 bytes = 130 hex chars + 2 for '0x' = 132
+    // 0x01aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+    // ↑   ↑                                                               ↑                                                               ↑
+    // 0   4                                                               68                                                              132
+    // 0         1         2         3         4         5         6         7         8         9         0         1         2         3
+    // 012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901
+
+    if (extraData.length != 132) {
+      throw new Error(`Invalid extraData length for v2: expected 132, got ${extraData.length}`);
+    }
+    // 32-byte contextId starts at byte 1 (hex chars 4..67)
+    const contextIdBytes32Hex: Bytes32Hex = asBytes32Hex(`0x${extraData.slice(4, 68)}`);
+    const kmsContextId = asUint256BigInt(BigInt(contextIdBytes32Hex));
+
+    // 32-byte epochId starts at byte 33 (hex chars 68..131)
+    const epochIdBytes32Hex: Bytes32Hex = asBytes32Hex(`0x${extraData.slice(68, 132)}`);
+    const kmsEpochId = asUint256BigInt(BigInt(epochIdBytes32Hex));
+
+    return { version, kmsContextId, kmsEpochId };
   }
 
   throw new Error(`Unsupported extraData version ${version}`);
 }
 
 /**
- * Encodes a KMS context ID into the v1 extraData format.
- * Layout: `0x` + `01` (version byte) + 32-byte big-endian context ID.
+ * Encodes a KMS context into an extraData format.
+ * v0: `0x00`
+ * v1: `0x01` + 32-byte big-endian context ID.
+ * v2: `0x02` + 32-byte big-endian context ID + 32-byte big-endian epoch ID.
  */
 export function toKmsExtraData(args: {
   readonly version: Uint8Number;
   readonly kmsContextId: Uint256BigInt;
+  readonly kmsEpochId: Uint256BigInt;
 }): BytesHex {
   const v = args.version.toString(16).padStart(2, '0');
-  const id = args.kmsContextId.toString(16).padStart(64, '0');
-  return `0x${v}${id}` as BytesHex;
+
+  switch (args.version) {
+    case EXTRA_DATA_V0: {
+      return `0x${v}` as BytesHex;
+    }
+    case EXTRA_DATA_V1: {
+      if (args.kmsContextId === 0n) {
+        throw new Error('kmsContextId cannot be 0 for v1 extraData');
+      }
+      if (args.kmsEpochId !== 0n) {
+        throw new Error('kmsEpochId must be 0 for v1 extraData');
+      }
+      const contextId = args.kmsContextId.toString(16).padStart(64, '0');
+      return `0x${v}${contextId}` as BytesHex;
+    }
+    case EXTRA_DATA_V2: {
+      if (args.kmsContextId === 0n) {
+        throw new Error('kmsContextId cannot be 0 for v2 extraData');
+      }
+      if (args.kmsEpochId === 0n) {
+        throw new Error('kmsEpochId cannot be 0 for v2 extraData');
+      }
+      const contextId = args.kmsContextId.toString(16).padStart(64, '0');
+      const epochId = args.kmsEpochId.toString(16).padStart(64, '0');
+      return `0x${v}${contextId}${epochId}` as BytesHex;
+    }
+    default:
+      throw new Error(`Unsupported extraData version ${args.version}`);
+  }
 }
 
 export function assertIsKmsExtraData(
@@ -53,6 +119,9 @@ export function assertIsKmsExtraData(
   if (value === '0x00' || value === '0x') {
     return;
   }
-  assertIsBytesHex(value, { ...options, byteLength: 33 as ByteLength });
+
+  assertIsBytesHex(value, options);
+
+  // Will valid extraData length too
   fromKmsExtraData(value);
 }