From e5d51e08688c81b8506d78d69a03eef848de691c Mon Sep 17 00:00:00 2001
From: Spencer Pogorzelski <34356756+Scoder12@users.noreply.github.com>
Date: Sun, 13 Aug 2023 22:40:50 -0700
Subject: [PATCH] websocket progress

---
 client/src/AdriftClient.ts | 27 ++++++++++++++++++++++++++-
 client/src/Connection.ts   | 21 +++++++++++++++------
 frontend/src/App.svelte    | 10 ++++++++++
 server/src/server.ts       | 24 +++++++++++++++++++++++-
 4 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/client/src/AdriftClient.ts b/client/src/AdriftClient.ts
index 80088eb..c3790c4 100644
--- a/client/src/AdriftClient.ts
+++ b/client/src/AdriftClient.ts
@@ -67,6 +67,7 @@ export class AdriftBareClient extends Client {
     webSocketImpl: WebSocketImpl
   ): WebSocket {
     const ws = new webSocketImpl("ws:null", protocols);
+    const closeEmitter = new EventTarget();
     // this will error. that's okay
 
     let { send, close } = this.connection.wsconnect(
@@ -77,7 +78,7 @@ export class AdriftBareClient extends Client {
       },
       () => {
         onReadyState(WebSocket.CLOSED);
-        ws.dispatchEvent(new Event("close"));
+        closeEmitter.dispatchEvent(new Event("close"));
       },
       (data) => {
         ws.dispatchEvent(
@@ -101,6 +102,30 @@ export class AdriftBareClient extends Client {
         close(code, reason);
       }
     );
+    let onClose = ws.onclose;
+    (ws as any).__defineGetter__("onclose", () => onClose);
+    (ws as any).__defineSetter__("onclose", (newOnClose: any) => {
+      onClose = newOnClose;
+    });
+    closeEmitter.addEventListener("close", () => {
+      if (onClose) onClose.call(ws, new Event("close"));
+    });
+    const _addEventListener = ws.addEventListener.bind(ws);
+    (ws as any).addEventListener = (evt: string, cb: any, options: any) => {
+      if (evt == "close") {
+        closeEmitter.addEventListener("close", cb, options);
+        return;
+      }
+      _addEventListener(evt, cb, options);
+    };
+    const _removeEventListener = ws.removeEventListener.bind(ws);
+    (ws as any).removeEventListener = (evt: string, cb: any, options: any) => {
+      if (evt == "close") {
+        closeEmitter.removeEventListener("close", cb, options);
+        return;
+      }
+      _removeEventListener(evt, cb, options);
+    };
 
     return ws;
   }
diff --git a/client/src/Connection.ts b/client/src/Connection.ts
index 9301c77..7473d46 100644
--- a/client/src/Connection.ts
+++ b/client/src/Connection.ts
@@ -9,14 +9,13 @@ import {
   Transport,
 } from "protocol";
 
+type OpenWSMeta = { onclose: () => void; onmessage: (data: any) => void };
+
 export class Connection {
   requestCallbacks: Record<number, Function> = {};
   openRequestStreams: Record<number, ReadableStreamDefaultController<any>> = {};
-  openingSockets: Record<number, () => void>;
-  openSockets: Record<
-    number,
-    { onclose: () => void; onmessage: (data: any) => void }
-  >;
+  openingSockets: Record<number, { onopen: () => void; rest: OpenWSMeta }> = {};
+  openSockets: Record<number, OpenWSMeta> = {};
 
   counter: number = 0;
 
@@ -68,6 +67,16 @@ export class Connection {
         this.openRequestStreams[requestID]?.close();
         delete this.openRequestStreams[requestID];
         break;
+
+      case S2CRequestTypes.WSOpen:
+        const { onopen, rest } = this.openingSockets[requestID];
+        delete this.openingSockets[requestID];
+        this.openSockets[requestID] = rest;
+        onopen();
+        break;
+
+      default:
+        break;
     }
   }
 
@@ -126,7 +135,7 @@ export class Connection {
     });
 
     // this can't be async, just call onopen when opened
-    this.openingSockets[seq] = onopen;
+    this.openingSockets[seq] = { onopen, rest: { onmessage, onclose } };
 
     return {
       send: (data) => {
diff --git a/frontend/src/App.svelte b/frontend/src/App.svelte
index 5463ad5..223e2eb 100644
--- a/frontend/src/App.svelte
+++ b/frontend/src/App.svelte
@@ -157,6 +157,16 @@
   }
 
   (window as any).bare = new BareClient();
+  (window as any).myWsTest = () => {
+    const ws = ((window as any).ws = (
+      (window as any).bare as BareClient
+    ).createWebSocket("wss://ws.postman-echo.com/raw", [], {}));
+    ws.onopen = () => console.log("onopen");
+    ws.addEventListener("open", () => console.log("open listener"));
+    ws.onclose = () => console.error(new Error("onclose"));
+    ws.addEventListener("close", () => console.log("close listener"));
+    ws.onmessage = console.log;
+  };
 </script>
 
 {#if ready}
diff --git a/server/src/server.ts b/server/src/server.ts
index 87a1d9a..378e40c 100644
--- a/server/src/server.ts
+++ b/server/src/server.ts
@@ -29,7 +29,7 @@ function bareErrorToResponse(e: BareError): {
 
 export class AdriftServer {
   send: (msg: ArrayBuffer) => void;
-  sockets: Record<number, WebSocket>;
+  sockets: Record<number, WebSocket> = {};
   events: EventEmitter;
 
   constructor(send: (msg: ArrayBuffer) => void) {
@@ -240,6 +240,7 @@ export class AdriftServer {
         const payload = AdriftServer.tryParseJSONPayload(msg.slice(cursor));
         const ws = (this.sockets[seq] = new WebSocket(payload.url));
         ws.binaryType = "arraybuffer";
+        // TODO v important: onerror
         ws.onopen = () => {
           this.sendWSOpen(seq);
         };
@@ -282,6 +283,27 @@ export class AdriftServer {
         break;
       }
 
+      case C2SRequestTypes.WSSendText: {
+        const socket = this.sockets[seq];
+        if (!socket) return;
+        socket.send(new TextDecoder().decode(msg.slice(cursor)));
+        break;
+      }
+
+      case C2SRequestTypes.WSSendBinary: {
+        const socket = this.sockets[seq];
+        if (!socket) return;
+        socket.send(msg.slice(cursor));
+        break;
+      }
+
+      case C2SRequestTypes.WSClose: {
+        const socket = this.sockets[seq];
+        if (!socket) return;
+        socket.close();
+        break;
+      }
+
       default:
         // not implemented
         break;