From 36c3fa8ccb492100242fccbde6ec4481c69622d8 Mon Sep 17 00:00:00 2001
From: morrigan <belma.dzi@gmail.com>
Date: Mon, 2 Dec 2019 10:48:19 +0100
Subject: [PATCH] Address PR comments

---
 package.json     |  2 +-
 src/crypto.ts    | 11 +++++++----
 src/handshake.ts | 12 +++++++-----
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/package.json b/package.json
index 6a89844..62758b8 100644
--- a/package.json
+++ b/package.json
@@ -28,7 +28,6 @@
     "@typescript-eslint/parser": "^2.6.0",
     "bn.js-typings": "^1.0.1",
     "chai": "^4.2.0",
-    "debug": "^4.1.1",
     "eslint": "^6.6.0",
     "libp2p-crypto": "^0.17.1",
     "mocha": "^6.2.2",
@@ -56,6 +55,7 @@
     "bcrypto": "^4.2.3",
     "bn.js": "^5.0.0",
     "buffer": "^5.4.3",
+    "debug": "^4.1.1",
     "it-buffer": "^0.1.1",
     "it-length-prefixed": "^3.0.0",
     "it-pair": "^1.0.0",
diff --git a/src/crypto.ts b/src/crypto.ts
index c27d056..fb8308a 100644
--- a/src/crypto.ts
+++ b/src/crypto.ts
@@ -1,15 +1,17 @@
 import { Duplex } from "it-pair";
 import { Handshake } from "./handshake";
+import { Buffer } from "buffer";
 
 interface ReturnEncryptionWrapper {
-  (source: any): any;
+  (source: Iterable<Uint8Array>): any;
 }
 
 // Returns generator that encrypts payload from the user
-export function encryptStream(handshake: Handshake): ReturnEncryptionWrapper {
+  export function encryptStream(handshake: Handshake): ReturnEncryptionWrapper {
   return async function * (source) {
     for await (const chunk of source) {
-      const data = await handshake.encrypt(chunk, handshake.session);
+      const chunkBuffer = Buffer.from(chunk);
+      const data = await handshake.encrypt(chunkBuffer, handshake.session);
       yield data;
     }
   }
@@ -20,7 +22,8 @@ export function encryptStream(handshake: Handshake): ReturnEncryptionWrapper {
 export function decryptStream(handshake: Handshake): ReturnEncryptionWrapper {
   return async function * (source) {
     for await (const chunk of source) {
-      const decrypted = await handshake.decrypt(chunk, handshake.session);
+      const chunkBuffer = Buffer.from(chunk);
+      const decrypted = await handshake.decrypt(chunkBuffer, handshake.session);
       yield decrypted
     }
   }
diff --git a/src/handshake.ts b/src/handshake.ts
index ed4f594..e680f3e 100644
--- a/src/handshake.ts
+++ b/src/handshake.ts
@@ -61,8 +61,8 @@ export class Handshake {
       logger("Stage 0 - Initiator finished proposing, sent signed NoiseHandshake payload.");
     } else {
       const receivedMessageBuffer = decodeMessageBuffer((await this.connection.readLP()).slice());
-
       const plaintext = await this.xx.recvMessage(this.session, receivedMessageBuffer);
+      // TODO: Verify payload
       logger("Stage 0 - Responder received proposed message and remote static public key.");
     }
   }
@@ -72,7 +72,8 @@ export class Handshake {
     if (this.isInitiator) {
       const receivedMessageBuffer = decodeMessageBuffer((await this.connection.readLP()).slice());
       const plaintext = await this.xx.recvMessage(this.session, receivedMessageBuffer);
-      logger('Stage 1 - Initiator received the message.');
+      // TODO: Verify payload
+      logger('Stage 1 - Initiator received the message. Got remote\'s static key.');
     } else {
       // create payload as responder
       const signedPayload = signPayload(this.libp2pPrivateKey, getHandshakePayload(this.staticKeys.publicKey));
@@ -97,16 +98,17 @@ export class Handshake {
     } else {
       const receivedMessageBuffer = (await this.connection.readLP()).slice();
       const plaintext = await this.xx.recvMessage(this.session, decodeMessageBuffer(receivedMessageBuffer));
-      logger('Stage 2 - Responder received the message, finished handshake.')
+      logger('Stage 2 - Responder received the message, finished handshake. Got remote\'s static key.')
     }
   }
 
-  encrypt(plaintext: bytes, session: NoiseSession): bytes {
+  public encrypt(plaintext: bytes, session: NoiseSession): bytes {
     const cs = this.getCS(session);
+
     return this.xx.encryptWithAd(cs, Buffer.alloc(0), plaintext);
   }
 
-  decrypt(ciphertext: bytes, session: NoiseSession): bytes {
+  public decrypt(ciphertext: bytes, session: NoiseSession): bytes {
     const cs = this.getCS(session, false);
     return this.xx.decryptWithAd(cs, Buffer.alloc(0), ciphertext);
   }