@cofhe/sdk 0.1.0 → 0.1.1

package/core/encrypt/zkPackProveVerify.ts

@@ -0,0 +1,273 @@
+/* eslint-disable no-unused-vars */
+
+import { CofhesdkError, CofhesdkErrorCode } from '../error.js';
+import { type EncryptableItem, FheTypes } from '../types.js';
+import { toBigIntOrThrow, validateBigIntInRange, toHexString, hexToBytes } from '../utils.js';
+
+// ===== TYPES =====
+
+export type VerifyResultRaw = {
+  ct_hash: string;
+  signature: string;
+  recid: number;
+};
+
+export type VerifyResult = {
+  ct_hash: string;
+  signature: string;
+};
+
+export type ZkProvenCiphertextList = {
+  serialize(): Uint8Array;
+};
+
+export type ZkCompactPkeCrs = {
+  free(): void;
+  serialize(compress: boolean): Uint8Array;
+  safe_serialize(serialized_size_limit: bigint): Uint8Array;
+};
+
+export type ZkCompactPkeCrsConstructor = {
+  deserialize(buffer: Uint8Array): ZkCompactPkeCrs;
+  safe_deserialize(buffer: Uint8Array, serialized_size_limit: bigint): ZkCompactPkeCrs;
+  from_config(config: unknown, max_num_bits: number): ZkCompactPkeCrs;
+  deserialize_from_public_params(buffer: Uint8Array): ZkCompactPkeCrs;
+  safe_deserialize_from_public_params(buffer: Uint8Array, serialized_size_limit: bigint): ZkCompactPkeCrs;
+};
+
+export type ZkCiphertextListBuilder = {
+  push_boolean(data: boolean): void;
+  push_u8(data: number): void;
+  push_u16(data: number): void;
+  push_u32(data: number): void;
+  push_u64(data: bigint): void;
+  push_u128(data: bigint): void;
+  push_u160(data: bigint): void;
+  build_with_proof_packed(
+    crs: ZkCompactPkeCrs,
+    metadata: Uint8Array,
+    computeLoad: 1 // ZkComputeLoad.Verify
+  ): ZkProvenCiphertextList;
+};
+
+export type ZkBuilderAndCrsGenerator = (
+  fhe: string,
+  crs: string
+) => { zkBuilder: ZkCiphertextListBuilder; zkCrs: ZkCompactPkeCrs };
+
+// ===== CONSTANTS =====
+
+export const MAX_UINT8: bigint = 255n;
+export const MAX_UINT16: bigint = 65535n;
+export const MAX_UINT32: bigint = 4294967295n;
+export const MAX_UINT64: bigint = 18446744073709551615n; // 2^64 - 1
+export const MAX_UINT128: bigint = 340282366920938463463374607431768211455n; // 2^128 - 1
+export const MAX_UINT256: bigint = 115792089237316195423570985008687907853269984665640564039457584007913129640319n; // 2^256 - 1
+export const MAX_UINT160: bigint = 1461501637330902918203684832716283019655932542975n; // 2^160 - 1
+export const MAX_ENCRYPTABLE_BITS: number = 2048;
+
+// ===== CORE FUNCTIONS =====
+
+export const zkPack = (items: EncryptableItem[], builder: ZkCiphertextListBuilder): ZkCiphertextListBuilder => {
+  let totalBits = 0;
+  for (const item of items) {
+    switch (item.utype) {
+      case FheTypes.Bool: {
+        builder.push_boolean(item.data);
+        totalBits += 1;
+        break;
+      }
+      case FheTypes.Uint8: {
+        const bint = toBigIntOrThrow(item.data);
+        validateBigIntInRange(bint, MAX_UINT8);
+        builder.push_u8(parseInt(bint.toString()));
+        totalBits += 8;
+        break;
+      }
+      case FheTypes.Uint16: {
+        const bint = toBigIntOrThrow(item.data);
+        validateBigIntInRange(bint, MAX_UINT16);
+        builder.push_u16(parseInt(bint.toString()));
+        totalBits += 16;
+        break;
+      }
+      case FheTypes.Uint32: {
+        const bint = toBigIntOrThrow(item.data);
+        validateBigIntInRange(bint, MAX_UINT32);
+        builder.push_u32(parseInt(bint.toString()));
+        totalBits += 32;
+        break;
+      }
+      case FheTypes.Uint64: {
+        const bint = toBigIntOrThrow(item.data);
+        validateBigIntInRange(bint, MAX_UINT64);
+        builder.push_u64(bint);
+        totalBits += 64;
+        break;
+      }
+      case FheTypes.Uint128: {
+        const bint = toBigIntOrThrow(item.data);
+        validateBigIntInRange(bint, MAX_UINT128);
+        builder.push_u128(bint);
+        totalBits += 128;
+        break;
+      }
+      // [U256-DISABLED]
+      // case FheTypes.Uint256: {
+      //   const bint = toBigIntOrThrow(item.data);
+      //   validateBigIntInRange(bint, MAX_UINT256);
+      //   builder.push_u256(bint);
+      //   totalBits += 256;
+      //   break;
+      // }
+      case FheTypes.Uint160: {
+        const bint = toBigIntOrThrow(item.data);
+        validateBigIntInRange(bint, MAX_UINT160);
+        builder.push_u160(bint);
+        totalBits += 160;
+        break;
+      }
+      default: {
+        throw new CofhesdkError({
+          code: CofhesdkErrorCode.ZkPackFailed,
+          message: `Invalid utype: ${(item as any).utype}`,
+          hint: `Ensure that the utype is valid, using the Encryptable type, for example: Encryptable.uint128(100n)`,
+          context: {
+            item,
+          },
+        });
+      }
+    }
+  }
+
+  if (totalBits > MAX_ENCRYPTABLE_BITS) {
+    throw new CofhesdkError({
+      code: CofhesdkErrorCode.ZkPackFailed,
+      message: `Total bits ${totalBits} exceeds ${MAX_ENCRYPTABLE_BITS}`,
+      hint: `Ensure that the total bits of the items to encrypt does not exceed ${MAX_ENCRYPTABLE_BITS}`,
+      context: {
+        totalBits,
+        maxBits: MAX_ENCRYPTABLE_BITS,
+        items,
+      },
+    });
+  }
+
+  return builder;
+};
+
+export const zkProve = async (
+  builder: ZkCiphertextListBuilder,
+  crs: ZkCompactPkeCrs,
+  address: string,
+  securityZone: number,
+  chainId: number
+): Promise<Uint8Array> => {
+  const metadata = constructZkPoKMetadata(address, securityZone, chainId);
+
+  return new Promise((resolve) => {
+    setTimeout(() => {
+      const compactList = builder.build_with_proof_packed(
+        crs,
+        metadata,
+        1 // ZkComputeLoad.Verify
+      );
+
+      resolve(compactList.serialize());
+    }, 0);
+  });
+};
+
+const constructZkPoKMetadata = (accountAddr: string, securityZone: number, chainId: number): Uint8Array => {
+  // Decode the account address from hex
+  const accountAddrNoPrefix = accountAddr.startsWith('0x') ? accountAddr.slice(2) : accountAddr;
+  const accountBytes = hexToBytes(accountAddrNoPrefix);
+
+  // Encode chainId as 32 bytes (u256) in big-endian format
+  const chainIdBytes = new Uint8Array(32);
+
+  // Since chain IDs are typically small numbers, we can just encode them
+  // directly without BigInt operations, filling only the necessary bytes
+  // from the right (least significant)
+  let value = chainId;
+  for (let i = 31; i >= 0 && value > 0; i--) {
+    chainIdBytes[i] = value & 0xff;
+    value = value >>> 8;
+  }
+
+  const metadata = new Uint8Array(1 + accountBytes.length + 32);
+  metadata[0] = securityZone;
+  metadata.set(accountBytes, 1);
+  metadata.set(chainIdBytes, 1 + accountBytes.length);
+
+  return metadata;
+};
+
+export const zkVerify = async (
+  verifierUrl: string,
+  serializedBytes: Uint8Array,
+  address: string,
+  securityZone: number,
+  chainId: number
+): Promise<VerifyResult[]> => {
+  // Convert bytearray to hex string
+  const packed_list = toHexString(serializedBytes);
+
+  const sz_byte = new Uint8Array([securityZone]);
+
+  // Construct request payload
+  const payload = {
+    packed_list,
+    account_addr: address,
+    security_zone: sz_byte[0],
+    chain_id: chainId,
+  };
+
+  const body = JSON.stringify(payload);
+
+  // Send request to verification server
+  try {
+    const response = await fetch(`${verifierUrl}/verify`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body,
+    });
+
+    if (!response.ok) {
+      // Get the response body as text for better error details
+      const errorBody = await response.text();
+      throw new CofhesdkError({
+        code: CofhesdkErrorCode.ZkVerifyFailed,
+        message: `HTTP error! ZK proof verification failed - ${errorBody}`,
+      });
+    }
+
+    const json = (await response.json()) as { status: string; data: VerifyResultRaw[]; error: string };
+
+    if (json.status !== 'success') {
+      throw new CofhesdkError({
+        code: CofhesdkErrorCode.ZkVerifyFailed,
+        message: `ZK proof verification response malformed - ${json.error}`,
+      });
+    }
+
+    return json.data.map(({ ct_hash, signature, recid }) => {
+      return {
+        ct_hash,
+        signature: concatSigRecid(signature, recid),
+      };
+    });
+  } catch (e) {
+    throw new CofhesdkError({
+      code: CofhesdkErrorCode.ZkVerifyFailed,
+      message: `ZK proof verification failed`,
+      cause: e instanceof Error ? e : undefined,
+    });
+  }
+};
+
+const concatSigRecid = (signature: string, recid: number): string => {
+  return signature + (recid + 27).toString(16).padStart(2, '0');
+};

package/core/error.ts

@@ -0,0 +1,170 @@
+/* eslint-disable no-unused-vars */
+
+export enum CofhesdkErrorCode {
+  InternalError = 'INTERNAL_ERROR',
+  UnknownEnvironment = 'UNKNOWN_ENVIRONMENT',
+  InitTfheFailed = 'INIT_TFHE_FAILED',
+  InitViemFailed = 'INIT_VIEM_FAILED',
+  InitEthersFailed = 'INIT_ETHERS_FAILED',
+  NotConnected = 'NOT_CONNECTED',
+  MissingPublicClient = 'MISSING_PUBLIC_CLIENT',
+  MissingWalletClient = 'MISSING_WALLET_CLIENT',
+  MissingProviderParam = 'MISSING_PROVIDER_PARAM',
+  EmptySecurityZonesParam = 'EMPTY_SECURITY_ZONES_PARAM',
+  InvalidPermitData = 'INVALID_PERMIT_DATA',
+  InvalidPermitDomain = 'INVALID_PERMIT_DOMAIN',
+  PermitNotFound = 'PERMIT_NOT_FOUND',
+  CannotRemoveLastPermit = 'CANNOT_REMOVE_LAST_PERMIT',
+  AccountUninitialized = 'ACCOUNT_UNINITIALIZED',
+  ChainIdUninitialized = 'CHAIN_ID_UNINITIALIZED',
+  SealOutputFailed = 'SEAL_OUTPUT_FAILED',
+  SealOutputReturnedNull = 'SEAL_OUTPUT_RETURNED_NULL',
+  InvalidUtype = 'INVALID_UTYPE',
+  DecryptFailed = 'DECRYPT_FAILED',
+  DecryptReturnedNull = 'DECRYPT_RETURNED_NULL',
+  ZkMocksInsertCtHashesFailed = 'ZK_MOCKS_INSERT_CT_HASHES_FAILED',
+  ZkMocksCalcCtHashesFailed = 'ZK_MOCKS_CALC_CT_HASHES_FAILED',
+  ZkMocksVerifySignFailed = 'ZK_MOCKS_VERIFY_SIGN_FAILED',
+  ZkMocksCreateProofSignatureFailed = 'ZK_MOCKS_CREATE_PROOF_SIGNATURE_FAILED',
+  ZkVerifyFailed = 'ZK_VERIFY_FAILED',
+  ZkPackFailed = 'ZK_PACK_FAILED',
+  ZkProveFailed = 'ZK_PROVE_FAILED',
+  EncryptRemainingInItems = 'ENCRYPT_REMAINING_IN_ITEMS',
+  ZkUninitialized = 'ZK_UNINITIALIZED',
+  ZkVerifierUrlUninitialized = 'ZK_VERIFIER_URL_UNINITIALIZED',
+  ThresholdNetworkUrlUninitialized = 'THRESHOLD_NETWORK_URL_UNINITIALIZED',
+  MissingConfig = 'MISSING_CONFIG',
+  UnsupportedChain = 'UNSUPPORTED_CHAIN',
+  MissingZkBuilderAndCrsGenerator = 'MISSING_ZK_BUILDER_AND_CRS_GENERATOR',
+  MissingTfhePublicKeyDeserializer = 'MISSING_TFHE_PUBLIC_KEY_DESERIALIZER',
+  MissingCompactPkeCrsDeserializer = 'MISSING_COMPACT_PKE_CRS_DESERIALIZER',
+  MissingFheKey = 'MISSING_FHE_KEY',
+  MissingCrs = 'MISSING_CRS',
+  FetchKeysFailed = 'FETCH_KEYS_FAILED',
+  PublicWalletGetChainIdFailed = 'PUBLIC_WALLET_GET_CHAIN_ID_FAILED',
+  PublicWalletGetAddressesFailed = 'PUBLIC_WALLET_GET_ADDRESSES_FAILED',
+  RehydrateKeysStoreFailed = 'REHYDRATE_KEYS_STORE_FAILED',
+}
+
+export type CofhesdkErrorParams = {
+  code: CofhesdkErrorCode;
+  message: string;
+  cause?: Error;
+  hint?: string;
+  context?: Record<string, unknown>;
+};
+
+/**
+ * CofhesdkError class
+ * This class is used to create errors that are specific to the CoFHE SDK
+ * It extends the Error class and adds a code, cause, hint, and context
+ * The code is used to identify the type of error
+ * The cause is used to indicate the inner error that caused the CofhesdkError
+ * The hint is used to provide a hint about how to fix the error
+ * The context is used to provide additional context about the state that caused the error
+ * The serialize method is used to serialize the error to a JSON string
+ * The toString method is used to provide a human-readable string representation of the error
+ */
+export class CofhesdkError extends Error {
+  public readonly code: CofhesdkErrorCode;
+  public readonly cause?: Error;
+  public readonly hint?: string;
+  public readonly context?: Record<string, unknown>;
+
+  constructor({ code, message, cause, hint, context }: CofhesdkErrorParams) {
+    // If there's a cause, append its message to provide full context
+    const fullMessage = cause ? `${message} | Caused by: ${cause.message}` : message;
+
+    super(fullMessage);
+    this.name = 'CofhesdkError';
+    this.code = code;
+    this.cause = cause;
+    this.hint = hint;
+    this.context = context;
+
+    // Maintains proper stack trace for where our error was thrown (only available on V8)
+    if (Error.captureStackTrace) {
+      Error.captureStackTrace(this, CofhesdkError);
+    }
+  }
+
+  /**
+   * Creates a CofhesdkError from an unknown error
+   * If the error is a CofhesdkError, it is returned unchanged, else a new CofhesdkError is created
+   * If a wrapperError is provided, it is used to create the new CofhesdkError, else a default is used
+   */
+  static fromError(error: unknown, wrapperError?: CofhesdkErrorParams): CofhesdkError {
+    if (isCofhesdkError(error)) return error;
+
+    const cause = error instanceof Error ? error : new Error(`${error}`);
+
+    return new CofhesdkError({
+      code: wrapperError?.code ?? CofhesdkErrorCode.InternalError,
+      message: wrapperError?.message ?? 'An internal error occurred',
+      hint: wrapperError?.hint,
+      context: wrapperError?.context,
+      cause: cause,
+    });
+  }
+
+  /**
+   * Serializes the error to JSON string with proper handling of Error objects
+   */
+  serialize(): string {
+    return bigintSafeJsonStringify({
+      name: this.name,
+      code: this.code,
+      message: this.message,
+      hint: this.hint,
+      context: this.context,
+      cause: this.cause
+        ? {
+            name: this.cause.name,
+            message: this.cause.message,
+            stack: this.cause.stack,
+          }
+        : undefined,
+      stack: this.stack,
+    });
+  }
+
+  /**
+   * Returns a human-readable string representation of the error
+   */
+  toString(): string {
+    const parts = [`${this.name} [${this.code}]: ${this.message}`];
+
+    if (this.hint) {
+      parts.push(`Hint: ${this.hint}`);
+    }
+
+    if (this.context && Object.keys(this.context).length > 0) {
+      parts.push(`Context: ${bigintSafeJsonStringify(this.context)}`);
+    }
+
+    if (this.stack) {
+      parts.push(`\nStack trace:`);
+      parts.push(this.stack);
+    }
+
+    if (this.cause) {
+      parts.push(`\nCaused by: ${this.cause.name}: ${this.cause.message}`);
+      if (this.cause.stack) {
+        parts.push(this.cause.stack);
+      }
+    }
+
+    return parts.join('\n');
+  }
+}
+
+const bigintSafeJsonStringify = (value: unknown): string => {
+  return JSON.stringify(value, (key, value) => {
+    if (typeof value === 'bigint') {
+      return `${value}n`;
+    }
+    return value;
+  });
+};
+
+export const isCofhesdkError = (error: unknown): error is CofhesdkError => error instanceof CofhesdkError;

package/core/fetchKeys.test.ts

@@ -0,0 +1,212 @@
+/* eslint-disable turbo/no-undeclared-env-vars */
+/* eslint-disable no-undef */
+
+import { sepolia, arbSepolia } from '@/chains';
+
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { fetchKeys, fetchMultichainKeys } from './fetchKeys.js';
+import { type CofhesdkConfig, createCofhesdkConfigBase } from './config.js';
+import { createKeysStore, type KeysStorage } from './keyStore.js';
+
+describe('fetchKeys', () => {
+  let config: CofhesdkConfig;
+  let mockTfhePublicKeyDeserializer: any;
+  let mockCompactPkeCrsDeserializer: any;
+  let keysStorage: KeysStorage;
+
+  beforeEach(async () => {
+    // Reset all mocks
+    vi.clearAllMocks();
+
+    // Setup config with real chains
+    config = createCofhesdkConfigBase({
+      supportedChains: [sepolia, arbSepolia],
+    });
+
+    // Setup mock serializers
+    mockTfhePublicKeyDeserializer = vi.fn();
+    mockCompactPkeCrsDeserializer = vi.fn();
+
+    // Create a fresh keysStorage instance for each test (non-persisted)
+    keysStorage = createKeysStore(null);
+  });
+
+  afterEach(async () => {
+    vi.restoreAllMocks();
+  });
+
+  it('should fetch and store FHE public key and CRS for Sepolia when not cached', async () => {
+    const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
+      config,
+      sepolia.id,
+      0,
+      mockTfhePublicKeyDeserializer,
+      mockCompactPkeCrsDeserializer,
+      keysStorage
+    );
+
+    expect(fheKeyFetchedFromCoFHE).toBe(true);
+    expect(crsFetchedFromCoFHE).toBe(true);
+
+    // Verify keys were stored
+    const storedFheKey = keysStorage.getFheKey(sepolia.id, 0);
+    const storedCrs = keysStorage.getCrs(sepolia.id);
+
+    expect(storedFheKey).toBeDefined();
+    expect(storedCrs).toBeDefined();
+    expect(mockTfhePublicKeyDeserializer).toHaveBeenCalled();
+    expect(mockCompactPkeCrsDeserializer).toHaveBeenCalled();
+  });
+
+  it('should fetch and store FHE public key and CRS for Arbitrum Sepolia when not cached', async () => {
+    const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
+      config,
+      arbSepolia.id,
+      0,
+      mockTfhePublicKeyDeserializer,
+      mockCompactPkeCrsDeserializer,
+      keysStorage
+    );
+
+    expect(fheKeyFetchedFromCoFHE).toBe(true);
+    expect(crsFetchedFromCoFHE).toBe(true);
+
+    // Verify keys were stored
+    const storedFheKey = keysStorage.getFheKey(arbSepolia.id, 0);
+    const storedCrs = keysStorage.getCrs(arbSepolia.id);
+
+    expect(storedFheKey).toBeDefined();
+    expect(storedCrs).toBeDefined();
+    expect(mockTfhePublicKeyDeserializer).toHaveBeenCalled();
+    expect(mockCompactPkeCrsDeserializer).toHaveBeenCalled();
+  });
+
+  it('should not fetch FHE key if already cached', async () => {
+    // Pre-populate with a cached key
+    const mockCachedKey = '0x1234567890';
+    keysStorage.setFheKey(sepolia.id, 0, mockCachedKey);
+
+    const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
+      config,
+      sepolia.id,
+      0,
+      mockTfhePublicKeyDeserializer,
+      mockCompactPkeCrsDeserializer,
+      keysStorage
+    );
+
+    expect(fheKeyFetchedFromCoFHE).toBe(false);
+    expect(crsFetchedFromCoFHE).toBe(true);
+
+    // Verify the cached key wasn't overwritten
+    const retrievedKey = keysStorage.getFheKey(sepolia.id, 0);
+    expect(retrievedKey).toEqual(mockCachedKey);
+
+    // Verify CRS was still fetched
+    const retrievedCrs = keysStorage.getCrs(sepolia.id);
+    expect(retrievedCrs).toBeDefined();
+  });
+
+  it('should not fetch CRS if already cached', async () => {
+    // Pre-populate with a cached CRS
+    const mockCachedCrs = '0x2345678901';
+    keysStorage.setCrs(sepolia.id, mockCachedCrs);
+
+    const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
+      config,
+      sepolia.id,
+      0,
+      mockTfhePublicKeyDeserializer,
+      mockCompactPkeCrsDeserializer,
+      keysStorage
+    );
+
+    expect(fheKeyFetchedFromCoFHE).toBe(true);
+    expect(crsFetchedFromCoFHE).toBe(false);
+
+    // Verify the cached CRS wasn't overwritten
+    const retrievedCrs = keysStorage.getCrs(sepolia.id);
+    expect(retrievedCrs).toEqual(mockCachedCrs);
+
+    // Verify FHE key was still fetched
+    const retrievedKey = keysStorage.getFheKey(sepolia.id, 0);
+    expect(retrievedKey).toBeDefined();
+  });
+
+  it('should not make any network calls if both keys are cached', async () => {
+    // Pre-populate both keys
+    const mockCachedKey = '0x1234567890';
+    const mockCachedCrs = '0x2345678901';
+    keysStorage.setFheKey(sepolia.id, 0, mockCachedKey);
+    keysStorage.setCrs(sepolia.id, mockCachedCrs);
+
+    const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
+      config,
+      sepolia.id,
+      0,
+      mockTfhePublicKeyDeserializer,
+      mockCompactPkeCrsDeserializer,
+      keysStorage
+    );
+
+    expect(fheKeyFetchedFromCoFHE).toBe(false);
+    expect(crsFetchedFromCoFHE).toBe(false);
+
+    // Verify both keys remain unchanged
+    const retrievedKey = keysStorage.getFheKey(sepolia.id, 0);
+    const retrievedCrs = keysStorage.getCrs(sepolia.id);
+
+    expect(retrievedKey).toEqual(mockCachedKey);
+    expect(retrievedCrs).toEqual(mockCachedCrs);
+  });
+
+  it('should throw error for unsupported chain ID', async () => {
+    await expect(
+      fetchKeys(
+        config,
+        999, // Non-existent chain
+        0,
+        mockTfhePublicKeyDeserializer,
+        mockCompactPkeCrsDeserializer,
+        keysStorage
+      )
+    ).rejects.toThrow('Config does not support chain <999>');
+  });
+
+  it('should throw error when FHE public key serialization fails', async () => {
+    mockTfhePublicKeyDeserializer.mockImplementation(() => {
+      throw new Error('Serialization failed');
+    });
+
+    await expect(
+      fetchKeys(config, sepolia.id, 0, mockTfhePublicKeyDeserializer, mockCompactPkeCrsDeserializer, keysStorage)
+    ).rejects.toThrow('Error serializing FHE publicKey; Error: Serialization failed');
+  });
+
+  it('should throw error when CRS serialization fails', async () => {
+    mockCompactPkeCrsDeserializer.mockImplementation(() => {
+      throw new Error('Serialization failed');
+    });
+
+    await expect(
+      fetchKeys(config, sepolia.id, 0, mockTfhePublicKeyDeserializer, mockCompactPkeCrsDeserializer, keysStorage)
+    ).rejects.toThrow('Error serializing CRS; Error: Serialization failed');
+  });
+
+  it('should fetch and store FHE public key and CRS for all chains in the config', async () => {
+    await fetchMultichainKeys(config, 0, mockTfhePublicKeyDeserializer, mockCompactPkeCrsDeserializer, keysStorage);
+
+    // Verify keys were stored
+    const storedFheKey = keysStorage.getFheKey(sepolia.id, 0);
+    const storedCrs = keysStorage.getCrs(sepolia.id);
+    const storedFheKeyArb = keysStorage.getFheKey(arbSepolia.id, 0);
+    const storedCrsArb = keysStorage.getCrs(arbSepolia.id);
+
+    expect(storedFheKey).toBeDefined();
+    expect(storedCrs).toBeDefined();
+    expect(storedFheKeyArb).toBeDefined();
+    expect(storedCrsArb).toBeDefined();
+    expect(mockTfhePublicKeyDeserializer).toHaveBeenCalled();
+    expect(mockCompactPkeCrsDeserializer).toHaveBeenCalled();
+  });
+});