From def870b610dbb824965b53551d0e2f031e6929b7 Mon Sep 17 00:00:00 2001
From: Levent DEMIR <levent.demir@zama.ai>
Date: Mon, 18 Mar 2024 11:01:19 +0100
Subject: [PATCH] chore(decryptAddress): add tests

---
 src/sdk/decrypt.test.ts | 28 ++++++++++++++++++++++++++++
 src/sdk/decrypt.ts      |  2 +-
 src/utils.ts            |  4 +---
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/sdk/decrypt.test.ts b/src/sdk/decrypt.test.ts
index 50407e1..44c8485 100644
--- a/src/sdk/decrypt.test.ts
+++ b/src/sdk/decrypt.test.ts
@@ -44,4 +44,32 @@ describe('decrypt', () => {
     const cleartext = decryptAddress(keypair, ciphertext);
     expect(cleartext).toBe(getAddress(value.toString(16)));
   });
+
+  it('decrypts an address Uint8Array value bigger than 160 bits', async () => {
+    const keypair = sodium.crypto_box_keypair();
+    const address = '0x9b8a8ba1f109551bd432803012645ac136ddd64dba72'
+    // Must truncate to 40-digit
+    const expected = '0x8ba1f109551bd432803012645ac136ddd64dba72'
+    const value = BigInt(address);
+    const ciphertext = sodium.crypto_box_seal(
+      bigIntToBytes(value),
+      keypair.publicKey,
+    );
+    const cleartext = decryptAddress(keypair, ciphertext);
+    expect(cleartext).toBe(getAddress(expected));
+  });
+
+  it.only('decrypts an address Uint8Array value lower than 160 bits', async () => {
+    const keypair = sodium.crypto_box_keypair();
+    const address = '0x8ba1f109551bd432803012645ac136ddd64d'
+    // Must add padding until to 40-digit
+    const expected = '0x00008ba1f109551bd432803012645ac136ddd64d'
+    const value = BigInt(address);
+    const ciphertext = sodium.crypto_box_seal(
+      bigIntToBytes(value),
+      keypair.publicKey,
+    );
+    const cleartext = decryptAddress(keypair, ciphertext);
+    expect(cleartext).toBe(getAddress(expected));
+  });
 });
diff --git a/src/sdk/decrypt.ts b/src/sdk/decrypt.ts
index 3ea9530..3ca24b5 100644
--- a/src/sdk/decrypt.ts
+++ b/src/sdk/decrypt.ts
@@ -35,7 +35,7 @@ let hexString = bytesToHex(decrypted);
 if (hexString.length > 40) {
   hexString = hexString.substring(hexString.length - 40);
 } else {
-  hexString = hexString.padStart(40, '0');
+  hexString = hexString.slice(2).padStart(40, '0');
 }
   return getAddress(hexString);
 };
diff --git a/src/utils.ts b/src/utils.ts
index 7283a85..a74dc06 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -19,8 +19,7 @@ export const bytesToHex = function (byteArray: Uint8Array): string {
     return '0x0';
   }
   const buffer = Buffer.from(byteArray);
-  const result = '0x' + buffer.toString('hex');
-  return result;
+  return `0x${buffer.toString('hex')}`;
 };
 
 export const bytesToBigInt = function (byteArray: Uint8Array): bigint {
@@ -30,7 +29,6 @@ export const bytesToBigInt = function (byteArray: Uint8Array): bigint {
 
   const buffer = Buffer.from(byteArray);
   const result = toBigIntBE(buffer);
-
   return result;
 };