From b2d058291cf33f38a217661d4e843a7ef9b9b59a Mon Sep 17 00:00:00 2001
From: morrigan <belma.dzi@gmail.com>
Date: Thu, 21 Nov 2019 13:38:39 +0100
Subject: [PATCH] Written handshake exchange

---
 src/handshake.ts | 49 +++++++++++++++++++++++++++++++++++-------------
 src/noise.ts     | 11 +++++------
 src/utils.ts     | 11 ++++++-----
 3 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/src/handshake.ts b/src/handshake.ts
index 78907fb..ecc3fda 100644
--- a/src/handshake.ts
+++ b/src/handshake.ts
@@ -2,56 +2,79 @@ import { bytes, bytes32 } from "./@types/basic";
 import { NoiseSession, XXHandshake } from "./xx";
 import { KeyPair } from "./@types/libp2p";
 import { Buffer } from "buffer";
+import {createHandshakePayload, getHandshakePayload, signPayload} from "./utils";
 
 type handshakeType = "XX";
 
 export class Handshake {
   private type: handshakeType;
   private remotePublicKey: bytes;
-  private signedPayload: bytes;
   private prologue: bytes32;
   private staticKeys: KeyPair;
   private connection: any;
+  private xx: XXHandshake;
 
   constructor(
     type: handshakeType,
     remotePublicKey: bytes,
     prologue: bytes32,
-    signedPayload: bytes,
     staticKeys: KeyPair,
     connection,
   ) {
     this.type = type;
     this.remotePublicKey = remotePublicKey;
-    this.signedPayload = signedPayload;
     this.prologue = prologue;
     this.staticKeys = staticKeys;
     this.connection = connection;
+
+    this.xx = new XXHandshake();
   }
 
   // stage 0
-  async propose(isInitiator: boolean) : Promise<NoiseSession> {
-    const xx = new XXHandshake();
-
-    const ns = await xx.initSession(isInitiator, this.prologue, this.staticKeys, this.remotePublicKey);
+  async propose(isInitiator: boolean, earlyData?: bytes) : Promise<NoiseSession> {
+    const ns = await this.xx.initSession(isInitiator, this.prologue, this.staticKeys, this.remotePublicKey);
 
     if (isInitiator) {
-      const message = Buffer.concat([Buffer.alloc(0), this.signedPayload]);
-      const messageBuffer = await xx.sendMessage(ns, message);
+      const signedPayload = signPayload(this.staticKeys.privateKey, getHandshakePayload(this.staticKeys.publicKey));
+      const handshakePayload = await createHandshakePayload(
+        this.staticKeys.publicKey,
+        signedPayload,
+        earlyData,
+        this.staticKeys.privateKey
+      );
+      const message = Buffer.concat([Buffer.alloc(0), handshakePayload]);
+      const messageBuffer = await this.xx.sendMessage(ns, message);
       this.connection.writeLP(messageBuffer);
     } else {
       const receivedMessageBuffer = (await this.connection.readLP()).slice();
-      const plaintext = await xx.recvMessage(ns, receivedMessageBuffer);
+      const plaintext = await this.xx.recvMessage(ns, receivedMessageBuffer);
     }
 
     return ns;
   }
 
-  async exchange() : Promise<NoiseSession> {
+  async exchange(isInitiator: boolean, session: NoiseSession) : Promise<void> {
+    if (isInitiator) {
+      const receivedMessageBuffer = (await this.connection.readLP()).slice();
+      const plaintext = await this.xx.recvMessage(session, receivedMessageBuffer);
+    } else {
+      // create payload as responder
+      const signedPayload = signPayload(this.staticKeys.privateKey, getHandshakePayload(this.staticKeys.publicKey));
+      const handshakePayload = await createHandshakePayload(this.remotePublicKey, signedPayload);
 
+      const message = Buffer.concat([Buffer.alloc(0), handshakePayload]);
+      const messageBuffer = await this.xx.sendMessage(session, message);
+      this.connection.writeLP(messageBuffer);
+    }
   }
 
-  async finish() : Promise<NoiseSession> {
-
+  async finish(isInitiator: boolean, session: NoiseSession) : Promise<void> {
+    if (isInitiator) {
+      const messageBuffer = await this.xx.sendMessage(session, Buffer.alloc(0));
+      this.connection.writeLP(messageBuffer);
+    } else {
+      const receivedMessageBuffer = (await this.connection.readLP()).slice();
+      const plaintext = await this.xx.recvMessage(session, receivedMessageBuffer);
+    }
   }
 }
diff --git a/src/noise.ts b/src/noise.ts
index d10c54c..263ef92 100644
--- a/src/noise.ts
+++ b/src/noise.ts
@@ -68,13 +68,12 @@ export class Noise implements NoiseConnection {
       this.staticKeys = await generateKeypair();
     }
 
-    const payload = getHandshakePayload(this.staticKeys.publicKey);
-    const signedPayload = signPayload(this.staticKeys.privateKey, payload);
-    const handshakePayload = await createHandshakePayload(this.staticKeys, signedPayload);
-
     const prologue = Buffer.from(this.protocol);
-    const handshake = new Handshake('XX', remotePublicKey, prologue, handshakePayload, this.staticKeys, connection);
-    const session = await handshake.propose(isInitiator);
+    const handshake = new Handshake('XX', remotePublicKey, prologue, this.staticKeys, connection);
+
+    const session = await handshake.propose(isInitiator, this.earlyData);
+    await handshake.exchange(isInitiator, session);
+    await handshake.finish(isInitiator, session);
 
     return await encryptStreams(connection, session);
   }
diff --git a/src/utils.ts b/src/utils.ts
index 61635cf..6c3b236 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -21,15 +21,16 @@ export async function generateKeypair() : Promise<KeyPair> {
 }
 
 export async function createHandshakePayload(
-  libp2pKeys: KeyPair,
+  libp2pPublicKey: bytes,
   signedPayload: bytes,
   earlyData?: bytes,
+  libp2pPrivateKey?: bytes,
 ) : Promise<bytes> {
   const NoiseHandshakePayload = await loadPayloadProto();
   const payloadInit = NoiseHandshakePayload.create({
-    libp2pKey: libp2pKeys.publicKey,
+    libp2pKey: libp2pPublicKey,
     noiseStaticKeySignature: signedPayload,
-    ...resolveEarlyDataPayload(libp2pKeys.privateKey, earlyData),
+    ...resolveEarlyDataPayload(libp2pPrivateKey, earlyData),
   });
 
   return Buffer.from(NoiseHandshakePayload.encode(payloadInit).finish());
@@ -44,8 +45,8 @@ export const getHandshakePayload = (publicKey: bytes ) => Buffer.concat([Buffer.
 
 export const getEarlyDataPayload = (earlyData: bytes) => Buffer.concat([Buffer.from("noise-libp2p-early-data:"), earlyData]);
 
-function resolveEarlyDataPayload(privateKey: bytes, earlyData?: bytes) : Object {
-  if (!earlyData) {
+function resolveEarlyDataPayload(privateKey?: bytes, earlyData?: bytes) : Object {
+  if (!earlyData || !privateKey) {
     return {};
   }