Address PR comments

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",

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 {
   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
     }
   }

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);
   }