From 04426e525d6dce69983ba02834e5f2d0c9a968e4 Mon Sep 17 00:00:00 2001
From: pubkey <8926560+pubkey@users.noreply.github.com>
Date: Wed, 20 Apr 2022 20:08:59 +0200
Subject: [PATCH] ADD options to make encryption deterministic

---
 src/encrypt-with-public-key.js |  5 ++--
 test/unit.test.js              | 44 ++++++++++++++++++++++++++++++++--
 typings/index.d.ts             |  6 ++++-
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/src/encrypt-with-public-key.js b/src/encrypt-with-public-key.js
index fb007b4..a34b10d 100644
--- a/src/encrypt-with-public-key.js
+++ b/src/encrypt-with-public-key.js
@@ -5,7 +5,7 @@ import {
     decompress
 } from './public-key';
 
-export default function encryptWithPublicKey(publicKey, message) {
+export default function encryptWithPublicKey(publicKey, message, opts) {
 
     // ensure its an uncompressed publicKey
     publicKey = decompress(publicKey);
@@ -16,7 +16,8 @@ export default function encryptWithPublicKey(publicKey, message) {
 
     return encrypt(
         Buffer.from(pubString, 'hex'),
-        Buffer.from(message)
+        Buffer.from(message),
+        opts ? opts : {}
     ).then(encryptedBuffers => {
         const encrypted = {
             iv: encryptedBuffers.iv.toString('hex'),
diff --git a/test/unit.test.js b/test/unit.test.js
index 902e848..25a7d31 100644
--- a/test/unit.test.js
+++ b/test/unit.test.js
@@ -2,6 +2,7 @@ const AsyncTestUtil = require('async-test-util');
 const assert = require('assert');
 const BN = require('bn.js');
 const EthCrypto = require('../dist/lib/index');
+const crypto = require('crypto');
 
 const TEST_DATA = {
     address: '0x3f243FdacE01Cfd9719f7359c94BA11361f32471',
@@ -124,7 +125,7 @@ describe('unit.test.js', () => {
 
             });
         });
-        describe('negative', () => {});
+        describe('negative', () => { });
     });
     describe('.recoverPublicKey()', () => {
         it('should recover the correct key', async () => {
@@ -162,6 +163,45 @@ describe('unit.test.js', () => {
                 );
                 assert.equal(decrypted, message);
             });
+            it('should use a random iv if none provided', async () => {
+                const message = AsyncTestUtil.randomString(12);
+                const encrypted = await EthCrypto.encryptWithPublicKey(
+                    TEST_DATA.publicKey,
+                    message
+                );
+                const encrypted2 = await EthCrypto.encryptWithPublicKey(
+                    TEST_DATA.publicKey,
+                    message
+                );
+
+                assert.ok(encrypted.ciphertext !== encrypted2.ciphertext);
+            });
+            it('should have deterministic output if iv is provided', async () => {
+                const message = AsyncTestUtil.randomString(12);
+
+                const iv = crypto.randomBytes(16);
+                const ephemPrivateKey = crypto.randomBytes(32);
+                console.dir(iv);
+                const encrypted = await EthCrypto.encryptWithPublicKey(
+                    TEST_DATA.publicKey,
+                    message,
+                    {
+                        iv,
+                        ephemPrivateKey
+                    }
+                );
+                const encrypted2 = await EthCrypto.encryptWithPublicKey(
+                    TEST_DATA.publicKey,
+                    message,
+                    {
+                        iv,
+                        ephemPrivateKey
+                    }
+                );
+
+                assert.strictEqual(encrypted.ciphertext, encrypted2.ciphertext);
+                assert.strictEqual(encrypted.mac, encrypted2.mac);
+            });
         });
         describe('negative', () => {
             it('should throw when non-key given', async () => {
@@ -204,7 +244,7 @@ describe('unit.test.js', () => {
                 assert.equal(decrypted, message);
             });
         });
-        describe('negative', () => {});
+        describe('negative', () => { });
     });
     describe('.cipher', () => {
         describe('.stringify()', () => {
diff --git a/typings/index.d.ts b/typings/index.d.ts
index 2d1aff8..5ae244f 100644
--- a/typings/index.d.ts
+++ b/typings/index.d.ts
@@ -56,7 +56,11 @@ type vrsType = {
 };
 export const vrs: vrsType;
 
-type encryptWithPublicKeyType = (publicKey: string, message: string) => Promise<Encrypted>;
+export type EncryptOptions = {
+    iv?: Buffer,
+    ephemPrivateKey?: Buffer
+};
+type encryptWithPublicKeyType = (publicKey: string, message: string, options?: EncryptOptions) => Promise<Encrypted>;
 export const encryptWithPublicKey: encryptWithPublicKeyType;
 
 type decryptWithPrivateKeyType = (privateKey: string, encrypted: Encrypted) => Promise<string>;