Write encryption generators

src/crypto.ts

@@ -1,13 +1,28 @@
 import { Duplex } from "it-pair";
 import { NoiseSession } from "./xx";
+import { Handshake } from "./handshake";
 
-// Send encrypted payload from the user to stream
-export async function encryptStreams(streams: Duplex, session: NoiseSession) : Promise<Duplex> {
+interface IReturnEncryptionWrapper {
+  (source: any): any;
+}
 
+// Returns generator that encrypts payload from the user
+export function encryptStream(handshake: Handshake, session: NoiseSession) : IReturnEncryptionWrapper {
+  return async function * (source) {
+    for await (const chunk of source) {
+      const data = await handshake.encrypt(chunk, session);
+      yield data;
+    }
+  }
 }
 
 
-// Decrypt received payload from the stream and pipe to user
-export async function decryptStreams(streams: Duplex, session: NoiseSession) : Promise<Duplex> {
-
+// Decrypt received payload to the user
+export function decryptStreams(handshake: Handshake, session: NoiseSession) : IReturnEncryptionWrapper {
+  return async function * (source) {
+    for await (const chunk of source) {
+      const decrypted = await handshake.decrypt(chunk, session);
+      yield decrypted
+    }
+  }
 }

src/handshake.ts

@@ -9,11 +9,13 @@ import {
   getHandshakePayload,
   signPayload
 } from "./utils";
-import { WrappedConnection } from "./noise";
+import {Noise, WrappedConnection} from "./noise";
 
 type handshakeType = "XX";
 
 export class Handshake {
+  public isInitiator: boolean;
+
   private type: handshakeType;
   private remotePublicKey: bytes;
   private prologue: bytes32;
@@ -23,12 +25,14 @@ export class Handshake {
 
   constructor(
     type: handshakeType,
+    isInitiator: boolean,
     remotePublicKey: bytes,
     prologue: bytes32,
     staticKeys: KeyPair,
     connection: WrappedConnection,
   ) {
     this.type = type;
+    this.isInitiator = isInitiator;
     this.remotePublicKey = remotePublicKey;
     this.prologue = prologue;
     this.staticKeys = staticKeys;
@@ -38,10 +42,10 @@ export class Handshake {
   }
 
   // stage 0
-  async propose(isInitiator: boolean, earlyData?: bytes) : Promise<NoiseSession> {
-    const ns = await this.xx.initSession(isInitiator, this.prologue, this.staticKeys, this.remotePublicKey);
+  async propose(earlyData?: bytes) : Promise<NoiseSession> {
+    const ns = await this.xx.initSession(this.isInitiator, this.prologue, this.staticKeys, this.remotePublicKey);
 
-    if (isInitiator) {
+    if (this.isInitiator) {
       const signedPayload = signPayload(this.staticKeys.privateKey, getHandshakePayload(this.staticKeys.publicKey));
       const handshakePayload = await createHandshakePayload(
         this.staticKeys.publicKey,
@@ -61,8 +65,8 @@ export class Handshake {
   }
 
   // stage 1
-  async exchange(isInitiator: boolean, session: NoiseSession) : Promise<void> {
-    if (isInitiator) {
+  async exchange(session: NoiseSession) : Promise<void> {
+    if (this.isInitiator) {
       const receivedMessageBuffer = (await this.connection.readLP()).slice();
       const plaintext = await this.xx.recvMessage(session, decodeMessageBuffer(receivedMessageBuffer));
     } else {
@@ -77,8 +81,8 @@ export class Handshake {
   }
 
   // stage 2
-  async finish(isInitiator: boolean, session: NoiseSession) : Promise<void> {
-    if (isInitiator) {
+  async finish(session: NoiseSession) : Promise<void> {
+    if (this.isInitiator) {
       const messageBuffer = await this.xx.sendMessage(session, Buffer.alloc(0));
       this.connection.writeLP(encodeMessageBuffer(messageBuffer));
     } else {
@@ -86,4 +90,26 @@ export class Handshake {
       const plaintext = await this.xx.recvMessage(session, decodeMessageBuffer(receivedMessageBuffer));
     }
   }
+
+  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 {
+    const cs = this.getCS(session, false);
+    return this.xx.decryptWithAd(cs, Buffer.alloc(0), ciphertext);
+  }
+
+  private getCS(session: NoiseSession, encryption = true) {
+    if (!session.cs1 || !session.cs2) {
+      throw new Error("Handshake not completed properly, cipher state does not exist.");
+    }
+
+    if (this.isInitiator) {
+      return encryption ? session.cs1 : session.cs2;
+    } else {
+      return encryption ? session.cs2 : session.cs1;
+    }
+  }
 }

src/noise.ts

@@ -4,7 +4,7 @@ import Wrap from 'it-pb-rpc';
 
 import { Handshake } from "./handshake";
 import { generateKeypair } from "./utils";
-import { decryptStreams, encryptStreams } from "./crypto";
+import { decryptStream, encryptStream } from "./crypto";
 import { bytes } from "./@types/basic";
 import { NoiseConnection, PeerId, KeyPair, SecureOutbound } from "./@types/libp2p";
 import { Duplex } from "./@types/it-pair";
@@ -71,14 +71,16 @@ export class Noise implements NoiseConnection {
     remotePublicKey: bytes,
     isInitiator: boolean,
     ) : Promise<Duplex> {
+    // Perform handshake
     const prologue = Buffer.from(this.protocol);
-    const handshake = new Handshake('XX', remotePublicKey, prologue, this.staticKeys, connection);
+    const handshake = new Handshake('XX', isInitiator, remotePublicKey, prologue, this.staticKeys, connection);
 
-    const session = await handshake.propose(isInitiator, this.earlyData);
-    await handshake.exchange(isInitiator, session);
-    await handshake.finish(isInitiator, session);
+    const session = await handshake.propose(this.earlyData);
+    await handshake.exchange(session);
+    await handshake.finish(session);
 
-    return await encryptStreams(connection, session);
+    // Create encryption box/unbox wrapper
+    return await encryptStream(handshake, session);
   }