chore(core-p2p): terminate on ws.ping() ws.pong() (#4194)

ArkEcosystem · Nov 24, 2020 · 9f9178d · 9f9178d
1 parent 21eb71a
commit 9f9178d
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 2 deletions.
diff --git a/__tests__/unit/core-p2p/hapi-nes/client.test.ts b/__tests__/unit/core-p2p/hapi-nes/client.test.ts
@@ -1004,5 +1004,38 @@ describe("Client", () => {
                 await server.stop();
             });
         });
+
+        describe("ping / pong", () => {
+            it.each([["ping"], ["pong"]])("terminates when receiving a ws.%s", async (method) => {
+                const server = await createServerWithPlugin({}, {}, true);
+
+                server.route({
+                    method: "POST",
+                    path: "/",
+                    handler: async (request) => {
+                        request.server.plugins.nes._listener._sockets._forEach((socket) => {
+                            setTimeout(() => socket._ws[method](), 100);
+                        });
+
+                        return "hello";
+                    },
+                });
+
+                await server.start();
+                const client = new Client("http://localhost:" + server.info.port);
+
+                await client.connect();
+
+                await client.request("/");
+
+                await Hoek.wait(500);
+
+                //@ts-ignore
+                expect(client._ws).toBeNull(); // null because _cleanup() in reconnect() method
+
+                await client.disconnect();
+                await server.stop();
+            });
+        })
     });
 });
diff --git a/__tests__/unit/core-p2p/hapi-nes/socket.test.ts b/__tests__/unit/core-p2p/hapi-nes/socket.test.ts
@@ -402,6 +402,35 @@ describe("Socket", () => {
             client.close();
             await server.stop();
         });
+
+        it.each([["ping"], ["pong"]])("terminates on ws.%s", async (method) => {
+            const server = Hapi.server();
+            await server.register({ plugin: plugin, options: {} });
+
+            server.route({
+                method: "POST",
+                path: "/",
+                handler: () => "hello",
+            });
+
+            await server.start();
+            const client = new Ws("http://localhost:" + server.info.port);
+            client.onerror = Hoek.ignore;
+
+            const sendPingOrPong = async () => new Promise((resolve, reject) => {
+                client.on("open", () => {
+                    client[method]("", true, () => resolve());
+                });
+            })
+
+            await sendPingOrPong();
+            await delay(1000);
+
+            expect(client.readyState).toEqual(client.CLOSED);
+
+            client.close();
+            await server.stop();
+        });
     });
 
     describe("_processRequest()", () => {

diff --git a/packages/core-p2p/src/hapi-nes/client.ts b/packages/core-p2p/src/hapi-nes/client.ts
@@ -282,9 +282,12 @@ export class Client {
         ws.onmessage = (message) => {
             return this._onMessage(message);
         };
+
+        ws.on("ping", () => this._disconnect(() => {}, true, true));
+        ws.on("pong", () => this._disconnect(() => {}, true, true));
     }
 
-    private _disconnect(next, isInternal) {
+    private _disconnect(next, isInternal, terminate = false) {
         this._reconnection = null;
         clearTimeout(this._reconnectionTimer);
         this._reconnectionTimer = null;
@@ -302,7 +305,12 @@ export class Client {
 
         this._disconnectRequested = requested;
         this._disconnectListeners = [next];
-        this._ws.close();
+
+        if (terminate) {
+            this._ws.terminate();
+        } else {
+            this._ws.close();
+        }
     }
 
     private _cleanup() {

diff --git a/packages/core-p2p/src/hapi-nes/socket.ts b/packages/core-p2p/src/hapi-nes/socket.ts
@@ -47,6 +47,8 @@ export class Socket {
         };
 
         this._ws.on("message", (message) => this._onMessage(message));
+        this._ws.on("ping", () => this.terminate());
+        this._ws.on("pong", () => this.terminate());
     }
 
     public disconnect() {