From 98dec1addb59c2a215cb4425a4927c5f33a78b4c Mon Sep 17 00:00:00 2001
From: enisdenjo <badurinadenis@gmail.com>
Date: Wed, 23 Feb 2022 15:36:12 +0100
Subject: [PATCH] fix(server): `handleProtocols` accepts arrays too and
 gracefully rejects other types

Closes #318
---
 docs/modules/server.md  |  2 +-
 src/__tests__/server.ts | 66 ++++++++++++++++++++++++++++++++++++++++-
 src/server.ts           | 26 +++++++++-------
 3 files changed, 81 insertions(+), 13 deletions(-)

diff --git a/docs/modules/server.md b/docs/modules/server.md
index 7bf019c8..2e5df074 100644
--- a/docs/modules/server.md
+++ b/docs/modules/server.md
@@ -42,7 +42,7 @@ to handle the connection afterwards.
 
 | Name | Type |
 | :------ | :------ |
-| `protocols` | `string` \| `Set`<`string`\> |
+| `protocols` | `string` \| `Set`<`string`\> \| `string`[] |
 
 #### Returns
 
diff --git a/src/__tests__/server.ts b/src/__tests__/server.ts
index 7155d4f5..6703a9e0 100644
--- a/src/__tests__/server.ts
+++ b/src/__tests__/server.ts
@@ -8,7 +8,7 @@ import {
   ExecutionResult,
   GraphQLSchema,
 } from 'graphql';
-import { makeServer } from '../server';
+import { handleProtocols, makeServer } from '../server';
 import {
   GRAPHQL_TRANSPORT_WS_PROTOCOL,
   CloseCode,
@@ -1705,3 +1705,67 @@ describe('Disconnect/close', () => {
     client.ws.close(4321, 'Byebye');
   });
 });
+
+it('should only accept a Set, Array or string in handleProtocols', () => {
+  for (const test of [
+    {
+      in: new Set(['not', 'me']),
+      out: false,
+    },
+    {
+      in: new Set(['maybe', 'me', GRAPHQL_TRANSPORT_WS_PROTOCOL + 'nah']),
+      out: false,
+    },
+    {
+      in: new Set(['almost', 'next', GRAPHQL_TRANSPORT_WS_PROTOCOL, 'one']),
+      out: GRAPHQL_TRANSPORT_WS_PROTOCOL,
+    },
+    {
+      in: [''],
+      out: false,
+    },
+    {
+      in: ['123', GRAPHQL_TRANSPORT_WS_PROTOCOL],
+      out: GRAPHQL_TRANSPORT_WS_PROTOCOL,
+    },
+    {
+      in: [GRAPHQL_TRANSPORT_WS_PROTOCOL, GRAPHQL_TRANSPORT_WS_PROTOCOL],
+      out: GRAPHQL_TRANSPORT_WS_PROTOCOL,
+    },
+    {
+      in: `some, ${GRAPHQL_TRANSPORT_WS_PROTOCOL}   , other-one,third`,
+      out: GRAPHQL_TRANSPORT_WS_PROTOCOL,
+    },
+    {
+      in: `no, graphql-TransPort-ws`,
+      out: false,
+    },
+    {
+      in: { iAm: 'unacceptable' },
+      out: false,
+    },
+    {
+      in: 123,
+      out: false,
+    },
+    {
+      in: null,
+      out: false,
+    },
+    {
+      in: undefined,
+      out: false,
+    },
+    {
+      in: () => {
+        // void
+      },
+      out: false,
+    },
+  ]) {
+    expect(
+      // @ts-expect-error for test purposes, in can be different from type
+      handleProtocols(test.in),
+    ).toBe(test.out);
+  }
+});
diff --git a/src/server.ts b/src/server.ts
index 2a2cbb4b..6f02ba1b 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -862,16 +862,20 @@ export function makeServer<
  * to handle the connection afterwards.
  */
 export function handleProtocols(
-  protocols: Set<string> | string,
+  protocols: Set<string> | string[] | string,
 ): typeof GRAPHQL_TRANSPORT_WS_PROTOCOL | false {
-  return (
-    typeof protocols === 'string'
-      ? protocols
-          .split(',')
-          .map((p) => p.trim())
-          .includes(GRAPHQL_TRANSPORT_WS_PROTOCOL)
-      : protocols.has(GRAPHQL_TRANSPORT_WS_PROTOCOL)
-  )
-    ? GRAPHQL_TRANSPORT_WS_PROTOCOL
-    : false;
+  switch (true) {
+    case protocols instanceof Set &&
+      protocols.has(GRAPHQL_TRANSPORT_WS_PROTOCOL):
+    case Array.isArray(protocols) &&
+      protocols.includes(GRAPHQL_TRANSPORT_WS_PROTOCOL):
+    case typeof protocols === 'string' &&
+      protocols
+        .split(',')
+        .map((p) => p.trim())
+        .includes(GRAPHQL_TRANSPORT_WS_PROTOCOL):
+      return GRAPHQL_TRANSPORT_WS_PROTOCOL;
+    default:
+      return false;
+  }
 }