Address PR comment and refactor setting remote peer

src/handshake-ik.ts

@@ -6,7 +6,7 @@ import {KeyPair} from "./@types/libp2p";
 import {IHandshake} from "./@types/handshake-interface";
 import {Buffer} from "buffer";
 import {decode0, decode1, encode0, encode1} from "./encoder";
-import {getPeerIdFromPayload, verifySignedPayload} from "./utils";
+import {decodePayload, getPeerIdFromPayload, verifySignedPayload} from "./utils";
 import {FailedIKError} from "./errors";
 import {logger} from "./logger";
 import PeerId from "peer-id";
@@ -57,7 +57,9 @@ export class IKHandshake implements IHandshake {
         const receivedMessageBuffer = decode1(receivedMsg);
         const plaintext = this.ik.recvMessage(this.session, receivedMessageBuffer);
         logger("IK Stage 0 - Responder got message, going to verify payload.");
-        this.remotePeer = await verifySignedPayload(receivedMessageBuffer.ns, plaintext, this.remotePeer);
+        const decodedPayload = await decodePayload(plaintext);
+        this.remotePeer = this.remotePeer || await getPeerIdFromPayload(decodedPayload);
+        await verifySignedPayload(receivedMessageBuffer.ns, decodedPayload, this.remotePeer);
         logger("IK Stage 0 - Responder successfully verified payload!");
       } catch (e) {
         logger("Responder breaking up with IK handshake in stage 0.");
@@ -76,7 +78,9 @@ export class IKHandshake implements IHandshake {
       logger("IK Stage 1 - Initiator got message, going to verify payload.");
 
       try {
-        await verifySignedPayload(receivedMessageBuffer.ns, plaintext, this.remotePeer);
+        const decodedPayload = await decodePayload(plaintext);
+        this.remotePeer = this.remotePeer || await getPeerIdFromPayload(decodedPayload);
+        await verifySignedPayload(receivedMessageBuffer.ns, decodedPayload, this.remotePeer);
         logger("IK Stage 1 - Initiator successfully verified payload!");
       } catch (e) {
         logger("Initiator breaking up with IK handshake in stage 1.");

src/handshake-xx-fallback.ts

@@ -3,7 +3,7 @@ import {XXHandshake} from "./handshake-xx";
 import {XX} from "./handshakes/xx";
 import {KeyPair} from "./@types/libp2p";
 import {bytes, bytes32} from "./@types/basic";
-import {getPeerIdFromPayload, verifySignedPayload,} from "./utils";
+import {decodePayload, getPeerIdFromPayload, verifySignedPayload,} from "./utils";
 import {logger} from "./logger";
 import {WrappedConnection} from "./noise";
 import {decode0, decode1} from "./encoder";
@@ -57,14 +57,16 @@ export class XXFallbackHandshake extends XXHandshake {
 
       logger("Initiator going to check remote's signature...");
       try {
-        this.remotePeer = await verifySignedPayload(receivedMessageBuffer.ns, plaintext, this.remotePeer);
+        const decodedPayload = await decodePayload(plaintext);
+        this.remotePeer = this.remotePeer || await getPeerIdFromPayload(decodedPayload);
+        await verifySignedPayload(receivedMessageBuffer.ns, decodedPayload, this.remotePeer);
       } catch (e) {
         throw new Error(`Error occurred while verifying signed payload from responder: ${e.message}`);
       }
       logger("All good with the signature!");
     } else {
       logger("XX Fallback Stage 1 - Responder start");
-      super.exchange();
+      await super.exchange();
       logger("XX Fallback Stage 1 - Responder end");
     }
   }

src/handshake-xx.ts

@@ -6,6 +6,7 @@ import { bytes, bytes32 } from "./@types/basic";
 import { NoiseSession } from "./@types/handshake";
 import {IHandshake} from "./@types/handshake-interface";
 import {
+  decodePayload,
   getPeerIdFromPayload,
   verifySignedPayload,
 } from "./utils";
@@ -72,7 +73,9 @@ export class XXHandshake implements IHandshake {
 
       logger("Initiator going to check remote's signature...");
       try {
-        this.remotePeer = await verifySignedPayload(receivedMessageBuffer.ns, plaintext, this.remotePeer);
+        const decodedPayload = await decodePayload(plaintext);
+        this.remotePeer = this.remotePeer || await getPeerIdFromPayload(decodedPayload);
+        this.remotePeer = await verifySignedPayload(receivedMessageBuffer.ns, decodedPayload, this.remotePeer);
       } catch (e) {
         throw new Error(`Error occurred while verifying signed payload: ${e.message}`);
       }
@@ -99,7 +102,9 @@ export class XXHandshake implements IHandshake {
       logger('Stage 2 - Responder received the message, finished handshake. Got remote\'s static key.');
 
       try {
-        this.remotePeer = await verifySignedPayload(receivedMessageBuffer.ns, plaintext, this.remotePeer);
+        const decodedPayload = await decodePayload(plaintext);
+        this.remotePeer = this.remotePeer || await getPeerIdFromPayload(decodedPayload);
+        await verifySignedPayload(receivedMessageBuffer.ns, decodedPayload, this.remotePeer);
       } catch (e) {
         throw new Error(`Error occurred while verifying signed payload: ${e.message}`);
       }

src/utils.ts

@@ -63,12 +63,11 @@ export async function signPayload(peerId: PeerId, payload: bytes): Promise<bytes
   return peerId.privKey.sign(payload);
 }
 
-export async function getPeerIdFromPayload(payload: bytes): Promise<PeerId> {
-  const decodedPayload = await decodePayload(payload);
-  return await PeerId.createFromPubKey(Buffer.from(decodedPayload.identityKey));
+export async function getPeerIdFromPayload(payload: INoisePayload): Promise<PeerId> {
+  return await PeerId.createFromPubKey(Buffer.from(payload.identityKey));
 }
 
-async function decodePayload(payload: bytes): Promise<INoisePayload> {
+export async function decodePayload(payload: bytes): Promise<INoisePayload> {
   const NoiseHandshakePayload = await loadPayloadProto();
   return NoiseHandshakePayload.toObject(
     NoiseHandshakePayload.decode(payload)
@@ -83,42 +82,35 @@ async function isValidPeerId(peerId: bytes, publicKeyProtobuf: bytes) {
 }
 
 /**
- * Verifies signed payload and returns peer id that has sent the payload.
+ * Verifies signed payload, throws on any irregularities.
  * @param {bytes} noiseStaticKey - owner's noise static key
- * @param {bytes} plaintext - encoded payload
- * @param {PeerId} remotePeer - (optional) owner's libp2p peer ID
+ * @param {bytes} payload - decoded payload
+ * @param {PeerId} remotePeer - owner's libp2p peer ID
  * @returns {Promise<PeerId>} - peer ID of payload owner
  */
 export async function verifySignedPayload(
   noiseStaticKey: bytes,
-  plaintext: bytes,
-  remotePeer?: PeerId
+  payload: INoisePayload,
+  remotePeer: PeerId
 ): Promise<PeerId> {
-  let receivedPayload;
   try {
-    const NoiseHandshakePayload = await loadPayloadProto();
-    receivedPayload = NoiseHandshakePayload.toObject(
-      NoiseHandshakePayload.decode(plaintext)
-    );
     //temporary fix until protobufsjs conversion options starts working
     //by default it ends up as Uint8Array
-    receivedPayload.identityKey = Buffer.from(receivedPayload.identityKey);
-    receivedPayload.identitySig = Buffer.from(receivedPayload.identitySig);
+    payload.identityKey = Buffer.from(payload.identityKey);
+    payload.identitySig = Buffer.from(payload.identitySig);
   } catch (e) {
     throw new Error("Failed to decode received payload. Reason: " + e.message);
   }
 
-  remotePeer = remotePeer || await getPeerIdFromPayload(plaintext);
-
-  if (!(await isValidPeerId(remotePeer.id, receivedPayload.identityKey)) ) {
+  if (!(await isValidPeerId(remotePeer.id, payload.identityKey)) ) {
     throw new Error("Peer ID doesn't match libp2p public key.");
   }
 
   const generatedPayload = getHandshakePayload(noiseStaticKey);
 
   // Unmarshaling from PublicKey protobuf
-  const publicKey = crypto.keys.unmarshalPublicKey(receivedPayload.identityKey);
-  if (!publicKey.verify(generatedPayload, receivedPayload.identitySig)) {
+  const publicKey = crypto.keys.unmarshalPublicKey(payload.identityKey);
+  if (!publicKey.verify(generatedPayload, payload.identitySig)) {
     throw new Error("Static key doesn't match to peer that signed payload!");
   }