From f794ae96d61e3581c577ff423e60d990dfd9925f Mon Sep 17 00:00:00 2001
From: cgeek <cem.moreau@gmail.com>
Date: Tue, 5 Sep 2017 17:04:22 +0200
Subject: [PATCH] [enh] #1084 WS2P: wrapping with WS2PCluster

---
 app/lib/ws2p/WS2PCluster.ts       | 36 +++++++++++++++++++++++++++++++
 app/lib/ws2p/WS2PConnection.ts    |  4 ++++
 app/lib/ws2p/WS2PServer.ts        |  7 +++++-
 test/integration/tools/toolbox.ts | 12 ++++++-----
 4 files changed, 53 insertions(+), 6 deletions(-)
 create mode 100644 app/lib/ws2p/WS2PCluster.ts

diff --git a/app/lib/ws2p/WS2PCluster.ts b/app/lib/ws2p/WS2PCluster.ts
new file mode 100644
index 000000000..bc774ea42
--- /dev/null
+++ b/app/lib/ws2p/WS2PCluster.ts
@@ -0,0 +1,36 @@
+import {WS2PServer} from "./WS2PServer"
+import {Server} from "../../../server"
+import {WS2PClient} from "./WS2PClient"
+import {WS2PConnection} from "./WS2PConnection"
+
+const nuuid = require('node-uuid')
+
+export class WS2PCluster {
+
+  private ws2pServer:WS2PServer|null = null
+  private ws2pClients:{[k:string]:WS2PClient} = {}
+
+  constructor(private server:Server) {}
+
+  async listen(host:string, port:number) {
+    if (this.ws2pServer) {
+      await this.ws2pServer.close()
+    }
+    this.ws2pServer = await WS2PServer.bindOn(this.server, host, port)
+    return this.ws2pServer
+  }
+
+  clientsCount() {
+    return Object.keys(this.ws2pClients).length
+  }
+
+  async connect(host: string, port: number): Promise<WS2PConnection> {
+    const uuid = nuuid.v4()
+    const ws2pc = await WS2PClient.connectTo(this.server, host, port)
+    this.ws2pClients[uuid] = ws2pc
+    ws2pc.closed.then(() => {
+      delete this.ws2pClients[uuid]
+    })
+    return ws2pc
+  }
+}
\ No newline at end of file
diff --git a/app/lib/ws2p/WS2PConnection.ts b/app/lib/ws2p/WS2PConnection.ts
index aac8b5f1a..9d598f47c 100644
--- a/app/lib/ws2p/WS2PConnection.ts
+++ b/app/lib/ws2p/WS2PConnection.ts
@@ -324,6 +324,10 @@ export class WS2PConnection {
     return this.onWsClosed
   }
 
+  close() {
+    return this.ws.close()
+  }
+
   async connect() {
     if (!this.connectp) {
       this.connectp = (async () => {
diff --git a/app/lib/ws2p/WS2PServer.ts b/app/lib/ws2p/WS2PServer.ts
index 79861876f..e02ee6ff4 100644
--- a/app/lib/ws2p/WS2PServer.ts
+++ b/app/lib/ws2p/WS2PServer.ts
@@ -17,6 +17,10 @@ export class WS2PServer {
     private port:number) {
   }
 
+  getConnexions() {
+    return this.connections.slice()
+  }
+
   private listenToWebSocketConnections() {
     const key = new Key(this.server.conf.pair.pub, this.server.conf.pair.sec)
     this.wss = new WebSocketServer({ host: this.host, port: this.port })
@@ -52,7 +56,8 @@ export class WS2PServer {
     })
   }
 
-  close() {
+  async close() {
+    await Promise.all(this.connections.map(c => c.close()))
     return this.wss.close()
   }
 
diff --git a/test/integration/tools/toolbox.ts b/test/integration/tools/toolbox.ts
index 55d99ab3c..9d6e46032 100644
--- a/test/integration/tools/toolbox.ts
+++ b/test/integration/tools/toolbox.ts
@@ -19,8 +19,8 @@ import {Key} from "../../../app/lib/common-libs/crypto/keyring"
 import {WS2PConnection, WS2PPubkeyLocalAuth, WS2PPubkeyRemoteAuth} from "../../../app/lib/ws2p/WS2PConnection"
 import {WS2PResponse} from "../../../app/lib/ws2p/impl/WS2PResponse"
 import {WS2PMessageHandler} from "../../../app/lib/ws2p/impl/WS2PMessageHandler"
+import {WS2PCluster} from "../../../app/lib/ws2p/WS2PCluster"
 import {WS2PServer} from "../../../app/lib/ws2p/WS2PServer"
-import {WS2PClient} from "../../../app/lib/ws2p/WS2PClient"
 
 const assert      = require('assert');
 const _           = require('underscore');
@@ -609,13 +609,15 @@ export async function newWS2PBidirectionnalConnection(k1:Key, k2:Key, serverHand
   })
 }
 
-export const simpleWS2PNetwork = async (s1:TestingServer, s2:TestingServer) => {
+export const simpleWS2PNetwork: (s1: TestingServer, s2: TestingServer) => Promise<{ w1: WS2PConnection; ws2pc: WS2PConnection; wss: WS2PServer }> = async (s1: TestingServer, s2: TestingServer) => {
   let port = PORT++
   const clientPub = s2.conf.pair.pub
-  let w1:WS2PConnection|null
+  let w1: WS2PConnection | null
 
-  const ws2ps = await WS2PServer.bindOn(s1._server, 'localhost', port)
-  const ws2pc = await WS2PClient.connectTo(s2._server, 'localhost', port)
+  const cluster1 = new WS2PCluster(s1._server)
+  const cluster2 = new WS2PCluster(s2._server)
+  const ws2ps = await cluster1.listen('localhost', port)
+  const ws2pc = await cluster2.connect('localhost', port)
 
   w1 = await ws2ps.getConnection(clientPub)
   if (!w1) {
-- 
GitLab