From eef68f7d5c5ed5f0757664c1ea088c8f7471702b Mon Sep 17 00:00:00 2001
From: Spencer Pogorzelski <34356756+Scoder12@users.noreply.github.com>
Date: Sat, 19 Aug 2023 13:14:23 -0700
Subject: [PATCH] populate WS protocol on open

---
 client/src/AdriftClient.ts |  3 ++-
 client/src/Connection.ts   | 12 ++++++++----
 protocol/src/index.ts      |  4 ++++
 server/src/server.ts       |  7 ++++---
 4 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/client/src/AdriftClient.ts b/client/src/AdriftClient.ts
index 472e099..c17a5e2 100644
--- a/client/src/AdriftClient.ts
+++ b/client/src/AdriftClient.ts
@@ -182,8 +182,9 @@ export class AdriftBareClient extends Client {
     let { send, close } = this.connection.wsconnect(
       remote,
       protocols,
-      () => {
+      (protocol: string) => {
         onReadyState(WebSocket.OPEN);
+        (ws as any).protocol = protocol;
         ws.dispatchEvent(new Event("open"));
       },
       (code: number, reason: string, wasClean: boolean) => {
diff --git a/client/src/Connection.ts b/client/src/Connection.ts
index 4f3cca6..16cf1a2 100644
--- a/client/src/Connection.ts
+++ b/client/src/Connection.ts
@@ -9,6 +9,7 @@ import {
   ProtoBareHeaders,
   S2CRequestType,
   S2CRequestTypes,
+  S2CWSOpenPayload,
   S2C_HELLO_ERR,
   S2C_HELLO_OK,
   Transport,
@@ -17,7 +18,7 @@ import {
 } from "protocol";
 
 type OpenWSMeta = {
-  onopen: () => void;
+  onopen: (protocol: string) => void;
   onclose: (code: number, reason: string, wasClean: boolean) => void;
   onmessage: (data: ReadableStream, isBinary: boolean) => void;
   onerror: (message: string) => void;
@@ -126,12 +127,15 @@ export class Connection {
         delete this.openRequestStreams[requestID];
         break;
 
-      case S2CRequestTypes.WSOpen:
+      case S2CRequestTypes.WSOpen: {
         const socketMeta = this.openingSockets[requestID];
+        if (!socketMeta) return;
+        const payload: S2CWSOpenPayload = msgJSON();
         delete this.openingSockets[requestID];
         this.openSockets[requestID] = socketMeta;
-        setTimeout(() => socketMeta.onopen());
+        setTimeout(() => socketMeta.onopen(payload.protocol));
         break;
+      }
 
       case S2CRequestTypes.WSBinaryStart:
       case S2CRequestTypes.WSTextStart: {
@@ -270,7 +274,7 @@ export class Connection {
   wsconnect(
     url: URL,
     protocols: string | string[],
-    onopen: () => void,
+    onopen: (protocol: string) => void,
     onclose: (code: number, reason: string, wasClean: boolean) => void,
     onmessage: (data: ReadableStream, isBinary: boolean) => void,
     onerror: (message: string) => void,
diff --git a/protocol/src/index.ts b/protocol/src/index.ts
index 282842a..c27dcdf 100644
--- a/protocol/src/index.ts
+++ b/protocol/src/index.ts
@@ -44,6 +44,10 @@ export type C2SWSOpenPayload = {
   protocols: string | string[];
 };
 
+export type S2CWSOpenPayload = {
+  protocol: string;
+};
+
 export type WSClosePayload = {
   code: number;
   reason: string;
diff --git a/server/src/server.ts b/server/src/server.ts
index a8d58ac..cfb5c77 100644
--- a/server/src/server.ts
+++ b/server/src/server.ts
@@ -11,6 +11,7 @@ import {
   ProtoBareHeaders,
   S2CRequestType,
   S2CRequestTypes,
+  S2CWSOpenPayload,
   S2C_HELLO_ERR,
   S2C_HELLO_OK,
   WSClosePayload,
@@ -199,8 +200,8 @@ export class AdriftServer {
     this._sendSimpleRes(seq, S2CRequestTypes.HTTPResponseEnd);
   }
 
-  sendWSOpen(seq: number) {
-    this._sendSimpleRes(seq, S2CRequestTypes.WSOpen);
+  sendWSOpen(seq: number, payload: S2CWSOpenPayload) {
+    this._sendJSONRes(seq, S2CRequestTypes.WSOpen, payload);
   }
 
   sendWSClose(seq: number, payload: WSClosePayload) {
@@ -318,7 +319,7 @@ export class AdriftServer {
           // onclose will be called after this with code 1006, reason "" and wasClean false
         };
         ws.onopen = () => {
-          this.sendWSOpen(seq);
+          this.sendWSOpen(seq, { protocol: ws.protocol });
         };
         ws.onclose = (e) => {
           this.sendWSClose(seq, {