Merge pull request #59 from zama-ai/decryptAddress

feat: add decryptAddress

src/sdk/decrypt.test.ts

@@ -1,6 +1,7 @@
 import sodium from 'libsodium-wrappers';
-import { decrypt } from './decrypt';
-import { numberToBytes, toHexString } from '../utils';
+import { decrypt, decryptAddress } from './decrypt';
+import { bigIntToBytes } from '../utils';
+import { getAddress } from 'ethers';
 
 describe('decrypt', () => {
   beforeAll(async () => {
@@ -10,9 +11,9 @@ describe('decrypt', () => {
   it('decrypts a hex value', async () => {
     const keypair = sodium.crypto_box_keypair();
 
-    const value = 28482;
+    const value = BigInt(28482);
     const ciphertext = sodium.crypto_box_seal(
-      numberToBytes(value),
+      bigIntToBytes(value),
       keypair.publicKey,
       'hex',
     );
@@ -23,12 +24,24 @@ describe('decrypt', () => {
   it('decrypts a Uint8Array value', async () => {
     const keypair = sodium.crypto_box_keypair();
 
-    const value = 10;
+    const value = BigInt('10000939393388484938938389392929298383');
     const ciphertext = sodium.crypto_box_seal(
-      numberToBytes(value),
+      bigIntToBytes(value),
       keypair.publicKey,
     );
     const cleartext = decrypt(keypair, ciphertext);
-    expect(cleartext.toString()).toBe(`${value}`);
+    expect(cleartext.toString()).toBe(value.toString());
+  });
+
+  it('decrypts an address Uint8Array value', async () => {
+    const keypair = sodium.crypto_box_keypair();
+
+    const value = BigInt('0x8ba1f109551bd432803012645ac136ddd64dba72');
+    const ciphertext = sodium.crypto_box_seal(
+      bigIntToBytes(value),
+      keypair.publicKey,
+    );
+    const cleartext = decryptAddress(keypair, ciphertext);
+    expect(cleartext).toBe(getAddress(value.toString(16)));
   });
 });

src/sdk/decrypt.ts

@@ -1,6 +1,7 @@
 import sodium from 'libsodium-wrappers';
-import { bytesToBigInt, fromHexString } from '../utils';
+import { bytesToBigInt, fromHexString, bytesToHex } from '../utils';
 import { ContractKeypair } from './types';
+import { getAddress } from 'ethers';
 
 export const decrypt = (
   keypair: ContractKeypair,
@@ -15,3 +16,17 @@ export const decrypt = (
   );
   return bytesToBigInt(decrypted);
 };
+
+export const decryptAddress = (
+  keypair: ContractKeypair,
+  ciphertext: string | Uint8Array,
+): string => {
+  const toDecrypt =
+    typeof ciphertext === 'string' ? fromHexString(ciphertext) : ciphertext;
+  const decrypted = sodium.crypto_box_seal_open(
+    toDecrypt,
+    keypair.publicKey,
+    keypair.privateKey,
+  );
+  return getAddress(bytesToHex(decrypted));
+};

src/sdk/index.test.ts

@@ -1,7 +1,7 @@
 import sodium from 'libsodium-wrappers';
 import { createInstance } from './index';
 import { createTfhePublicKey } from '../tfhe';
-import { fromHexString, toHexString, numberToBytes } from '../utils';
+import { fromHexString, toHexString, bigIntToBytes } from '../utils';
 
 describe('index', () => {
   let tfhePublicKey: string;
@@ -23,6 +23,7 @@ describe('index', () => {
     expect(instance.encrypt32).toBeDefined();
     expect(instance.generatePublicKey).toBeDefined();
     expect(instance.decrypt).toBeDefined();
+    expect(instance.decryptAddress).toBeDefined();
     expect(instance.serializeKeypairs).toBeDefined();
     expect(instance.getPublicKey).toBeDefined();
     expect(instance.hasKeypair).toBeDefined();
@@ -38,6 +39,7 @@ describe('index', () => {
     expect(instance.encrypt64).toBeDefined();
     expect(instance.generatePublicKey).toBeDefined();
     expect(instance.decrypt).toBeDefined();
+    expect(instance.decryptAddress).toBeDefined();
     expect(instance.serializeKeypairs).toBeDefined();
     expect(instance.getPublicKey).toBeDefined();
     expect(instance.hasKeypair).toBeDefined();
@@ -86,15 +88,93 @@ describe('index', () => {
       },
     });
 
-    const value = 937387;
+    const value = BigInt(937387);
     const ciphertext = sodium.crypto_box_seal(
-      numberToBytes(value),
+      bigIntToBytes(value),
       fromHexString(keypair.publicKey),
       'hex',
     );
 
     const cleartext = instance.decrypt(contractAddress, ciphertext);
-    expect(cleartext.toString()).toBe(`${value}`);
+    expect(cleartext.toString()).toBe(value.toString());
+
+    const address = BigInt('0xD115BFFAbbdd893A6f7ceA402e7338643Ced44a6');
+    const ciphertextAddress = sodium.crypto_box_seal(
+      bigIntToBytes(address),
+      fromHexString(keypair.publicKey),
+      'hex',
+    );
+
+    const cleartextAddress = instance.decryptAddress(
+      contractAddress,
+      ciphertextAddress,
+    );
+    expect(cleartextAddress).toBe('0xD115BFFAbbdd893A6f7ceA402e7338643Ced44a6');
+  });
+
+  it('controls decrypt', async () => {
+    const instance = await createInstance({
+      chainId: 1234,
+      publicKey: tfhePublicKey,
+    });
+
+    const keypair = instance.generatePublicKey({
+      verifyingContract: '0xD115BFFAbbdd893A6f7ceA402e7338643Ced44a6',
+    });
+
+    const value = BigInt(937387);
+    const ciphertext = sodium.crypto_box_seal(
+      bigIntToBytes(value),
+      keypair.publicKey,
+      'hex',
+    );
+
+    const address = BigInt('0xD115BFFAbbdd893A6f7ceA402e7338643Ced44a6');
+    const ciphertextAddress = sodium.crypto_box_seal(
+      bigIntToBytes(address),
+      keypair.publicKey,
+      'hex',
+    );
+
+    expect(() => instance.decrypt(undefined as any, ciphertext)).toThrow(
+      'Missing contract address.',
+    );
+
+    expect(() =>
+      instance.decryptAddress(undefined as any, ciphertextAddress),
+    ).toThrow('Missing contract address.');
+
+    expect(() =>
+      instance.decrypt(
+        '0xD115BFFAbbdd893A6f7ceA402e7338643Ced44a6',
+        undefined as any,
+      ),
+    ).toThrow('Missing ciphertext.');
+
+    expect(() =>
+      instance.decryptAddress(
+        '0xD115BFFAbbdd893A6f7ceA402e7338643Ced44a6',
+        undefined as any,
+      ),
+    ).toThrow('Missing ciphertext.');
+
+    expect(() =>
+      instance.decrypt(
+        '0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb',
+        ciphertext,
+      ),
+    ).toThrow(
+      'Missing keypair for 0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb.',
+    );
+
+    expect(() =>
+      instance.decryptAddress(
+        '0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb',
+        ciphertextAddress,
+      ),
+    ).toThrow(
+      'Missing keypair for 0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb.',
+    );
   });
 
   it('controls encrypt', async () => {
@@ -262,13 +342,40 @@ describe('index', () => {
     const kp = instance.getPublicKey(contractAddress);
     expect(kp!.publicKey).toBe(publicKey);
 
-    const value = 89290;
+    const value = BigInt(89290);
     const ciphertext = sodium.crypto_box_seal(
-      numberToBytes(value),
+      bigIntToBytes(value),
       publicKey,
       'hex',
     );
     const cleartext = instance.decrypt(contractAddress, ciphertext);
-    expect(cleartext.toString()).toBe(`${value}`);
+    expect(cleartext.toString()).toBe(value.toString());
+  });
+
+  it('decrypts address', async () => {
+    const instance = await createInstance({
+      chainId: 1234,
+      publicKey: tfhePublicKey,
+    });
+
+    const contractAddress = '0x1c786b8ca49D932AFaDCEc00827352B503edf16c';
+
+    const { eip712, publicKey } = instance.generatePublicKey({
+      verifyingContract: contractAddress,
+    });
+
+    instance.setSignature(contractAddress, 'signnnn');
+
+    const kp = instance.getPublicKey(contractAddress);
+    expect(kp!.publicKey).toBe(publicKey);
+
+    const value = BigInt('0x1c786b8ca49D932AFaDCEc00827352B503edf16c');
+    const ciphertext = sodium.crypto_box_seal(
+      bigIntToBytes(value),
+      publicKey,
+      'hex',
+    );
+    const cleartext = instance.decryptAddress(contractAddress, ciphertext);
+    expect(cleartext).toBe('0x1c786b8ca49D932AFaDCEc00827352B503edf16c');
   });
 });

src/sdk/index.ts

@@ -14,7 +14,7 @@ import {
   GeneratePublicKeyParams,
   generatePublicKey,
 } from './publicKey';
-import { decrypt } from './decrypt';
+import { decrypt, decryptAddress } from './decrypt';
 import { fromHexString, isAddress, toHexString } from '../utils';
 import { ContractKeypairs } from './types';
 
@@ -40,6 +40,7 @@ export type FhevmInstance = {
   ) => { publicKey: Uint8Array; signature: string } | null;
   hasKeypair: (contractAddress: string) => boolean;
   decrypt: (contractAddress: string, ciphertext: string) => bigint;
+  decryptAddress: (contractAddress: string, ciphertext: string) => string;
   serializeKeypairs: () => ExportedContractKeypairs;
 };
 
@@ -251,6 +252,14 @@ export const createInstance = async (
       return decrypt(kp, ciphertext);
     },
 
+    decryptAddress(contractAddress, ciphertext) {
+      if (!ciphertext) throw new Error('Missing ciphertext.');
+      if (!contractAddress) throw new Error('Missing contract address.');
+      const kp = contractKeypairs[contractAddress];
+      if (!kp) throw new Error(`Missing keypair for ${contractAddress}.`);
+      return decryptAddress(kp, ciphertext);
+    },
+
     serializeKeypairs() {
       const stringKeypairs: ExportedContractKeypairs = {};
       Object.keys(contractKeypairs).forEach((contractAddress) => {

src/utils.test.ts

@@ -1,23 +1,27 @@
-import { numberToBytes, bytesToBigInt } from './utils';
+import { bigIntToBytes, bytesToBigInt } from './utils';
 
 describe('decrypt', () => {
   it('converts a number to bytes', async () => {
-    const value = 28482;
-    const bytes = numberToBytes(value);
+    const value = BigInt(28482);
+    const bytes = bigIntToBytes(value);
     expect(bytes).toEqual(new Uint8Array([111, 66]));
 
-    const value2 = 255;
-    const bytes2 = numberToBytes(value2);
+    const value2 = BigInt(255);
+    const bytes2 = bigIntToBytes(value2);
     expect(bytes2).toEqual(new Uint8Array([255]));
   });
 
   it('converts bytes to number', async () => {
     const value = new Uint8Array([23, 200, 15]);
-    const bytes = bytesToBigInt(value);
-    expect(bytes.toString()).toBe('1558543');
+    const bigint1 = bytesToBigInt(value);
+    expect(bigint1.toString()).toBe('1558543');
 
-    const value2 = new Uint8Array();
-    const bytes2 = bytesToBigInt(value2);
-    expect(bytes2.toString()).toBe('0');
+    const value2 = new Uint8Array([37, 6, 210, 166, 239]);
+    const bigint2 = bytesToBigInt(value2);
+    expect(bigint2.toString()).toBe('159028258543');
+
+    const value0 = new Uint8Array();
+    const bigint0 = bytesToBigInt(value0);
+    expect(bigint0.toString()).toBe('0');
   });
 });

src/utils.ts

@@ -1,4 +1,4 @@
-import { toBigIntBE } from 'bigint-buffer';
+import { toBigIntBE, toBufferBE } from 'bigint-buffer';
 
 export const fromHexString = (hexString: string): Uint8Array => {
   const arr = hexString.replace(/^(0x)/, '').match(/.{1,2}/g);
@@ -9,16 +9,22 @@ export const fromHexString = (hexString: string): Uint8Array => {
 export const toHexString = (bytes: Uint8Array) =>
   bytes.reduce((str, byte) => str + byte.toString(16).padStart(2, '0'), '');
 
-export const numberToBytes = (uint32Value: number) => {
-  const byteArrayLength = Math.ceil(Math.log2(uint32Value + 1) / 8);
-  const byteArray = new Uint8Array(byteArrayLength);
+export const bigIntToBytes = (value: bigint) => {
+  const byteArrayLength = Math.ceil(value.toString(2).length / 8);
+  return new Uint8Array(toBufferBE(value, byteArrayLength));
+};
 
-  for (let i = byteArrayLength - 1; i >= 0; i--) {
-    byteArray[i] = uint32Value & 0xff;
-    uint32Value >>= 8;
+export const bytesToHex = function (byteArray: Uint8Array): string {
+  if (!byteArray || byteArray?.length === 0) {
+    return '0x0';
   }
 
-  return byteArray;
+  const length = byteArray.length;
+
+  const buffer = Buffer.from(byteArray);
+  const result = buffer.toString('hex');
+
+  return result;
 };
 
 export const bytesToBigInt = function (byteArray: Uint8Array): bigint {