Use crs/key ID in createEncryptedInput

zama-ai · Nov 4, 2024 · b13ea79 · b13ea79
1 parent 7e29f5b
commit b13ea79
Show file tree

Hide file tree

Showing 16 changed files with 87 additions and 61 deletions.
diff --git a/src/sdk/config.ts b/src/sdk/config.ts
@@ -7,7 +7,7 @@ import {
 } from 'ethers';
 import { PublicParams } from './encrypt';
 import { getKeysFromGateway } from './network';
-import { fromHexString, cleanURL } from '../utils';
+import { fromHexString, cleanURL, SERIALIZED_SIZE_LIMIT_PK, SERIALIZED_SIZE_LIMIT_CRS } from '../utils';
 import { CompactPkePublicParams, TfheCompactPublicKey } from 'node-tfhe';
 import { abi } from '../abi/kmsVerifier.json';
 
@@ -16,6 +16,7 @@ export type FhevmInstanceConfig = {
   aclContractAddress: string;
   chainId?: number;
   publicKey?: string;
+  publicKeyId?: string;
   gatewayUrl?: string;
   network?: Eip1193Provider;
   networkUrl?: string;
@@ -50,16 +51,16 @@ export const getChainId = async (
 export const getTfheCompactPublicKey = async (config: FhevmInstanceConfig) => {
   if (config.gatewayUrl && !config.publicKey) {
     const inputs = await getKeysFromGateway(cleanURL(config.gatewayUrl));
-    return inputs.publicKey;
-  } else if (config.publicKey) {
+    return { publicKey: inputs.publicKey, publicKeyId: inputs.publicKeyId };
+  } else if (config.publicKey && config.publicKeyId) {
     const buff = fromHexString(config.publicKey);
     try {
-      return TfheCompactPublicKey.deserialize(buff);
+      return { publicKey: TfheCompactPublicKey.safe_deserialize(buff, SERIALIZED_SIZE_LIMIT_PK), publicKeyId: config.publicKeyId };
     } catch (e) {
       throw new Error('Invalid public key (deserialization failed)');
     }
   } else {
-    throw new Error('You must provide a public key.');
+    throw new Error('You must provide a public key with its public key ID.');
   }
 };
 
@@ -68,16 +69,19 @@ export const getPublicParams = async (config: FhevmInstanceConfig) => {
     const inputs = await getKeysFromGateway(cleanURL(config.gatewayUrl));
     return inputs.publicParams;
   } else if (config.publicParams && config.publicParams['2048']) {
-    const buff = fromHexString(config.publicParams['2048']);
+    const buff = fromHexString(config.publicParams['2048'].publicParams);
     try {
       return {
-        '2048': CompactPkePublicParams.deserialize(buff),
+        2048: { 
+          publicParams: CompactPkePublicParams.safe_deserialize(buff, SERIALIZED_SIZE_LIMIT_CRS),
+          publicParamsId: config.publicParams['2048'].publicParamsId
+        },
       };
     } catch (e) {
       throw new Error('Invalid public key (deserialization failed)');
     }
   } else {
-    throw new Error('You must provide a valid CRS.');
+    throw new Error('You must provide a valid CRS with its CRS ID.');
   }
 };
 

diff --git a/src/sdk/encrypt.test.ts b/src/sdk/encrypt.test.ts
@@ -1,14 +1,14 @@
 import { ProvenCompactCiphertextList } from 'node-tfhe';
 import { createEncryptedInput } from './encrypt';
-import { publicKey, publicParams } from '../test';
+import { publicKey, publicKeyId, publicParams } from '../test';
 import fetchMock from '@fetch-mock/core';
 
-fetchMock.post('https://test-gateway.net/zkp', {
+fetchMock.post('https://test-gateway.net/verify_proven_ct', {
   response: {
-    coprocessor: 'COPROCESSOR',
-    kms_signatures: ['0x32'],
-    coproc_signature: '0x54',
-    handles: ['0x2323', '0x2234'],
+    proof_of_storage: 'deadbeef',
+    handles: ['2323beef', '2234beef'],
+    kms_signatures: ['dead3232'],
+    coproc_signature: '1122beef',
   },
   status: 'success',
 });
@@ -20,6 +20,7 @@ describe('encrypt', () => {
       1234,
       'https://test-gateway.net/',
       publicKey,
+      publicKeyId,
       publicParams,
     )(
       '0x8ba1f109551bd432803012645ac136ddd64dba72',
@@ -48,14 +49,15 @@ describe('encrypt', () => {
     // types.forEach((val, i) => {
     //   expect(val).toBe(expectedTypes[i]);
     // });
-  });
+  }, 60000);
 
   it('encrypt/decrypt one 0 value', async () => {
     const input = createEncryptedInput(
       '0x325ea1b59F28e9e1C51d3B5b47b7D3965CC5D8C8',
       1234,
       'https://test-gateway.net/',
       publicKey,
+      publicKeyId,
       publicParams,
     )(
       '0x8ba1f109551bd432803012645ac136ddd64dba72',
@@ -83,6 +85,7 @@ describe('encrypt', () => {
       1234,
       'https://test-gateway.net/',
       publicKey,
+      publicKeyId,
       publicParams,
     )(
       '0x8ba1f109551bd432803012645ac136ddd64dba72',
@@ -113,6 +116,7 @@ describe('encrypt', () => {
         1234,
         'https://test-gateway.net/',
         publicKey,
+        publicKeyId,
         publicParams,
       )('0xa5e1defb98EFe38EBb2D958CEe052410247F4c80', '0'),
     ).toThrow('User address is not a valid address.');
@@ -122,6 +126,7 @@ describe('encrypt', () => {
         1234,
         'https://test-gateway.net/',
         publicKey,
+        publicKeyId,
         publicParams,
       )('0x0', '0xa5e1defb98EFe38EBb2D958CEe052410247F4c80'),
     ).toThrow('Contract address is not a valid address.');
@@ -132,6 +137,7 @@ describe('encrypt', () => {
         1234,
         'https://test-gateway.net/',
         publicKey,
+        publicKeyId,
         publicParams,
       )(
         '0x8ba1f109551bd432803012645ac136ddd64dba72',
@@ -144,6 +150,7 @@ describe('encrypt', () => {
       1234,
       'https://test-gateway.net/',
       publicKey,
+      publicKeyId,
       publicParams,
     )(
       '0x8ba1f109551bd432803012645ac136ddd64dba72',
@@ -190,6 +197,7 @@ describe('encrypt', () => {
       1234,
       'https://test-gateway.net/',
       publicKey,
+      publicKeyId,
       publicParams,
     )(
       '0x8ba1f109551bd432803012645ac136ddd64dba72',

diff --git a/src/sdk/encrypt.ts b/src/sdk/encrypt.ts
@@ -68,7 +68,7 @@ const checkEncryptedValue = (value: number | bigint, bits: number) => {
 };
 
 export type PublicParams<T = CompactPkePublicParams> = {
-  [key in EncryptionTypes]?: T;
+  [key in EncryptionTypes]?: {publicParams: T, publicParamsId: string};
 };
 
 export const createEncryptedInput =
@@ -77,6 +77,7 @@ export const createEncryptedInput =
     chainId: number,
     gateway: string,
     tfheCompactPublicKey: TfheCompactPublicKey,
+    publicKeyId: string,
     publicParams: PublicParams,
   ) =>
     (contractAddress: string, callerAddress: string): ZKInput => {
@@ -235,7 +236,8 @@ export const createEncryptedInput =
               }.`,
             );
           }
-          const pp = publicParams[closestPP]!;
+          const pp = publicParams[closestPP]!.publicParams;
+          const ppId = publicParams[closestPP]!.publicParamsId;
           const buffContract = fromHexString(contractAddress);
           const buffUser = fromHexString(callerAddress);
           const buffAcl = fromHexString(aclContractAddress);
@@ -260,15 +262,12 @@ export const createEncryptedInput =
             SERIALIZED_SIZE_LIMIT_CIPHERTEXT,
           ));
 
-          // TODO get keyID
-          const keyID = '408d8cbaa51dece7f782fe04ba0b1c1d017b1088';
-          const crsID = 'd8d94eb3a23d22d3eb6b5e7b694e8afcd571d906';
           const payload = {
             contract_address: contractAddress,
             caller_address: callerAddress,
             ct_proof: ciphertext.toString('hex'),
-            key_id: keyID,
-            crs_id: crsID,
+            key_id: publicKeyId,
+            crs_id: ppId,
           };
 
           const options = {

diff --git a/src/sdk/index.test.ts b/src/sdk/index.test.ts
@@ -1,6 +1,6 @@
 import { createInstance } from './index';
-import { publicKey, publicParams } from '../test';
-import { bytesToHex } from '../utils';
+import { publicKey, publicKeyId, publicParams } from '../test';
+import { bytesToHex, SERIALIZED_SIZE_LIMIT_CRS, SERIALIZED_SIZE_LIMIT_PK } from '../utils';
 
 jest.mock('ethers', () => ({
   JsonRpcProvider: () => ({
@@ -14,16 +14,18 @@ jest.mock('ethers', () => ({
 
 describe('index', () => {
   it('creates an instance', async () => {
-    const serializedPublicKey = bytesToHex(publicKey.serialize());
+    const serializedPublicKey = bytesToHex(publicKey.safe_serialize(SERIALIZED_SIZE_LIMIT_PK));
     const serializedPublicParams = bytesToHex(
-      publicParams[2048].serialize(false),
+      publicParams[2048].publicParams.safe_serialize(SERIALIZED_SIZE_LIMIT_CRS),
     );
+    const publicParamsId = publicParams[2048].publicParamsId;
     const instance = await createInstance({
       aclContractAddress: '0x4c102C7cA99d3079fEFF08114d3bad888b9794d9',
       kmsContractAddress: '0x325ea1b59F28e9e1C51d3B5b47b7D3965CC5D8C8',
       chainId: 1234,
       publicKey: serializedPublicKey,
-      publicParams: { 2048: serializedPublicParams },
+      publicKeyId,
+      publicParams: { 2048: { publicParams: serializedPublicParams, publicParamsId, } },
       networkUrl: 'https://network.com/',
     });
     expect(instance.reencrypt).toBeDefined();
@@ -36,15 +38,17 @@ describe('index', () => {
   it('fails to create an instance', async () => {
     const serializedPublicKey = bytesToHex(publicKey.serialize());
     const serializedPublicParams = bytesToHex(
-      publicParams[2048].serialize(false),
+      publicParams[2048].publicParams.safe_serialize(SERIALIZED_SIZE_LIMIT_CRS),
     );
+    const publicParamsId = publicParams[2048].publicParamsId;
     await expect(
       createInstance({
         aclContractAddress: '0x4c102C7cA99d3079fEFF08114d3bad888b9794d9',
         kmsContractAddress: '0x325ea1b59F28e9e1C51d3B5b47b7D3965CC5D8C8',
         chainId: BigInt(1234) as any,
         publicKey: serializedPublicKey,
-        publicParams: { 2048: serializedPublicParams },
+        publicKeyId,
+        publicParams: { 2048: { publicParams: serializedPublicParams, publicParamsId, } },
         networkUrl: 'https://',
       }),
     ).rejects.toThrow('chainId must be a number');
@@ -55,6 +59,7 @@ describe('index', () => {
         kmsContractAddress: '0x325ea1b59F28e9e1C51d3B5b47b7D3965CC5D8C8',
         chainId: 9000,
         publicKey: 43 as any,
+        publicKeyId,
       }),
     ).rejects.toThrow('publicKey must be a string');
   });

diff --git a/src/sdk/index.ts b/src/sdk/index.ts
@@ -11,7 +11,7 @@ import { PublicParams, ZKInput } from './encrypt';
 import { createEncryptedInput } from './encrypt';
 import { generateKeypair, createEIP712, EIP712 } from './keypair';
 import { reencryptRequest } from './reencrypt';
-import { isAddress, ethers } from 'ethers';
+import { isAddress } from 'ethers';
 
 export type FhevmInstance = {
   createEncryptedInput: (
@@ -62,9 +62,9 @@ export const createInstance = async (
 
   const chainId = await getChainId(provider, config);
 
-  const tfheCompactPublicKey = await getTfheCompactPublicKey(config);
+  const publicKeyData = await getTfheCompactPublicKey(config);
 
-  const pkePublicParams: PublicParams = await getPublicParams(config);
+  const publicParamsData = await getPublicParams(config);
 
   const kmsSigners = await getKMSSigners(provider, config);
 
@@ -73,8 +73,9 @@ export const createInstance = async (
       aclContractAddress,
       chainId,
       cleanURL(config.gatewayUrl),
-      tfheCompactPublicKey,
-      pkePublicParams,
+      publicKeyData.publicKey,
+      publicKeyData.publicKeyId,
+      publicParamsData,
     ),
     generateKeypair,
     createEIP712: createEIP712(chainId),
@@ -87,6 +88,6 @@ export const createInstance = async (
       provider,
     ),
     getPublicKey: () => publicKey || null,
-    getPublicParams: () => pkePublicParams || null,
+    getPublicParams: () => publicParamsData || null,
   };
 };
diff --git a/src/sdk/network.test.ts b/src/sdk/network.test.ts
@@ -1,6 +1,6 @@
 import { GatewayKeys, getKeysFromGateway } from './network';
 import { publicKey, publicParams } from '../test';
-import { bytesToHex } from '../utils';
+import { bytesToHex, SERIALIZED_SIZE_LIMIT_CRS, SERIALIZED_SIZE_LIMIT_PK } from '../utils';
 import fetchMock from '@fetch-mock/core';
 
 const payload: GatewayKeys = {
@@ -86,18 +86,18 @@ fetchMock.get('https://test-gateway.net/keyurl', payload);
 
 fetchMock.get(
   'https://s3.amazonaws.com/bucket-name-1/PUB-p1/PublicKey/408d8cbaa51dece7f782fe04ba0b1c1d017b1088',
-  bytesToHex(publicKey.serialize()),
+  publicKey.safe_serialize(SERIALIZED_SIZE_LIMIT_PK),
 );
+
 fetchMock.get(
   'https://s3.amazonaws.com/bucket-name-1/PUB-p1/CRS/d8d94eb3a23d22d3eb6b5e7b694e8afcd571d906',
-  bytesToHex(publicParams[2048].serialize(false)),
+  publicParams[2048].publicParams.safe_serialize(SERIALIZED_SIZE_LIMIT_CRS),
 );
 
 describe('network', () => {
-  // TODO: fix this test by returning valid safe serialized keys
   it.skip('getInputsFromGateway', async () => {
     const material = await getKeysFromGateway('https://test-gateway.net/');
 
-    expect(material.publicKey.serialize()).toStrictEqual(publicKey.serialize());
+    expect(material.publicKey.safe_serialize(SERIALIZED_SIZE_LIMIT_PK)).toStrictEqual(publicKey.safe_serialize(SERIALIZED_SIZE_LIMIT_PK));
   });
 });
diff --git a/src/sdk/network.ts b/src/sdk/network.ts
@@ -1,4 +1,5 @@
 import { CompactPkePublicParams, TfheCompactPublicKey } from 'node-tfhe';
+import { SERIALIZED_SIZE_LIMIT_PK, SERIALIZED_SIZE_LIMIT_CRS } from '../utils';
 
 export type GatewayKeysItem = {
   data_id: string;
@@ -33,7 +34,11 @@ export type GatewayKeys = {
   status: string;
 };
 
+const keyurlCache: { [key: string]: any } = {};
 export const getKeysFromGateway = async (url: string) => {
+  if (keyurlCache[url]) {
+    return keyurlCache[url];
+  }
   try {
     const response = await fetch(`${url}keyurl`);
     if (!response.ok) {
@@ -42,23 +47,32 @@ 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;
       console.log('pubKeyUrl', pubKeyUrl);
       const publicKeyResponse = await fetch(pubKeyUrl);
       const publicKey = await publicKeyResponse.arrayBuffer();
-      const crsUrl = data.response.crs['2048'].urls[0];
-      const crs2048 = await (await fetch(crsUrl)).arrayBuffer();
-      return {
+      const publicParamsUrl = data.response.crs['2048'].urls[0];
+      const publicParamsId = data.response.crs['2048'].data_id;
+      const publicParams2048 = await (await fetch(publicParamsUrl)).arrayBuffer();
+
+      const result = {
         publicKey: TfheCompactPublicKey.safe_deserialize(
           new Uint8Array(publicKey),
-          BigInt(1024) * BigInt(1024) * BigInt(16),
+          SERIALIZED_SIZE_LIMIT_PK,
         ),
+        publicKeyId,
         publicParams: {
-          2048: CompactPkePublicParams.safe_deserialize(
-            new Uint8Array(crs2048),
-            BigInt(1024) * BigInt(1024) * BigInt(512),
-          ),
-        },
+          2048: {
+            publicParams: CompactPkePublicParams.safe_deserialize(
+              new Uint8Array(publicParams2048),
+              SERIALIZED_SIZE_LIMIT_CRS,
+            ),
+            publicParamsId,
+          }
+        }
       };
+      keyurlCache[url] = result;
+      return result;
     } else {
       throw new Error('No public key available');
     }