From 15f7a6e700a69c9a40abb82d989a55032d5cf687 Mon Sep 17 00:00:00 2001
From: John Turpish <97759690+John-LittleBearLabs@users.noreply.github.com>
Date: Wed, 10 Aug 2022 10:02:54 -0400
Subject: [PATCH] fix: Allow the caller to set the prologue (#181) (#182)

* fix: Allow the caller to set the prologue (#181)

* The easier PR comments.

* Add a test.

* refix formatting
---
 src/noise.ts       |  5 +++--
 test/noise.spec.ts | 26 ++++++++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/noise.ts b/src/noise.ts
index 5b52b9d..d049913 100644
--- a/src/noise.ts
+++ b/src/noise.ts
@@ -27,7 +27,7 @@ export class Noise implements INoiseConnection {
   public protocol = '/noise'
   public crypto: ICryptoInterface
 
-  private readonly prologue = new Uint8Array(0)
+  private readonly prologue: Uint8Array
   private readonly staticKeys: KeyPair
   private readonly earlyData?: bytes
 
@@ -35,7 +35,7 @@ export class Noise implements INoiseConnection {
    * @param {bytes} staticNoiseKey - x25519 private key, reuse for faster handshakes
    * @param {bytes} earlyData
    */
-  constructor (staticNoiseKey?: bytes, earlyData?: bytes, crypto: ICryptoInterface = stablelib) {
+  constructor (staticNoiseKey?: bytes, earlyData?: bytes, crypto: ICryptoInterface = stablelib, prologueBytes?: Uint8Array) {
     this.earlyData = earlyData ?? new Uint8Array(0)
     this.crypto = crypto
 
@@ -45,6 +45,7 @@ export class Noise implements INoiseConnection {
     } else {
       this.staticKeys = this.crypto.generateX25519KeyPair()
     }
+    this.prologue = prologueBytes ?? new Uint8Array(0)
   }
 
   /**
diff --git a/test/noise.spec.ts b/test/noise.spec.ts
index 0409a69..b446650 100644
--- a/test/noise.spec.ts
+++ b/test/noise.spec.ts
@@ -119,9 +119,9 @@ describe('Noise', () => {
       const wrappedInbound = pbStream(inbound.conn)
       const wrappedOutbound = pbStream(outbound.conn)
 
-      const largePlaintext = randomBytes(100000)
+      const largePlaintext = randomBytes(60000)
       wrappedOutbound.writeLP(Buffer.from(largePlaintext))
-      const response = await wrappedInbound.read(100000)
+      const response = await wrappedInbound.read(60000)
 
       expect(response.length).equals(largePlaintext.length)
     } catch (e) {
@@ -183,4 +183,26 @@ describe('Noise', () => {
       assert(false, err.message)
     }
   })
+
+  it('should accept a prologue', async () => {
+    try {
+      const noiseInit = new Noise(undefined, undefined, stablelib, Buffer.from('Some prologue'))
+      const noiseResp = new Noise(undefined, undefined, stablelib, Buffer.from('Some prologue'))
+
+      const [inboundConnection, outboundConnection] = duplexPair<Uint8Array>()
+      const [outbound, inbound] = await Promise.all([
+        noiseInit.secureOutbound(localPeer, outboundConnection, remotePeer),
+        noiseResp.secureInbound(remotePeer, inboundConnection, localPeer)
+      ])
+      const wrappedInbound = pbStream(inbound.conn)
+      const wrappedOutbound = pbStream(outbound.conn)
+
+      wrappedOutbound.writeLP(Buffer.from('test'))
+      const response = await wrappedInbound.readLP()
+      expect(uint8ArrayToString(response.slice())).equal('test')
+    } catch (e) {
+      const err = e as Error
+      assert(false, err.message)
+    }
+  })
 })