feat() debug decryption

zama-ai · Jun 23, 2023 · 119489f · 119489f
1 parent 509545a
commit 119489f
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 41 deletions.
diff --git a/src/sdk/decrypt.ts b/src/sdk/decrypt.ts
@@ -1,11 +1,10 @@
 import sodium from 'libsodium-wrappers';
-import { bytesToNumber, fromHexString } from '../utils';
+import { bytesToNumber, fromHexString, toHexString } from '../utils';
 import { ContractKeypair } from './types';
 
 export const decrypt = async (keypair: ContractKeypair, ciphertext: string | Uint8Array): Promise<number> => {
   await sodium.ready;
   const toDecrypt = typeof ciphertext === 'string' ? fromHexString(ciphertext) : ciphertext;
-
   const decrypted = sodium.crypto_box_seal_open(toDecrypt, keypair.publicKey, keypair.privateKey);
   return bytesToNumber(decrypted);
 };
diff --git a/src/sdk/encrypt.test.ts b/src/sdk/encrypt.test.ts
@@ -5,13 +5,23 @@ import {
   CompactFheUint8List,
   CompactFheUint16List,
   CompactFheUint32List,
+  TfheCompactPublicKey,
+  TfheClientKey,
 } from 'node-tfhe';
-import { createTFHEKey } from '../tfhe';
+import { createTfheKeypair } from '../tfhe';
 import { encrypt8, encrypt16, encrypt32 } from './encrypt';
 
 describe('encrypt8', () => {
-  it('encrypt/decrypt', async () => {
-    const { publicKey, clientKey } = await createTFHEKey();
+  let clientKey: TfheClientKey;
+  let publicKey: TfheCompactPublicKey;
+
+  beforeAll(async () => {
+    const keypair = createTfheKeypair();
+    clientKey = keypair.clientKey;
+    publicKey = keypair.publicKey;
+  });
+
+  it('encrypt/decrypt 8bits', async () => {
     const buffer = encrypt8(34, publicKey);
     const compactList = CompactFheUint8List.deserialize(buffer);
     let encryptedList = compactList.expand();
@@ -20,11 +30,8 @@ describe('encrypt8', () => {
       expect(decrypted).toBe(34);
     });
   });
-});
 
-describe('encrypt16', () => {
-  it('encrypt/decrypt', async () => {
-    const { publicKey, clientKey } = await createTFHEKey();
+  it('encrypt/decrypt 16bits', async () => {
     const buffer = encrypt16(434, publicKey);
     const compactList = CompactFheUint16List.deserialize(buffer);
     let encryptedList = compactList.expand();
@@ -33,11 +40,8 @@ describe('encrypt16', () => {
       expect(decrypted).toBe(434);
     });
   });
-});
 
-describe('encrypt32', () => {
-  it('encrypt/decrypt', async () => {
-    const { publicKey, clientKey } = await createTFHEKey();
+  it('encrypt/decrypt 32bits', async () => {
     const buffer = encrypt32(30210, publicKey);
     const compactList = CompactFheUint32List.deserialize(buffer);
     let encryptedList = compactList.expand();

diff --git a/src/sdk/index.test.ts b/src/sdk/index.test.ts
@@ -1,18 +1,20 @@
 import { createInstance } from './index';
-import { createTFHEKey } from '../tfhe';
+import { createTfhePublicKey } from '../tfhe';
 import sodium from 'libsodium-wrappers';
 import { fromHexString, toHexString, numberToBytes } from '../utils';
 
 describe('token', () => {
+  let tfhePublicKey: string;
+
   beforeAll(async () => {
     await sodium.ready;
+    tfhePublicKey = createTfhePublicKey();
   });
 
   it('creates an instance', async () => {
-    const TFHEkeypair = createTFHEKey();
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
     });
     expect(instance.encrypt8).toBeDefined();
     expect(instance.encrypt16).toBeDefined();
@@ -25,14 +27,13 @@ describe('token', () => {
   });
 
   it('creates an instance with keypairs', async () => {
-    const TFHEkeypair = createTFHEKey();
     const keypair = sodium.crypto_box_keypair('hex');
 
     const contractAddress = '0x1c786b8ca49D932AFaDCEc00827352B503edf16c';
 
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
       keypairs: {
         [contractAddress]: {
           privateKey: keypair.privateKey,
@@ -50,10 +51,9 @@ describe('token', () => {
   });
 
   it('controls encrypt', async () => {
-    const TFHEkeypair = createTFHEKey();
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
     });
 
     expect(() => instance.encrypt8(undefined as any)).toThrow('Missing value');
@@ -66,20 +66,18 @@ describe('token', () => {
   });
 
   it('controls generateToken', async () => {
-    const TFHEkeypair = createTFHEKey();
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
     });
     await expect(instance.generateToken(undefined as any)).rejects.toThrow('Missing contract address');
     await expect(instance.generateToken({ verifyingContract: '' })).rejects.toThrow('Missing contract address');
   });
 
   it('save generated token', async () => {
-    const TFHEkeypair = createTFHEKey();
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
     });
 
     const contractAddress = '0x1c786b8ca49D932AFaDCEc00827352B503edf16c';
@@ -95,10 +93,9 @@ describe('token', () => {
   });
 
   it("don't export keys without signature", async () => {
-    const TFHEkeypair = createTFHEKey();
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
     });
 
     const contractAddress = '0x1c786b8ca49D932AFaDCEc00827352B503edf16c';
@@ -112,10 +109,9 @@ describe('token', () => {
   });
 
   it('decrypts data', async () => {
-    const TFHEkeypair = createTFHEKey();
     const instance = createInstance({
       chainId: 1234,
-      publicKey: TFHEkeypair.publicKey,
+      publicKey: tfhePublicKey,
     });
 
     const contractAddress = '0x1c786b8ca49D932AFaDCEc00827352B503edf16c';

diff --git a/src/sdk/index.ts b/src/sdk/index.ts
@@ -32,12 +32,14 @@ export type ExportedContractKeypairs = {
 
 export type FhevmInstanceParams = {
   chainId: number;
-  publicKey: TfheCompactPublicKey;
+  publicKey: string;
   keypairs?: ExportedContractKeypairs;
 };
 
 export const createInstance = (params: FhevmInstanceParams): FhevmInstance => {
   const { chainId, publicKey, keypairs } = params;
+  const buff = fromHexString(publicKey);
+  const tfheCompactPublicKey = TfheCompactPublicKey.deserialize(buff);
 
   let contractKeypairs: ContractKeypairs = {};
 
@@ -61,18 +63,18 @@ export const createInstance = (params: FhevmInstanceParams): FhevmInstance => {
     encrypt8(value) {
       if (!value) throw new Error('Missing value');
       if (typeof value !== 'number') throw new Error('Value must be a number');
-      return encrypt8(value, publicKey);
+      return encrypt8(value, tfheCompactPublicKey);
     },
     encrypt16(value) {
       if (!value) throw new Error('Missing value');
       if (typeof value !== 'number') throw new Error('Value must be a number');
-      return encrypt16(value, publicKey);
+      return encrypt16(value, tfheCompactPublicKey);
     },
 
     encrypt32(value) {
       if (!value) throw new Error('Missing value');
       if (typeof value !== 'number') throw new Error('Value must be a number');
-      return encrypt32(value, publicKey);
+      return encrypt32(value, tfheCompactPublicKey);
     },
 
     // Reencryption
@@ -123,7 +125,6 @@ export const createInstance = (params: FhevmInstanceParams): FhevmInstance => {
       if (!contractAddress) throw new Error('Missing contract address');
       const kp = contractKeypairs[contractAddress];
       if (!kp) throw new Error(`Missing keypair for ${contractAddress}`);
-
       return decrypt(kp, ciphertext);
     },
 

diff --git a/src/sdk/token.ts b/src/sdk/token.ts
@@ -1,4 +1,4 @@
-import sodium, { KeyPair } from 'libsodium-wrappers';
+import sodium from 'libsodium-wrappers';
 import { toHexString } from '../utils';
 import { ContractKeypair } from './types';
 
@@ -29,13 +29,13 @@ export type GenerateTokenParams = {
 };
 
 export type FhevmToken = {
-  keypair: KeyPair;
+  keypair: ContractKeypair;
   token: EIP712;
 };
 
 export const generateToken = async (params: GenerateTokenParams): Promise<FhevmToken> => {
   await sodium.ready;
-  const keypair = sodium.crypto_box_keypair();
+  const keypair = params.keypair || sodium.crypto_box_keypair();
   const msgParams: EIP712 = {
     types: {
       // This refers to the domain the contract is hosted on.
@@ -67,9 +67,11 @@ export const generateToken = async (params: GenerateTokenParams): Promise<FhevmT
       publicKey: `0x${toHexString(keypair.publicKey)}`,
     },
   };
-
   return {
-    keypair,
+    keypair: {
+      publicKey: keypair.publicKey,
+      privateKey: keypair.privateKey,
+    },
     token: msgParams,
   };
 };
diff --git a/src/tfhe.ts b/src/tfhe.ts
@@ -5,11 +5,18 @@ import {
   ShortintParameters,
   ShortintParametersName,
 } from 'node-tfhe';
+import { fromHexString, toHexString } from './utils';
 
-export const createTFHEKey = () => {
+export const createTfheKeypair = () => {
   const block_params = new ShortintParameters(ShortintParametersName.PARAM_SMALL_MESSAGE_2_CARRY_2_COMPACT_PK);
   let config = TfheConfigBuilder.all_disabled().enable_custom_integers(block_params).build();
   let clientKey = TfheClientKey.generate(config);
   let publicKey = TfheCompactPublicKey.new(clientKey);
-  return { publicKey, clientKey };
+  publicKey = TfheCompactPublicKey.deserialize(publicKey.serialize());
+  return { clientKey, publicKey };
+};
+
+export const createTfhePublicKey = (): string => {
+  const { publicKey } = createTfheKeypair();
+  return toHexString(publicKey.serialize());
 };
diff --git a/src/utils.ts b/src/utils.ts
@@ -1,7 +1,7 @@
 import sha3 from 'sha3';
 
 export const fromHexString = (hexString: string): Uint8Array => {
-  const arr = hexString.match(/.{1,2}/g);
+  const arr = hexString.replace(/^(0x)/, '').match(/.{1,2}/g);
   if (!arr) return new Uint8Array();
   return Uint8Array.from(arr.map((byte) => parseInt(byte, 16)));
 };