fix(sdk): cache the CRS per security zone

.changeset/crs-per-security-zone.md

@@ -0,0 +1,11 @@
+---
+'@cofhe/sdk': minor
+---
+
+fix(sdk): cache the CRS per security zone
+
+`fetchCrs` requested the CRS for a given `securityZone` but cached it under the chain id alone, while `fetchFhePublicKey` next to it keys by `(chainId, securityZone)`. After encrypting on one zone, `encryptInputs` on a different zone read the first zone's CRS back out of the store. It deserializes fine, so the stale CRS passed the validity check and was used to build the proof, which the ZK verifier then rejects.
+
+CoFHE serves a distinct CRS per zone: `POST /GetCrs` with `securityZone: 0` and `securityZone: 1` against `testnet-cofhe.fhenix.zone` returns two different 7,968,512-byte values.
+
+`KeysStore.crs` is now keyed by chain and zone like `KeysStore.fhe`, `getCrs(chainId, securityZone = 0)` takes the zone, and `setCrs(chainId, securityZone, crs)` matches `setFheKey`'s argument order. A `crs` persisted under the old shape holds a bare string per chain; those entries are dropped on rehydrate so the CRS is refetched per zone instead of being indexed as a string.

packages/sdk/core/fetchKeys.ts

@@ -77,7 +77,7 @@ const fetchCrs = async (
   keysStorage?: KeysStorage | null
 ): Promise<[string, boolean]> => {
   // Escape if key already exists
-  const storedKey = keysStorage?.getCrs(chainId);
+  const storedKey = keysStorage?.getCrs(chainId, securityZone);
   const [storedKeyValid] = checkKeyValidity(storedKey, compactPkeCrsDeserializer);
   if (storedKeyValid) return [storedKey!, false];
 
@@ -108,7 +108,7 @@ const fetchCrs = async (
     throw new Error(`Error serializing CRS; ${err}`);
   }
 
-  keysStorage?.setCrs(chainId, crs_data);
+  keysStorage?.setCrs(chainId, securityZone, crs_data);
 
   return [crs_data, true];
 };

packages/sdk/core/keyStore.ts

@@ -11,15 +11,15 @@ 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>;
+  crs: ChainRecord<SecurityZoneRecord<string | undefined>>;
 };
 
 export type KeysStorage = {
   store: StoreApi<KeysStore>;
   getFheKey: (chainId: number | undefined, securityZone?: number) => string | undefined;
-  getCrs: (chainId: number | undefined) => string | undefined;
+  getCrs: (chainId: number | undefined, securityZone?: number) => string | undefined;
   setFheKey: (chainId: number, securityZone: number, key: string) => void;
-  setCrs: (chainId: number, crs: string) => void;
+  setCrs: (chainId: number, securityZone: number, crs: string) => void;
   clearKeysStorage: () => Promise<void>;
   rehydrateKeysStore: () => Promise<void>;
 };
@@ -45,6 +45,10 @@ const DEFAULT_KEYS_STORE: KeysStore = {
   crs: {},
 };
 
+function isSecurityZoneRecord(value: unknown): value is SecurityZoneRecord<string | undefined> {
+  return value != null && typeof value === 'object';
+}
+
 type StoreWithPersist = ReturnType<typeof createStoreWithPersit>;
 
 function isStoreWithPersist(store: StoreApi<KeysStore> | StoreWithPersist): store is StoreWithPersist {
@@ -72,9 +76,9 @@ export function createKeysStore(storage: IStorage | null): KeysStorage {
     return stored;
   };
 
-  const getCrs = (chainId: number | undefined) => {
-    if (chainId == null) return undefined;
-    const stored = keysStore.getState().crs[chainId];
+  const getCrs = (chainId: number | undefined, securityZone = 0) => {
+    if (chainId == null || securityZone == null) return undefined;
+    const stored = keysStore.getState().crs[chainId]?.[securityZone];
     return stored;
   };
 
@@ -87,10 +91,11 @@ export function createKeysStore(storage: IStorage | null): KeysStorage {
     );
   };
 
-  const setCrs = (chainId: number, crs: string) => {
+  const setCrs = (chainId: number, securityZone: number, crs: string) => {
     keysStore.setState(
       produce<KeysStore>((state: KeysStore) => {
-        state.crs[chainId] = crs;
+        if (state.crs[chainId] == null) state.crs[chainId] = {};
+        state.crs[chainId][securityZone] = crs;
       })
     );
   };
@@ -143,8 +148,16 @@ function createStoreWithPersit(storage: IStorage) {
           mergedFhe[chainId] = { ...persistedZones, ...currentZones };
         }
 
-        // Deep merge for crs
-        const mergedCrs: KeysStore['crs'] = { ...persisted.crs, ...current.crs };
+        // Deep merge for crs. A `crs` persisted before it was keyed by security zone holds a
+        // bare string per chain; those entries are dropped so the CRS is refetched per zone
+        // rather than indexed as a string.
+        const mergedCrs: KeysStore['crs'] = {};
+        const allCrsChainIds = new Set([...Object.keys(current.crs), ...Object.keys(persisted.crs)]);
+        for (const chainId of allCrsChainIds) {
+          const persistedZones = isSecurityZoneRecord(persisted.crs[chainId]) ? persisted.crs[chainId] : {};
+          const currentZones = isSecurityZoneRecord(current.crs[chainId]) ? current.crs[chainId] : {};
+          mergedCrs[chainId] = { ...persistedZones, ...currentZones };
+        }
 
         return {
           fhe: mergedFhe,

packages/sdk/core/test/encryptInputsBuilder.test.ts

@@ -175,7 +175,7 @@ let keysStorage: KeysStorage;
 
 const insertMockKeys = (chainId: number, securityZone: number) => {
   keysStorage.setFheKey(chainId, securityZone, '0x1234567890');
-  keysStorage.setCrs(chainId, '0x1234567890');
+  keysStorage.setCrs(chainId, securityZone, '0x1234567890');
 };
 
 const mockTfhePublicKeyDeserializer: FheKeyDeserializer = (buff: string) => {

packages/sdk/core/test/fetchKeys.test.ts

@@ -110,7 +110,7 @@ describe('fetchKeys', () => {
   it('should not fetch CRS if already cached', async () => {
     // Pre-populate with a cached CRS
     const mockCachedCrs = '0x2345678901';
-    keysStorage.setCrs(sepolia.id, mockCachedCrs);
+    keysStorage.setCrs(sepolia.id, 0, mockCachedCrs);
 
     const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
       config,
@@ -133,12 +133,33 @@ describe('fetchKeys', () => {
     expect(retrievedKey).toBeDefined();
   });
 
+  it('should not reuse a cached CRS from a different security zone', async () => {
+    // CoFHE serves a distinct CRS per security zone, so a zone 0 CRS must not satisfy zone 1
+    const zone0Crs = '0x2345678901';
+    keysStorage.setCrs(sepolia.id, 0, zone0Crs);
+
+    const [, [crs, crsFetchedFromCoFHE]] = await fetchKeys(
+      config,
+      sepolia.id,
+      1,
+      mockTfhePublicKeyDeserializer,
+      mockCompactPkeCrsDeserializer,
+      keysStorage
+    );
+
+    expect(crsFetchedFromCoFHE).toBe(true);
+    expect(crs).not.toEqual(zone0Crs);
+
+    // The zone 0 entry is untouched
+    expect(keysStorage.getCrs(sepolia.id, 0)).toEqual(zone0Crs);
+  });
+
   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);
+    keysStorage.setCrs(sepolia.id, 0, mockCachedCrs);
 
     const [[fheKey, fheKeyFetchedFromCoFHE], [crs, crsFetchedFromCoFHE]] = await fetchKeys(
       config,

packages/sdk/core/test/keyStore.test.ts

@@ -98,7 +98,7 @@ describe('KeyStore', () => {
     const testCrs = '0x1234567890';
 
     it('should set and get CRS', () => {
-      keysStorage.setCrs(testChainId, testCrs);
+      keysStorage.setCrs(testChainId, 0, testCrs);
 
       const retrievedCrs = keysStorage.getCrs(testChainId);
 
@@ -109,16 +109,28 @@ describe('KeyStore', () => {
       const crs1 = '0x1234567890';
       const crs2 = '0x2345678901';
 
-      keysStorage.setCrs(1, crs1);
-      keysStorage.setCrs(2, crs2);
+      keysStorage.setCrs(1, 0, crs1);
+      keysStorage.setCrs(2, 0, crs2);
 
       expect(keysStorage.getCrs(1)).toEqual(crs1);
       expect(keysStorage.getCrs(2)).toEqual(crs2);
     });
 
+    it('should keep CRS separate per security zone', () => {
+      const crsZone0 = '0xaaaaaaaaaa';
+      const crsZone1 = '0xbbbbbbbbbb';
+
+      keysStorage.setCrs(testChainId, 0, crsZone0);
+      keysStorage.setCrs(testChainId, 1, crsZone1);
+
+      expect(keysStorage.getCrs(testChainId, 0)).toEqual(crsZone0);
+      expect(keysStorage.getCrs(testChainId, 1)).toEqual(crsZone1);
+    });
+
     it('should return undefined for non-existent CRS', () => {
       expect(keysStorage.getCrs(999)).toBeUndefined();
       expect(keysStorage.getCrs(undefined)).toBeUndefined();
+      expect(keysStorage.getCrs(testChainId, 999)).toBeUndefined();
     });
   });
 
@@ -162,7 +174,7 @@ describe('KeyStore', () => {
       const testCrs = '0x2345678901';
 
       keysStorage.setFheKey(testChainId, 0, testKey);
-      keysStorage.setCrs(testChainId, testCrs);
+      keysStorage.setCrs(testChainId, 0, testCrs);
 
       expect(keysStorage.getFheKey(testChainId, 0)).toEqual(testKey);
       expect(keysStorage.getCrs(testChainId)).toEqual(testCrs);
@@ -194,13 +206,13 @@ describe('KeyStore', () => {
 
       keysStorage.setFheKey(1, 0, key1);
       keysStorage.setFheKey(2, 0, key2);
-      keysStorage.setCrs(1, crs1);
+      keysStorage.setCrs(1, 0, 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);
+      expect(state.crs[1][0]).toEqual(crs1);
     });
   });