@cofhe/sdk 0.1.0 → 0.2.0

package/core/fetchKeys.ts

@@ -0,0 +1,144 @@
+import { hardhat } from '@/chains';
+
+import { type CofhesdkConfig, getCoFheUrlOrThrow } from './config.js';
+import { type KeysStorage } from './keyStore.js';
+
+const PUBLIC_KEY_LENGTH_MIN = 15_000;
+export type FheKeyDeserializer = (buff: string) => void;
+
+const checkKeyValidity = (key: string | undefined, serializer: FheKeyDeserializer) => {
+  if (key == null || key.length === 0) return [false, `Key is null or empty <${key}>`];
+  try {
+    serializer(key);
+    return [true, `Key is valid`];
+  } catch (err) {
+    return [false, `Serialization failed <${err}> key length <${key.length}>`];
+  }
+};
+
+const fetchFhePublicKey = async (
+  coFheUrl: string,
+  chainId: number,
+  securityZone: number,
+  tfhePublicKeyDeserializer: FheKeyDeserializer,
+  keysStorage?: KeysStorage | null
+): Promise<[string, boolean]> => {
+  // Escape if key already exists
+  const storedKey = keysStorage?.getFheKey(chainId, securityZone);
+  const [storedKeyValid] = checkKeyValidity(storedKey, tfhePublicKeyDeserializer);
+  if (storedKeyValid) return [storedKey!, false];
+
+  let pk_data: string | undefined = undefined;
+
+  try {
+    const pk_res = await fetch(`${coFheUrl}/GetNetworkPublicKey`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ securityZone }),
+    });
+    const json = (await pk_res.json()) as { publicKey: string };
+    pk_data = json.publicKey;
+  } catch (err) {
+    throw new Error(`Error fetching FHE publicKey; fetching from CoFHE failed with error ${err}`);
+  }
+
+  if (pk_data == null || typeof pk_data !== 'string') {
+    throw new Error(`Error fetching FHE publicKey; fetched result invalid: missing or not a string`);
+  }
+
+  if (pk_data === '0x') {
+    throw new Error('Error fetching FHE publicKey; provided chain is not FHE enabled / not found');
+  }
+
+  if (pk_data.length < PUBLIC_KEY_LENGTH_MIN) {
+    throw new Error(
+      `Error fetching FHE publicKey; got shorter than expected key length: ${pk_data.length}. Expected length >= ${PUBLIC_KEY_LENGTH_MIN}`
+    );
+  }
+
+  // Check validity by serializing
+  try {
+    tfhePublicKeyDeserializer(pk_data);
+  } catch (err) {
+    throw new Error(`Error serializing FHE publicKey; ${err}`);
+  }
+
+  // Store result
+  keysStorage?.setFheKey(chainId, securityZone, pk_data);
+
+  return [pk_data, true];
+};
+
+const fetchCrs = async (
+  coFheUrl: string,
+  chainId: number,
+  securityZone: number,
+  compactPkeCrsDeserializer: FheKeyDeserializer,
+  keysStorage?: KeysStorage | null
+): Promise<[string, boolean]> => {
+  // Escape if key already exists
+  const storedKey = keysStorage?.getCrs(chainId);
+  const [storedKeyValid] = checkKeyValidity(storedKey, compactPkeCrsDeserializer);
+  if (storedKeyValid) return [storedKey!, false];
+
+  let crs_data: string | undefined = undefined;
+
+  try {
+    const crs_res = await fetch(`${coFheUrl}/GetCrs`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ securityZone }),
+    });
+    const json = (await crs_res.json()) as { crs: string };
+    crs_data = json.crs;
+  } catch (err) {
+    throw new Error(`Error fetching CRS; fetching failed with error ${err}`);
+  }
+
+  if (crs_data == null || typeof crs_data !== 'string') {
+    throw new Error(`Error fetching CRS; invalid: missing or not a string`);
+  }
+
+  try {
+    compactPkeCrsDeserializer(crs_data);
+  } catch (err) {
+    console.error(`Error serializing CRS ${err}`);
+    throw new Error(`Error serializing CRS; ${err}`);
+  }
+
+  keysStorage?.setCrs(chainId, crs_data);
+
+  return [crs_data, true];
+};
+
+/**
+ * Retrieves the FHE public key and the CRS from the provider.
+ * If the key/crs already exists in the store it is returned, else it is fetched, stored, and returned
+ * @param {CofhesdkConfig} config - The configuration object for the CoFHE SDK
+ * @param {number} chainId - The chain to fetch the FHE key for, if no chainId provided, undefined is returned
+ * @param securityZone - The security zone for which to retrieve the key (default 0).
+ * @param tfhePublicKeyDeserializer - The serializer for the FHE public key (used for validation).
+ * @param compactPkeCrsDeserializer - The serializer for the CRS (used for validation).
+ * @param keysStorage - The keys storage instance to use (optional)
+ * @returns {Promise<[[string, boolean], [string, boolean]]>} - A promise that resolves to [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]]
+ */
+export const fetchKeys = async (
+  config: CofhesdkConfig,
+  chainId: number,
+  securityZone: number = 0,
+  tfhePublicKeyDeserializer: FheKeyDeserializer,
+  compactPkeCrsDeserializer: FheKeyDeserializer,
+  keysStorage?: KeysStorage | null
+): Promise<[[string, boolean], [string, boolean]]> => {
+  // Get cofhe url from config
+  const coFheUrl = getCoFheUrlOrThrow(config, chainId);
+
+  return await Promise.all([
+    fetchFhePublicKey(coFheUrl, chainId, securityZone, tfhePublicKeyDeserializer, keysStorage),
+    fetchCrs(coFheUrl, chainId, securityZone, compactPkeCrsDeserializer, keysStorage),
+  ]);
+};

package/core/index.ts

@@ -0,0 +1,89 @@
+// Client (base implementations)
+export { createCofhesdkClientBase, InitialConnectStore as CONNECT_STORE_DEFAULTS } from './client.js';
+
+// Configuration (base implementations)
+export { createCofhesdkConfigBase, getCofhesdkConfigItem } from './config.js';
+export type { CofhesdkConfig, CofhesdkInputConfig, CofhesdkInternalConfig } from './config.js';
+
+// Types
+export type {
+  // Client types
+  CofhesdkClient,
+  CofhesdkClientParams,
+  CofhesdkClientConnectionState,
+  CofhesdkClientPermits,
+} from './clientTypes.js';
+
+export type {
+  IStorage,
+  // Primitive types
+  Primitive,
+  LiteralToPrimitive,
+  // Encryptable types
+  EncryptableItem,
+  EncryptableBool,
+  EncryptableUint8,
+  EncryptableUint16,
+  EncryptableUint32,
+  EncryptableUint64,
+  EncryptableUint128,
+  EncryptableAddress,
+  // Encrypted types
+  EncryptedNumber,
+  EncryptedItemInput,
+  EncryptedBoolInput,
+  EncryptedUint8Input,
+  EncryptedUint16Input,
+  EncryptedUint32Input,
+  EncryptedUint64Input,
+  EncryptedUint128Input,
+  EncryptedAddressInput,
+  EncryptedItemInputs,
+  EncryptableToEncryptedItemInputMap,
+  FheTypeValue,
+  // Decryption types
+  UnsealedItem,
+  // Util types
+  EncryptStepCallbackFunction as EncryptSetStateFn,
+  EncryptStepCallbackContext,
+} from './types.js';
+export {
+  FheTypes,
+  FheUintUTypes,
+  FheAllUTypes,
+  Encryptable,
+  isEncryptableItem,
+  EncryptStep,
+  isLastEncryptionStep,
+  assertCorrectEncryptedItemInput,
+} from './types.js';
+
+// Error handling
+export { CofhesdkError, CofhesdkErrorCode, isCofhesdkError } from './error.js';
+export type { CofhesdkErrorParams } from './error.js';
+
+// Key fetching
+export { fetchKeys } from './fetchKeys.js';
+export type { FheKeyDeserializer } from './fetchKeys.js';
+
+// Key storage
+export { createKeysStore } from './keyStore.js';
+export type { KeysStorage, KeysStore } from './keyStore.js';
+
+// Builders (exported via client, but can be imported directly for typing)
+export { EncryptInputsBuilder } from './encrypt/encryptInputsBuilder.js';
+export { DecryptHandlesBuilder } from './decrypt/decryptHandleBuilder.js';
+
+// ZK utilities
+export type {
+  ZkBuilderAndCrsGenerator,
+  ZkProveWorkerFunction,
+  ZkProveWorkerRequest,
+  ZkProveWorkerResponse,
+} from './encrypt/zkPackProveVerify.js';
+export { zkProveWithWorker } from './encrypt/zkPackProveVerify.js';
+
+export { MOCKS_ZK_VERIFIER_SIGNER_ADDRESS } from './encrypt/cofheMocksZkVerifySign.js';
+
+// Utils
+export { fheTypeToString } from './utils.js';

package/core/keyStore.test.ts

@@ -0,0 +1,226 @@
+/* eslint-disable turbo/no-undeclared-env-vars */
+/* eslint-disable no-undef */
+
+import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
+import { createKeysStore, type KeysStore, type KeysStorage } from './keyStore.js';
+
+// Mock the storage module
+const mockStorage = {
+  getItem: vi.fn(),
+  setItem: vi.fn(),
+  removeItem: vi.fn(),
+};
+
+vi.mock('../src/storage', () => ({
+  getStorage: () => mockStorage,
+}));
+
+describe('KeyStore', () => {
+  let keysStorage: KeysStorage;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    // Create a fresh keysStorage instance for each test (with mock storage)
+    keysStorage = createKeysStore(mockStorage as any);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe('Store Structure', () => {
+    it('should have initial empty state', () => {
+      const state = keysStorage.store.getState();
+      expect(state).toEqual({
+        fhe: {},
+        crs: {},
+      });
+    });
+
+    it('should be a Zustand store with persist', () => {
+      expect(keysStorage.store).toBeDefined();
+      expect(keysStorage.store.getState).toBeDefined();
+      expect(keysStorage.store.setState).toBeDefined();
+      // In test environment, persist might not be fully initialized
+      if ('persist' in keysStorage.store) {
+        expect((keysStorage.store as any).persist).toBeDefined();
+      }
+    });
+  });
+
+  describe('FHE Key Management', () => {
+    const testChainId = 1337;
+    const testSecurityZone = 0;
+    const testKey = '0x1234567890';
+
+    it('should set and get FHE key', () => {
+      keysStorage.setFheKey(testChainId, testSecurityZone, testKey);
+
+      const retrievedKey = keysStorage.getFheKey(testChainId, testSecurityZone);
+
+      expect(retrievedKey).toEqual(testKey);
+    });
+
+    it('should handle multiple security zones', () => {
+      const key0 = '0x1234567890';
+      const key1 = '0x2345678901';
+
+      keysStorage.setFheKey(testChainId, 0, key0);
+      keysStorage.setFheKey(testChainId, 1, key1);
+
+      expect(keysStorage.getFheKey(testChainId, 0)).toEqual(key0);
+      expect(keysStorage.getFheKey(testChainId, 1)).toEqual(key1);
+    });
+
+    it('should handle multiple chains', () => {
+      const chain1Key = '0x1234567890';
+      const chain2Key = '0x2345678901';
+
+      keysStorage.setFheKey(1, testSecurityZone, chain1Key);
+      keysStorage.setFheKey(2, testSecurityZone, chain2Key);
+
+      expect(keysStorage.getFheKey(1, testSecurityZone)).toEqual(chain1Key);
+      expect(keysStorage.getFheKey(2, testSecurityZone)).toEqual(chain2Key);
+    });
+
+    it('should return undefined for non-existent keys', () => {
+      expect(keysStorage.getFheKey(999, 0)).toBeUndefined();
+      expect(keysStorage.getFheKey(testChainId, 999)).toBeUndefined();
+      expect(keysStorage.getFheKey(undefined, 0)).toBeUndefined();
+      expect(keysStorage.getFheKey(testChainId, undefined as any)).toBeUndefined();
+    });
+  });
+
+  describe('CRS Management', () => {
+    const testChainId = 1337;
+    const testCrs = '0x1234567890';
+
+    it('should set and get CRS', () => {
+      keysStorage.setCrs(testChainId, testCrs);
+
+      const retrievedCrs = keysStorage.getCrs(testChainId);
+
+      expect(retrievedCrs).toEqual(testCrs);
+    });
+
+    it('should handle multiple chains for CRS', () => {
+      const crs1 = '0x1234567890';
+      const crs2 = '0x2345678901';
+
+      keysStorage.setCrs(1, crs1);
+      keysStorage.setCrs(2, crs2);
+
+      expect(keysStorage.getCrs(1)).toEqual(crs1);
+      expect(keysStorage.getCrs(2)).toEqual(crs2);
+    });
+
+    it('should return undefined for non-existent CRS', () => {
+      expect(keysStorage.getCrs(999)).toBeUndefined();
+      expect(keysStorage.getCrs(undefined)).toBeUndefined();
+    });
+  });
+
+  describe('Storage Utilities', () => {
+    it('should clear keys storage', async () => {
+      await keysStorage.clearKeysStorage();
+
+      expect(mockStorage.removeItem).toHaveBeenCalledWith('cofhesdk-keys');
+    });
+
+    it('should rehydrate keys store', async () => {
+      const mockRehydrate = vi.fn();
+
+      // Mock the persist object if it doesn't exist
+      if (!('persist' in keysStorage.store)) {
+        (keysStorage.store as any).persist = {};
+      }
+      (keysStorage.store as any).persist.rehydrate = mockRehydrate;
+
+      await keysStorage.rehydrateKeysStore();
+
+      expect(mockRehydrate).toHaveBeenCalled();
+    });
+  });
+
+  describe('keysStorage Object', () => {
+    it('should have all required methods', () => {
+      expect(keysStorage).toBeDefined();
+      expect(keysStorage.store).toBeDefined();
+      expect(keysStorage.getFheKey).toBeDefined();
+      expect(keysStorage.getCrs).toBeDefined();
+      expect(keysStorage.setFheKey).toBeDefined();
+      expect(keysStorage.setCrs).toBeDefined();
+      expect(keysStorage.clearKeysStorage).toBeDefined();
+      expect(keysStorage.rehydrateKeysStore).toBeDefined();
+    });
+
+    it('should work through keysStorage object', () => {
+      const testChainId = 1337;
+      const testKey = '0x1234567890';
+      const testCrs = '0x2345678901';
+
+      keysStorage.setFheKey(testChainId, 0, testKey);
+      keysStorage.setCrs(testChainId, testCrs);
+
+      expect(keysStorage.getFheKey(testChainId, 0)).toEqual(testKey);
+      expect(keysStorage.getCrs(testChainId)).toEqual(testCrs);
+    });
+  });
+
+  describe('State Management', () => {
+    it('should update state immutably', () => {
+      const initialState = keysStorage.store.getState();
+      const testChainId = 1337;
+      const testKey = '0x1234567890';
+
+      keysStorage.setFheKey(testChainId, 0, testKey);
+
+      const newState = keysStorage.store.getState();
+
+      // State should be different objects
+      expect(newState).not.toBe(initialState);
+      expect(newState.fhe).not.toBe(initialState.fhe);
+
+      // But should contain the new key
+      expect(newState.fhe[testChainId][0]).toEqual(testKey);
+    });
+
+    it('should preserve existing data when adding new keys', () => {
+      const key1 = '0x1234567890';
+      const key2 = '0x2345678901';
+      const crs1 = '0x3456789012';
+
+      keysStorage.setFheKey(1, 0, key1);
+      keysStorage.setFheKey(2, 0, key2);
+      keysStorage.setCrs(1, crs1);
+
+      const state = keysStorage.store.getState();
+
+      expect(state.fhe[1][0]).toEqual(key1);
+      expect(state.fhe[2][0]).toEqual(key2);
+      expect(state.crs[1]).toEqual(crs1);
+    });
+  });
+
+  describe('Type Safety', () => {
+    it('should have correct TypeScript types', () => {
+      const state: KeysStore = keysStorage.store.getState();
+
+      // These should compile without TypeScript errors
+      expect(typeof state.fhe).toBe('object');
+      expect(typeof state.crs).toBe('object');
+
+      // Test that the types allow proper access patterns
+      const chainId = 1337;
+      const securityZone = 0;
+
+      if (state.fhe[chainId] && state.fhe[chainId][securityZone]) {
+        expect(state.fhe[chainId][securityZone]).toBeInstanceOf(String);
+      }
+
+      if (state.crs[chainId]) {
+        expect(state.crs[chainId]).toBeInstanceOf(String);
+      }
+    });
+  });
+});

package/core/keyStore.ts

@@ -0,0 +1,154 @@
+import { createStore, type StoreApi } from 'zustand/vanilla';
+import { persist, createJSONStorage } from 'zustand/middleware';
+import { produce } from 'immer';
+import { type IStorage } from './types.js';
+
+// Type definitions
+type ChainRecord<T> = Record<string, T>;
+type SecurityZoneRecord<T> = Record<number, T>;
+
+// Keys store for FHE keys and CRS
+export type KeysStore = {
+  fhe: ChainRecord<SecurityZoneRecord<string | undefined>>;
+  crs: ChainRecord<string | undefined>;
+};
+
+export type KeysStorage = {
+  store: StoreApi<KeysStore>;
+  getFheKey: (chainId: number | undefined, securityZone?: number) => string | undefined;
+  getCrs: (chainId: number | undefined) => string | undefined;
+  setFheKey: (chainId: number, securityZone: number, key: string) => void;
+  setCrs: (chainId: number, crs: string) => void;
+  clearKeysStorage: () => Promise<void>;
+  rehydrateKeysStore: () => Promise<void>;
+};
+
+function isValidPersistedState(state: unknown): state is KeysStore {
+  if (state && typeof state === 'object') {
+    if ('fhe' in state && 'crs' in state) {
+      return true;
+    } else {
+      throw new Error(
+        "Invalid persisted state structure for KeysStore. Is object but doesn't contain required fields 'fhe' and 'crs'."
+      );
+    }
+  }
+
+  return false;
+}
+
+const DEFAULT_KEYS_STORE: KeysStore = {
+  fhe: {},
+  crs: {},
+};
+
+type StoreWithPersist = ReturnType<typeof createStoreWithPersit>;
+
+function isStoreWithPersist(store: StoreApi<KeysStore> | StoreWithPersist): store is StoreWithPersist {
+  return 'persist' in store;
+}
+/**
+ * Creates a keys storage instance using the provided storage implementation
+ * @param storage - The storage implementation to use (IStorage interface), or null for non-persisted store
+ * @returns A KeysStorage instance with all utility methods
+ */
+export function createKeysStore(storage: IStorage | null): KeysStorage {
+  // Conditionally create store with or without persist wrapper
+  const keysStore = storage
+    ? createStoreWithPersit(storage)
+    : createStore<KeysStore>()(() => ({
+        fhe: {},
+        crs: {},
+      }));
+
+  // Utility functions
+
+  const getFheKey = (chainId: number | undefined, securityZone = 0) => {
+    if (chainId == null || securityZone == null) return undefined;
+    const stored = keysStore.getState().fhe[chainId]?.[securityZone];
+    return stored;
+  };
+
+  const getCrs = (chainId: number | undefined) => {
+    if (chainId == null) return undefined;
+    const stored = keysStore.getState().crs[chainId];
+    return stored;
+  };
+
+  const setFheKey = (chainId: number, securityZone: number, key: string) => {
+    keysStore.setState(
+      produce<KeysStore>((state: KeysStore) => {
+        if (state.fhe[chainId] == null) state.fhe[chainId] = {};
+        state.fhe[chainId][securityZone] = key;
+      })
+    );
+  };
+
+  const setCrs = (chainId: number, crs: string) => {
+    keysStore.setState(
+      produce<KeysStore>((state: KeysStore) => {
+        state.crs[chainId] = crs;
+      })
+    );
+  };
+
+  const clearKeysStorage = async () => {
+    if (storage) {
+      await storage.removeItem('cofhesdk-keys');
+    }
+    // If no storage, this is a no-op
+  };
+
+  const rehydrateKeysStore = async () => {
+    if (!isStoreWithPersist(keysStore)) return;
+    if (keysStore.persist.hasHydrated()) return;
+    await keysStore.persist.rehydrate();
+  };
+
+  return {
+    store: keysStore,
+    getFheKey,
+    getCrs,
+    setFheKey,
+    setCrs,
+    clearKeysStorage,
+    rehydrateKeysStore,
+  };
+}
+
+function createStoreWithPersit(storage: IStorage) {
+  const result = createStore<KeysStore>()(
+    persist(() => DEFAULT_KEYS_STORE, {
+      // because earleir tests were written with on-init hydration skipped (due to the error suppression in zustand), returning this flag to fix test (i.e. KeyStore > Storage Utilities > should rehydrate keys store)
+      skipHydration: true,
+      // if onRehydrateStorage is not passed here, the errors thrown by storage layer are swallowed by zustand here: https://github.com/pmndrs/zustand/blob/39a391b6c1ff9aa89b81694d9bdb21da37dd4ac6/src/middleware/persist.ts#L321
+      onRehydrateStorage: () => (_state?, _error?) => {
+        if (_error) throw new Error(`onRehydrateStorage: Error rehydrating keys store: ${_error}`);
+      },
+      name: 'cofhesdk-keys',
+      storage: createJSONStorage(() => storage),
+      merge: (persistedState, currentState) => {
+        const persisted = isValidPersistedState(persistedState) ? persistedState : DEFAULT_KEYS_STORE;
+        const current = currentState as KeysStore;
+
+        // Deep merge for fhe
+        const mergedFhe: KeysStore['fhe'] = { ...persisted.fhe };
+        const allChainIds = new Set([...Object.keys(current.fhe), ...Object.keys(persisted.fhe)]);
+        for (const chainId of allChainIds) {
+          const persistedZones = persisted.fhe[chainId] || {};
+          const currentZones = current.fhe[chainId] || {};
+          mergedFhe[chainId] = { ...persistedZones, ...currentZones };
+        }
+
+        // Deep merge for crs
+        const mergedCrs: KeysStore['crs'] = { ...persisted.crs, ...current.crs };
+
+        return {
+          fhe: mergedFhe,
+          crs: mergedCrs,
+        };
+      },
+    })
+  );
+  return result;
+}