fix: propagate publicKeyId to getKeysFromGateway (#135)

* fix: propagate publicKeyId to getKeysFromGateway
zama-ai · Nov 18, 2024 · a223bd9 · a223bd9
1 parent 71d95f8
commit a223bd9
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/src/sdk/config.ts b/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);

diff --git a/src/sdk/network.ts b/src/sdk/network.ts
@@ -35,7 +35,7 @@ export type GatewayKeys = {
 };
 
 const keyurlCache: { [key: string]: any } = {};
-export const getKeysFromGateway = async (url: string) => {
+export const getKeysFromGateway = async (url: string, publicKeyId?: string) => {
   if (keyurlCache[url]) {
     return keyurlCache[url];
   }
@@ -46,8 +46,26 @@ export const getKeysFromGateway = async (url: string) => {
     }
     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];
+      }
+
       const publicKeyResponse = await fetch(pubKeyUrl);
       const publicKey = await publicKeyResponse.arrayBuffer();
       const publicParamsUrl = data.response.crs['2048'].urls[0];