fix(uWebSockets): Handle premature and abrupt socket closes

Closes #186
enisdenjo · May 28, 2021 · 9d3ff52 · 9d3ff52
1 parent 5e37be8
commit 9d3ff52
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 42 deletions.
diff --git a/src/tests/server.ts b/src/tests/server.ts
@@ -21,47 +21,6 @@ import { createTClient, startWSTServer as startTServer } from './utils';
  * Tests
  */
 
-it('should allow connections with valid protocols only', async () => {
-  const { url } = await startTServer();
-
-  const warn = console.warn;
-  console.warn = () => {
-    /* hide warnings for test */
-  };
-
-  let client = await createTClient(url, '');
-  await client.waitForClose((event) => {
-    expect(event.code).toBe(1002);
-    expect(event.reason).toBe('Protocol Error');
-    expect(event.wasClean).toBeTruthy();
-  });
-
-  client = await createTClient(url, ['graphql', 'json']);
-  await client.waitForClose((event) => {
-    expect(event.code).toBe(1002);
-    expect(event.reason).toBe('Protocol Error');
-    expect(event.wasClean).toBeTruthy();
-  });
-
-  client = await createTClient(
-    url,
-    GRAPHQL_TRANSPORT_WS_PROTOCOL + 'gibberish',
-  );
-  await client.waitForClose((event) => {
-    expect(event.code).toBe(1002);
-    expect(event.reason).toBe('Protocol Error');
-    expect(event.wasClean).toBeTruthy();
-  });
-
-  client = await createTClient(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
-  await client.waitForClose(
-    () => fail('shouldnt close for valid protocol'),
-    30, // should be kicked off within this time
-  );
-
-  console.warn = warn;
-});
-
 it('should use the schema resolved from a promise on subscribe', async (done) => {
   expect.assertions(2);
 

diff --git a/src/tests/use.ts b/src/tests/use.ts
@@ -5,11 +5,53 @@ import {
   stringifyMessage,
   parseMessage,
   SubscribePayload,
+  GRAPHQL_TRANSPORT_WS_PROTOCOL,
 } from '../common';
 import { createTClient, tServers, WSExtra, UWSExtra } from './utils';
 
 for (const { tServer, startTServer } of tServers) {
   describe(tServer, () => {
+    it('should allow connections with valid protocols only', async () => {
+      const { url } = await startTServer();
+
+      const warn = console.warn;
+      console.warn = () => {
+        /* hide warnings for test */
+      };
+
+      let client = await createTClient(url, '');
+      await client.waitForClose((event) => {
+        expect(event.code).toBe(1002);
+        expect(event.reason).toBe('Protocol Error');
+        expect(event.wasClean).toBeTruthy();
+      });
+
+      client = await createTClient(url, ['graphql', 'json']);
+      await client.waitForClose((event) => {
+        expect(event.code).toBe(1002);
+        expect(event.reason).toBe('Protocol Error');
+        expect(event.wasClean).toBeTruthy();
+      });
+
+      client = await createTClient(
+        url,
+        GRAPHQL_TRANSPORT_WS_PROTOCOL + 'gibberish',
+      );
+      await client.waitForClose((event) => {
+        expect(event.code).toBe(1002);
+        expect(event.reason).toBe('Protocol Error');
+        expect(event.wasClean).toBeTruthy();
+      });
+
+      client = await createTClient(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
+      await client.waitForClose(
+        () => fail('shouldnt close for valid protocol'),
+        30, // should be kicked off within this time
+      );
+
+      console.warn = warn;
+    });
+
     it('should gracefully go away when disposing', async () => {
       const server = await startTServer();
 

diff --git a/src/use/uWebSockets.ts b/src/use/uWebSockets.ts
@@ -98,12 +98,18 @@ export function makeBehavior(
         {
           protocol: request.getHeader('sec-websocket-protocol'),
           send: async (message) => {
+            // the socket might have been destroyed in the meantime
+            if (!clients.has(socket)) return;
             if (!socket.send(message))
               // if backpressure is built up wait for drain
               await new Promise<void>((resolve) => (onDrain = resolve));
           },
           close: (code, reason) => {
-            socket.end(code, reason);
+            // end socket in next tick making sure the client is registered
+            setImmediate(() => {
+              // the socket might have been destroyed before issuing a close
+              if (clients.has(socket)) socket.end(code, reason);
+            });
           },
           onMessage: (cb) => (client.handleMessage = cb),
         },