fix: propagate publicKeyId to getKeysFromGateway

src/sdk/config.ts

@@ -50,7 +50,7 @@ export const getChainId = async (
 
 export const getTfheCompactPublicKey = async (config: FhevmInstanceConfig) => {
   if (config.gatewayUrl && !config.publicKey) {
-    const inputs = await getKeysFromGateway(cleanURL(config.gatewayUrl));
+    const inputs = await getKeysFromGateway(cleanURL(config.gatewayUrl), config.publicKeyId);
     return { publicKey: inputs.publicKey, publicKeyId: inputs.publicKeyId };
   } else if (config.publicKey && config.publicKeyId) {
     const buff = fromHexString(config.publicKey);
@@ -66,7 +66,7 @@ export const getTfheCompactPublicKey = async (config: FhevmInstanceConfig) => {
 
 export const getPublicParams = async (config: FhevmInstanceConfig) => {
   if (config.gatewayUrl && !config.publicParams) {
-    const inputs = await getKeysFromGateway(cleanURL(config.gatewayUrl));
+    const inputs = await getKeysFromGateway(cleanURL(config.gatewayUrl), config.publicKeyId);
     return inputs.publicParams;
   } else if (config.publicParams && config.publicParams['2048']) {
     const buff = fromHexString(config.publicParams['2048'].publicParams);

src/sdk/network.test.ts

@@ -99,7 +99,7 @@ fetchMock.get(
 
 describe('network', () => {
   it('getInputsFromGateway', async () => {
-    const material = await getKeysFromGateway('https://test-gateway.net/');
+    const material = await getKeysFromGateway('https://test-gateway.net/', undefined);
 
     expect(
       material.publicKey.safe_serialize(SERIALIZED_SIZE_LIMIT_PK),

src/sdk/network.ts

@@ -35,19 +35,41 @@ export type GatewayKeys = {
 };
 
 const keyurlCache: { [key: string]: any } = {};
-export const getKeysFromGateway = async (url: string) => {
+export const getKeysFromGateway = async (url: string, publicKeyId: string | undefined) => {
   if (keyurlCache[url]) {
     return keyurlCache[url];
   }
   try {
+    console.log('fetching keys from', url);
     const response = await fetch(`${url}keyurl`);
     if (!response.ok) {
       throw new Error(`HTTP error! status: ${response.status}`);
     }
     const data: GatewayKeys = await response.json();
     if (data) {
-      const pubKeyUrl = data.response.fhe_key_info[0].fhe_public_key.urls[0];
-      const publicKeyId = data.response.fhe_key_info[0].fhe_public_key.data_id;
+      let pubKeyUrl: string;
+
+      // If no publicKeyId is provided, use the first one
+      // Warning: if there are multiple keys available, the first one will most likely never be the
+      // same between several calls (fetching the infos is non-deterministic)
+      if (!publicKeyId) {
+        pubKeyUrl = data.response.fhe_key_info[0].fhe_public_key.urls[0];
+        publicKeyId = data.response.fhe_key_info[0].fhe_public_key.data_id;
+      } else {
+        // If a publicKeyId is provided, get the corresponding info
+        const keyInfo = data.response.fhe_key_info.find(info => info.fhe_public_key.data_id === publicKeyId);
+
+        if (!keyInfo) {
+          throw new Error(`Could not find FHE key info with data_id ${publicKeyId}`);
+        }
+
+        // TODO: Get a given party's public key url instead of the first one
+        pubKeyUrl = keyInfo.fhe_public_key.urls[0];
+      }
+
+      console.log('publicKeyId', publicKeyId);
+      console.log('pubKeyUrl', pubKeyUrl);
+
       const publicKeyResponse = await fetch(pubKeyUrl);
       const publicKey = await publicKeyResponse.arrayBuffer();
       const publicParamsUrl = data.response.crs['2048'].urls[0];