Merge pull request #53 from zama-ai/feature/encrypt4

Feature/encrypt4

bin/fhevm.js

@@ -6,6 +6,8 @@ import { program } from 'commander';
 import { toHexString, prependHttps, throwError } from './utils.js';
 import { createInstance } from '../lib/node.cjs';
 
+const allowedBits = [8, 16, 32, 64];
+
 const FHE_LIB_ADDRESS = '0x000000000000000000000000000000000000005d';
 
 let _instance;
@@ -20,23 +22,29 @@ const getInstance = async (provider) => {
     throwError('Network is unreachable');
   }
   const chainId = +network.chainId.toString();
-  const ret = await provider.call({
-    to: FHE_LIB_ADDRESS,
-    // first four bytes of keccak256('fhePubKey(bytes1)') + 1 byte for library
-    data: '0xd9d47bb001',
-  });
-  const decoded = AbiCoder.defaultAbiCoder().decode(['bytes'], ret);
-  const publicKey = decoded[0];
-  _instance = await createInstance({ chainId, publicKey });
-
+  try {
+    const ret = await provider.call({
+      to: FHE_LIB_ADDRESS,
+      // first four bytes of keccak256('fhePubKey(bytes1)') + 1 byte for library
+      data: '0xd9d47bb001',
+    });
+    const decoded = AbiCoder.defaultAbiCoder().decode(['bytes'], ret);
+    const publicKey = decoded[0];
+    _instance = await createInstance({ chainId, publicKey });
+  } catch (e) {
+    return throwError(
+      "This network doesn't seem to use fhEVM or use an incompatible version.",
+    );
+  }
   return _instance;
 };
 
 program
   .command('encrypt')
-  .argument('<bits>', 'number of bits (8, 16, 32)')
+  .argument('<bits>', 'number of bits (4, 8, 16, 32, 64)')
   .argument('<value>', 'integer to encrypt')
   .action(async (bits, value, options) => {
+    if (!allowedBits.includes(+bits)) throwError('Invalid number of bits');
     const host = prependHttps(options.node);
     const provider = new JsonRpcProvider(host);
     const instance = await getInstance(provider);

package.json

@@ -1,6 +1,6 @@
 {
   "name": "fhevmjs",
-  "version": "0.4.0-3",
+  "version": "0.4.0-4",
   "description": "fhEVM SDK for blockchain using TFHE",
   "main": "lib/node.cjs",
   "types": "lib/node.d.ts",

src/sdk/encrypt.test.ts

@@ -1,8 +1,10 @@
 import {
+  FheUint4,
   FheUint8,
   FheUint16,
   FheUint32,
   FheUint64,
+  CompactFheUint4List,
   CompactFheUint8List,
   CompactFheUint16List,
   CompactFheUint32List,
@@ -11,7 +13,7 @@ import {
   TfheClientKey,
 } from 'node-tfhe';
 import { createTfheKeypair } from '../tfhe';
-import { encrypt8, encrypt16, encrypt32, encrypt64 } from './encrypt';
+import { encrypt4, encrypt8, encrypt16, encrypt32, encrypt64 } from './encrypt';
 
 describe('encrypt8', () => {
   let clientKey: TfheClientKey;
@@ -23,6 +25,26 @@ describe('encrypt8', () => {
     publicKey = keypair.publicKey;
   });
 
+  it('encrypt/decrypt 0 4bits', async () => {
+    const buffer = encrypt4(0, publicKey);
+    const compactList = CompactFheUint4List.deserialize(buffer);
+    let encryptedList = compactList.expand();
+    encryptedList.forEach((v: FheUint4) => {
+      const decrypted = v.decrypt(clientKey);
+      expect(decrypted).toBe(0);
+    });
+  });
+
+  it('encrypt/decrypt 4bits', async () => {
+    const buffer = encrypt4(7, publicKey);
+    const compactList = CompactFheUint4List.deserialize(buffer);
+    let encryptedList = compactList.expand();
+    encryptedList.forEach((v: FheUint4) => {
+      const decrypted = v.decrypt(clientKey);
+      expect(decrypted).toBe(7);
+    });
+  });
+
   it('encrypt/decrypt 0 8bits', async () => {
     const buffer = encrypt8(0, publicKey);
     const compactList = CompactFheUint8List.deserialize(buffer);
@@ -83,7 +105,17 @@ describe('encrypt8', () => {
     });
   });
 
-  it('encrypt/decrypt 32bits', async () => {
+  it('encrypt/decrypt 0 64bits', async () => {
+    const buffer = encrypt64(0, publicKey);
+    const compactList = CompactFheUint64List.deserialize(buffer);
+    let encryptedList = compactList.expand();
+    encryptedList.forEach((v: FheUint64) => {
+      const decrypted = v.decrypt(clientKey);
+      expect(decrypted.toString()).toBe('0');
+    });
+  });
+
+  it('encrypt/decrypt 64bits', async () => {
     const buffer = encrypt64(3021094839202949, publicKey);
     const compactList = CompactFheUint64List.deserialize(buffer);
     let encryptedList = compactList.expand();
@@ -92,4 +124,14 @@ describe('encrypt8', () => {
       expect(decrypted.toString()).toBe('3021094839202949');
     });
   });
+
+  it('encrypt/decrypt 64bits', async () => {
+    const buffer = encrypt64(BigInt('18446744073709551615'), publicKey);
+    const compactList = CompactFheUint64List.deserialize(buffer);
+    let encryptedList = compactList.expand();
+    encryptedList.forEach((v: FheUint64) => {
+      const decrypted = v.decrypt(clientKey);
+      expect(decrypted.toString()).toBe('18446744073709551615');
+    });
+  });
 });

src/sdk/encrypt.ts

@@ -1,11 +1,24 @@
 import {
   TfheCompactPublicKey,
+  CompactFheUint4List,
   CompactFheUint8List,
   CompactFheUint16List,
   CompactFheUint32List,
   CompactFheUint64List,
 } from 'node-tfhe';
 
+export const encrypt4 = (
+  value: number,
+  publicKey: TfheCompactPublicKey,
+): Uint8Array => {
+  const uint8Array = new Uint8Array([value]);
+  const encrypted = CompactFheUint4List.encrypt_with_compact_public_key(
+    uint8Array,
+    publicKey,
+  );
+  return encrypted.serialize();
+};
+
 export const encrypt8 = (
   value: number,
   publicKey: TfheCompactPublicKey,
@@ -43,7 +56,7 @@ export const encrypt32 = (
 };
 
 export const encrypt64 = (
-  value: number,
+  value: number | bigint,
   publicKey: TfheCompactPublicKey,
 ): Uint8Array => {
   const uint64Array = new BigUint64Array([BigInt(value)]);

src/sdk/index.test.ts

@@ -3,7 +3,7 @@ import { createInstance } from './index';
 import { createTfhePublicKey } from '../tfhe';
 import { fromHexString, toHexString, numberToBytes } from '../utils';
 
-describe('token', () => {
+describe('index', () => {
   let tfhePublicKey: string;
 
   beforeAll(async () => {
@@ -100,10 +100,13 @@ describe('token', () => {
       chainId: 1234,
       publicKey: tfhePublicKey,
     });
+    expect(instance.encrypt64(BigInt(34))).toBeTruthy();
 
+    expect(() => instance.encrypt4(undefined as any)).toThrow('Missing value');
     expect(() => instance.encrypt8(undefined as any)).toThrow('Missing value');
     expect(() => instance.encrypt16(undefined as any)).toThrow('Missing value');
     expect(() => instance.encrypt32(undefined as any)).toThrow('Missing value');
+    expect(() => instance.encrypt64(undefined as any)).toThrow('Missing value');
 
     expect(() => instance.encrypt8('wrong value' as any)).toThrow(
       'Value must be a number',
@@ -114,22 +117,22 @@ describe('token', () => {
     expect(() => instance.encrypt32('wrong value' as any)).toThrow(
       'Value must be a number',
     );
-  });
+    expect(() => instance.encrypt64('wrong value' as any)).toThrow(
+      'Value must be a number',
+    );
 
-  it('controls generateToken', async () => {
-    const instance = await createInstance({
-      chainId: 1234,
-      publicKey: tfhePublicKey,
-    });
-    expect(() => instance.generateToken(undefined as any)).toThrow(
-      'Missing contract address',
+    expect(() => instance.encrypt4(BigInt(34) as any)).toThrow(
+      'Value must be a number',
     );
-    expect(() => instance.generateToken({ verifyingContract: '' })).toThrow(
-      'Missing contract address',
+    expect(() => instance.encrypt8(BigInt(34) as any)).toThrow(
+      'Value must be a number',
+    );
+    expect(() => instance.encrypt16(BigInt(34) as any)).toThrow(
+      'Value must be a number',
+    );
+    expect(() => instance.encrypt32(BigInt(34) as any)).toThrow(
+      'Value must be a number',
     );
-    expect(() =>
-      instance.generateToken({ verifyingContract: '0x847473829d' }),
-    ).toThrow('Invalid contract address');
   });
 
   it('controls generatePublicKey', async () => {
@@ -148,7 +151,7 @@ describe('token', () => {
     ).toThrow('Invalid contract address');
   });
 
-  it('save generated token', async () => {
+  it('save generated public key', async () => {
     const instance = await createInstance({
       chainId: 1234,
       publicKey: tfhePublicKey,
@@ -168,26 +171,6 @@ describe('token', () => {
     expect(kp!.publicKey).toBe(publicKey);
   });
 
-  it('save generated token (deprecated)', async () => {
-    const instance = await createInstance({
-      chainId: 1234,
-      publicKey: tfhePublicKey,
-    });
-
-    const contractAddress = '0x1c786b8ca49D932AFaDCEc00827352B503edf16c';
-
-    const { token, publicKey } = instance.generateToken({
-      verifyingContract: contractAddress,
-    });
-
-    instance.setTokenSignature(contractAddress, 'signnnn');
-
-    expect(instance.hasKeypair(contractAddress)).toBeTruthy();
-
-    const kp = instance.getTokenSignature(contractAddress);
-    expect(kp!.publicKey).toBe(publicKey);
-  });
-
   it("don't export keys without signature", async () => {
     const instance = await createInstance({
       chainId: 1234,
@@ -218,7 +201,7 @@ describe('token', () => {
       verifyingContract: contractAddress,
     });
 
-    instance.setTokenSignature(contractAddress, 'signnnn');
+    instance.setSignature(contractAddress, 'signnnn');
 
     const kp = instance.getPublicKey(contractAddress);
     expect(kp!.publicKey).toBe(publicKey);