feat(server): Optional onPing and onPong message type listeners

enisdenjo · Jun 21, 2021 · f36066f · f36066f
1 parent 79bd429
commit f36066f
Show file tree

Hide file tree

Showing 3 changed files with 122 additions and 1 deletion.
diff --git a/docs/interfaces/server.websocket.md b/docs/interfaces/server.websocket.md
@@ -14,6 +14,8 @@
 
 - [close](server.websocket.md#close)
 - [onMessage](server.websocket.md#onmessage)
+- [onPing](server.websocket.md#onping)
+- [onPong](server.websocket.md#onpong)
 - [send](server.websocket.md#send)
 
 ## Properties
@@ -79,6 +81,50 @@ to your clients however you wish.
 
 ___
 
+### onPing
+
+▸ `Optional` **onPing**(`payload`): `void` \| `Promise`<void\>
+
+Implement a listener for the `PingMessage` sent from the client to the server.
+If the client sent the ping with a payload, it will be passed through the
+first argument.
+
+If this listener is implemented, the server will NOT automatically reply
+to any pings from the client. Implementing it makes it your resposibility
+to decide how and when to respond.
+
+#### Parameters
+
+| Name | Type |
+| :------ | :------ |
+| `payload` | `undefined` \| `Record`<string, unknown\> |
+
+#### Returns
+
+`void` \| `Promise`<void\>
+
+___
+
+### onPong
+
+▸ `Optional` **onPong**(`payload`): `void` \| `Promise`<void\>
+
+Implement a listener for the `PongMessage` sent from the client to the server.
+If the client sent the pong with a payload, it will be passed through the
+first argument.
+
+#### Parameters
+
+| Name | Type |
+| :------ | :------ |
+| `payload` | `undefined` \| `Record`<string, unknown\> |
+
+#### Returns
+
+`void` \| `Promise`<void\>
+
+___
+
 ### send
 
 ▸ **send**(`data`): `void` \| `Promise`<void\>

diff --git a/src/server.ts b/src/server.ts
@@ -33,6 +33,8 @@ import {
   CompleteMessage,
   JSONMessageReplacer,
   JSONMessageReviver,
+  PingMessage,
+  PongMessage,
 } from './common';
 import { isObject, isAsyncIterable, areGraphQLErrors } from './utils';
 
@@ -449,6 +451,22 @@ export interface WebSocket {
    * to your clients however you wish.
    */
   onMessage(cb: (data: string) => Promise<void>): void;
+  /**
+   * Implement a listener for the `PingMessage` sent from the client to the server.
+   * If the client sent the ping with a payload, it will be passed through the
+   * first argument.
+   *
+   * If this listener is implemented, the server will NOT automatically reply
+   * to any pings from the client. Implementing it makes it your resposibility
+   * to decide how and when to respond.
+   */
+  onPing?(payload: PingMessage['payload']): Promise<void> | void;
+  /**
+   * Implement a listener for the `PongMessage` sent from the client to the server.
+   * If the client sent the pong with a payload, it will be passed through the
+   * first argument.
+   */
+  onPong?(payload: PongMessage['payload']): Promise<void> | void;
 }
 
 /** @category Server */
@@ -586,6 +604,10 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
             return;
           }
           case MessageType.Ping: {
+            if (socket.onPing)
+              // if the onPing listener is registered, automatic pong is disabled
+              return await socket.onPing(message.payload);
+
             await socket.send(
               stringifyMessage(
                 message.payload
@@ -599,7 +621,7 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
             return;
           }
           case MessageType.Pong:
-            return;
+            return await socket.onPong?.(message.payload);
           case MessageType.Subscribe: {
             if (!ctx.acknowledged) return socket.close(4401, 'Unauthorized');
 

diff --git a/src/tests/server.ts b/src/tests/server.ts
@@ -8,6 +8,7 @@ import {
   ExecutionResult,
   GraphQLSchema,
 } from 'graphql';
+import { makeServer } from '../server';
 import {
   GRAPHQL_TRANSPORT_WS_PROTOCOL,
   MessageType,
@@ -682,6 +683,58 @@ describe('Ping/Pong', () => {
       fail('Shouldt have closed');
     }, 20);
   });
+
+  it('should invoke the websocket callback on ping and not reply automatically', async (done) => {
+    const payload = { not: 'relevant' };
+
+    const closed = makeServer({}).opened(
+      {
+        protocol: GRAPHQL_TRANSPORT_WS_PROTOCOL,
+        send: () => fail('Shouldnt have responded to a ping'),
+        close: () => {
+          /**/
+        },
+        onMessage: (cb) => {
+          cb(stringifyMessage({ type: MessageType.Ping, payload }));
+        },
+        onPing: (pyld) => {
+          setImmediate(() => {
+            expect(pyld).toEqual(payload);
+            closed(1000, '');
+            done();
+          });
+        },
+        onPong: () => fail('Nothing shouldve ponged'),
+      },
+      {},
+    );
+  });
+
+  it('should invoke the websocket callback on pong', async (done) => {
+    const payload = { not: 'relevant' };
+
+    const closed = makeServer({}).opened(
+      {
+        protocol: GRAPHQL_TRANSPORT_WS_PROTOCOL,
+        send: () => Promise.resolve(),
+        close: () => {
+          /**/
+        },
+        onMessage: (cb) => {
+          cb(stringifyMessage({ type: MessageType.Pong, payload }));
+        },
+        onPing: () => fail('Nothing shouldve pinged'),
+        onPong: (pyld) => {
+          setImmediate(() => {
+            expect(pyld).toEqual(payload);
+            closed(1000, '');
+            done();
+          });
+        },
+      },
+      {},
+    );
+  });
 });
 
 describe('Subscribe', () => {