From 2031418a24e6df65a600349db940a555b1220088 Mon Sep 17 00:00:00 2001
From: crowlKats <13135287+crowlKats@users.noreply.github.com>
Date: Thu, 26 Nov 2020 16:38:15 +0100
Subject: [PATCH] feat(std/ws): protocol & version support (#8505)

Co-authored-by: Tom Wieland <tom.wieland@gmail.com>
---
 std/ws/mod.ts  | 19 ++++++++++++++-----
 std/ws/test.ts | 41 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/std/ws/mod.ts b/std/ws/mod.ts
index 1e40c288bd..a86c184827 100644
--- a/std/ws/mod.ts
+++ b/std/ws/mod.ts
@@ -427,13 +427,22 @@ export async function acceptWebSocket(req: {
       throw new Error("sec-websocket-key is not provided");
     }
     const secAccept = createSecAccept(secKey);
+    const newHeaders = new Headers({
+      Upgrade: "websocket",
+      Connection: "Upgrade",
+      "Sec-WebSocket-Accept": secAccept,
+    });
+    const secProtocol = headers.get("sec-websocket-protocol");
+    if (typeof secProtocol === "string") {
+      newHeaders.set("Sec-WebSocket-Protocol", secProtocol);
+    }
+    const secVersion = headers.get("sec-websocket-version");
+    if (typeof secVersion === "string") {
+      newHeaders.set("Sec-WebSocket-Version", secVersion);
+    }
     await writeResponse(bufWriter, {
       status: 101,
-      headers: new Headers({
-        Upgrade: "websocket",
-        Connection: "Upgrade",
-        "Sec-WebSocket-Accept": secAccept,
-      }),
+      headers: newHeaders,
     });
     return sock;
   }
diff --git a/std/ws/test.ts b/std/ws/test.ts
index e8fa4f8b13..9b7e7a710c 100644
--- a/std/ws/test.ts
+++ b/std/ws/test.ts
@@ -1,10 +1,16 @@
 // Copyright 2018-2020 the Deno authors. All rights reserved. MIT license.
 import { BufReader, BufWriter } from "../io/bufio.ts";
-import { assert, assertEquals, assertThrowsAsync } from "../testing/asserts.ts";
+import {
+  assert,
+  assertEquals,
+  assertThrowsAsync,
+  fail,
+} from "../testing/asserts.ts";
 import { TextProtoReader } from "../textproto/mod.ts";
 import * as bytes from "../bytes/mod.ts";
 import {
   acceptable,
+  acceptWebSocket,
   createSecAccept,
   createSecKey,
   createWebSocket,
@@ -16,6 +22,8 @@ import {
 } from "./mod.ts";
 import { decode, encode } from "../encoding/utf8.ts";
 import { delay } from "../async/delay.ts";
+import { serve } from "../http/server.ts";
+import { deferred } from "../async/deferred.ts";
 
 Deno.test("[ws] read unmasked text frame", async () => {
   // unmasked single text frame with payload "Hello"
@@ -454,3 +462,34 @@ Deno.test("[ws] WebSocket should act as asyncIterator", async () => {
   assertEquals(events[1], "Hello");
   assertEquals(events[2], { code: 1011, reason: "42" });
 });
+
+Deno.test("[ws] WebSocket protocol", async () => {
+  const promise = deferred();
+  const server = serve({ port: 5839 });
+
+  const ws = new WebSocket("ws://localhost:5839", ["foo", "bar"]);
+  ws.onopen = () => {
+    assertEquals(ws.protocol, "foo, bar");
+    ws.close();
+  };
+  ws.onerror = () => fail();
+  ws.onclose = () => {
+    server.close();
+    promise.resolve();
+  };
+
+  const x = await server[Symbol.asyncIterator]().next();
+  if (!x.done) {
+    const { conn, r: bufReader, w: bufWriter, headers } = x.value;
+    await acceptWebSocket({
+      conn,
+      bufReader,
+      bufWriter,
+      headers,
+    });
+
+    await promise;
+  } else {
+    fail();
+  }
+});