From 3b5a401a1c6456cd1b9fd2df473d052c38f8b28a Mon Sep 17 00:00:00 2001
From: morrigan <belma.dzi@gmail.com>
Date: Wed, 6 Nov 2019 16:09:15 +0100
Subject: [PATCH] Complete test

---
 src/xx.ts       |  5 +++--
 test/xx.test.ts | 12 ++++++++----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/xx.ts b/src/xx.ts
index 4947f63..aac3850 100644
--- a/src/xx.ts
+++ b/src/xx.ts
@@ -127,13 +127,14 @@ export class XXHandshake {
     cs.n = nonce;
   }
 
-  private encryptWithAd(cs: CipherState, ad: bytes, plaintext: bytes) : bytes {
+  public encryptWithAd(cs: CipherState, ad: bytes, plaintext: bytes) : bytes {
     const e = this.encrypt(cs.k, cs.n, ad, plaintext);
     this.setNonce(cs, this.incrementNonce(cs.n));
+
     return e;
   }
 
-  private decryptWithAd(cs: CipherState, ad: bytes, ciphertext: bytes) : bytes {
+  public decryptWithAd(cs: CipherState, ad: bytes, ciphertext: bytes) : bytes {
     const plaintext = this.decrypt(cs.k, cs.n, ad, ciphertext);
     this.setNonce(cs, this.incrementNonce(cs.n));
 
diff --git a/test/xx.test.ts b/test/xx.test.ts
index c5232b8..6eabb45 100644
--- a/test/xx.test.ts
+++ b/test/xx.test.ts
@@ -36,8 +36,7 @@ describe("Index", () => {
     return await crypto.keys.generateKeyPair('ed25519');
   }
 
-  async function doHandshake() {
-    const xx = new XXHandshake();
+  async function doHandshake(xx) {
     const kpInit = await xx.generateKeypair();
     const kpResp = await xx.generateKeypair();
     const payloadString = Buffer.from("noise-libp2p-static-key:");
@@ -99,14 +98,19 @@ describe("Index", () => {
 
     // initiator send message
     const messageBuffer3 = await xx.sendMessage(nsInit, Buffer.alloc(0));
+
     // responder receive message
     const plaintext3 = await xx.recvMessage(nsResp, messageBuffer3);
     console.log("Stage 2 responder payload: ", plaintext3);
 
+    assert(nsInit.cs1.k.equals(nsResp.cs1.k));
+    assert(nsInit.cs2.k.equals(nsResp.cs2.k));
+
     return { nsInit, nsResp };
   }
 
   it("Test handshake", async () => {
-    await doHandshake();
-  })
+    const xx = new XXHandshake();
+    await doHandshake(xx);
+  });
 });