From 73e6830b1fa27b1b3b1050c671ef6442b11805ba Mon Sep 17 00:00:00 2001
From: ProgrammerIn-wonderland <3838shah@gmail.com>
Date: Fri, 16 May 2025 13:45:55 -0400
Subject: [PATCH] Add reauthentication for unauthenticated web clients

---
 .../src/modules/networking/PSocket.js         | 62 ++++++++++++++-----
 .../src/modules/networking/PWispHandler.js    |  4 ++
 2 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/src/puter-js/src/modules/networking/PSocket.js b/src/puter-js/src/modules/networking/PSocket.js
index bfd413ed1..3f09a63f0 100644
--- a/src/puter-js/src/modules/networking/PSocket.js
+++ b/src/puter-js/src/modules/networking/PSocket.js
@@ -1,5 +1,6 @@
 import EventListener from "../../lib/EventListener.js";
 import { errors } from "./parsers.js";
+import { PWispHandler } from "./PWispHandler.js";
 const texten = new TextEncoder();
 
 export let wispInfo = {
@@ -12,23 +13,52 @@ export class PSocket extends EventListener {
     _streamID;
     constructor(host, port) {
         super(["data", "drain", "open", "error", "close", "tlsdata", "tlsopen"]);
-        const callbacks = {
-            dataCallBack: (data) => {
-                this.emit("data", data);
-            },
-            closeCallBack: (reason) => {
-                if (reason !== 0x02) {
-                    this.emit("error", new Error(errors[reason]));
-                    this.emit("close", true);
-                    return;    
-                }
-                this.emit("close", false);
-            }
-        }
 
-        this._streamID = wispInfo.handler.register(host, port, callbacks);
-        setTimeout(() => {this.emit("open", undefined)}, 0);
-        
+        (async () => {
+            if(!puter.authToken && puter.env === 'web'){
+                try{
+                    await puter.ui.authenticateWithPuter();
+                    console.log("auth'd", puter.authToken)
+                    
+                    // We have to remake the handler since the old one was done with improper auth, so we'll just talk with the auth server directly
+                    const { token: wispToken, server: wispServer } = (await (await fetch(puter.APIOrigin + '/wisp/relay-token/create', {
+                        method: 'POST',
+                        headers: {
+                            Authorization: `Bearer ${puter.authToken}`,
+                            'Content-Type': 'application/json',
+                        },
+                        body: JSON.stringify({}),
+                    })).json());
+
+                    wispInfo.handler = new PWispHandler(wispServer, wispToken);
+
+                    // Wait for websocket to fully open
+                    await new Promise((res, req) => {
+                        wispInfo.handler.onReady = res;
+                    });
+                }catch(e){
+                    // if authentication fails, throw an error
+                    throw (e);
+                }
+            }
+            const callbacks = {
+                dataCallBack: (data) => {
+                    this.emit("data", data);
+                },
+                closeCallBack: (reason) => {
+                    if (reason !== 0x02) {
+                        this.emit("error", new Error(errors[reason]));
+                        this.emit("close", true);
+                        return;    
+                    }
+                    this.emit("close", false);
+                }
+            }
+    
+            this._streamID = wispInfo.handler.register(host, port, callbacks);
+            setTimeout(() => {this.emit("open", undefined)}, 0);
+
+        })();
     }
     addListener(...args) {
         this.on(...args);
diff --git a/src/puter-js/src/modules/networking/PWispHandler.js b/src/puter-js/src/modules/networking/PWispHandler.js
index 66559419b..0b1d472c3 100644
--- a/src/puter-js/src/modules/networking/PWispHandler.js
+++ b/src/puter-js/src/modules/networking/PWispHandler.js
@@ -4,6 +4,7 @@ export class PWispHandler {
     _ws;
     _nextStreamID = 1;
     _bufferMax;
+    onReady = undefined;
     streamMap = new Map();
     constructor(wispURL, puterAuth) {
         const setup = () => {
@@ -21,6 +22,9 @@ export class PWispHandler {
                             this._ws.onclose = () => {
                                 setTimeout(setup(), 1000);
                             }
+                            if (this.onReady) {
+                                this.onReady();
+                            }
                             return;
                         }
                         this.streamMap.get(parsed.streamID).buffer = parsed.remainingBuffer;