fix wisp issues

src/puter-js/src/modules/networking/PSocket.js

@@ -50,9 +50,17 @@ export class PSocket extends EventListener {
 
                 wispInfo.handler = new PWispHandler(wispServer, wispToken);
                 // Wait for websocket to fully open
-                await new Promise((res, req) => {
-                    wispInfo.handler.onReady = res;
-                });
+                try {
+                    await new Promise((res, rej) => {
+                        wispInfo.handler.onReady = res;
+                        wispInfo.handler.onError = rej;
+                    });
+                } catch (e) {
+                    // Drop the dead handler so the next socket redials instead
+                    // of registering streams on a relay that never opened.
+                    wispInfo.handler = undefined;
+                    throw e;
+                }
             }
 
             const callbacks = {
@@ -74,7 +82,12 @@ export class PSocket extends EventListener {
                 this.emit('open', undefined);
             }, 0);
 
-        })();
+        })().catch((e) => {
+            // Nothing awaits this body, so a failure to connect has to reach
+            // the caller as an 'error' event rather than an unhandled rejection.
+            this.emit('error', e instanceof Error ? e : new Error(String(e)));
+            this.emit('close', true);
+        });
     }
     /**
      * Registers a handler for a socket event, the same as `on`.
@@ -100,7 +113,7 @@ export class PSocket extends EventListener {
             wispInfo.handler.write(this._streamID, data);
             if ( callback ) callback();
         } else if ( data.resize ) { // ArrayBuffer
-            data.write(this._streamID, new Uint8Array(data));
+            wispInfo.handler.write(this._streamID, new Uint8Array(data));
             if ( callback ) callback();
         } else if ( typeof (data) === 'string' ) {
             wispInfo.handler.write(this._streamID, texten.encode(data));

src/puter-js/src/modules/networking/PSocket.test.js

@@ -0,0 +1,161 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+// The relay handshake is the thing under test, so both the token fetch and the
+// wisp handler are stubbed: no network, no websocket.
+const mockFetchUrl = vi.fn();
+vi.mock('../../lib/networkUtils.js', () => ({
+    fetchUrl: (...args) => mockFetchUrl(...args),
+}));
+
+const handlerInstances = [];
+class FakeWispHandler {
+    onReady = undefined;
+    onError = undefined;
+
+    constructor (url, auth) {
+        this.url = url;
+        this.auth = auth;
+        this.write = vi.fn();
+        this.close = vi.fn();
+        this.register = vi.fn(() => 7);
+        handlerInstances.push(this);
+    }
+}
+vi.mock('./PWispHandler.js', () => ({
+    PWispHandler: class {
+        constructor (...args) {
+            return new FakeWispHandler(...args);
+        }
+    },
+}));
+
+const { PSocket, wispInfo } = await import('./PSocket.js');
+
+const tokenResponse = () => Promise.resolve({
+    json: () => Promise.resolve({ token: 'wisp-token', server: 'wss://relay.test/' }),
+});
+
+// Lets a test act after the constructor's async body has reached its await.
+const flush = () => new Promise(r => setTimeout(r, 0));
+
+const origPuter = globalThis.puter;
+
+beforeEach(() => {
+    handlerInstances.length = 0;
+    mockFetchUrl.mockReset().mockImplementation(tokenResponse);
+    wispInfo.handler = undefined;
+    globalThis.puter = { authToken: 'tok', APIOrigin: 'https://api.test', env: 'nodejs' };
+});
+
+afterEach(() => {
+    wispInfo.handler = undefined;
+    globalThis.puter = origPuter;
+});
+
+describe('PSocket relay handshake', () => {
+    it('opens once the handler reports ready', async () => {
+        const socket = new PSocket('example.com', 80);
+        const onOpen = vi.fn();
+        socket.on('open', onOpen);
+
+        await flush();
+        handlerInstances[0].onReady();
+        // `open` is emitted from a setTimeout of its own, one tick after the
+        // constructor body resumes.
+        await flush();
+        await flush();
+
+        expect(onOpen).toHaveBeenCalledTimes(1);
+        expect(handlerInstances[0].register).toHaveBeenCalledWith(
+            'example.com', 80, expect.any(Object),
+        );
+    });
+
+    // The reject parameter was typo'd `req` and never called, so a failed
+    // handshake hung instead of surfacing.
+    it('emits error instead of hanging when the handshake fails', async () => {
+        const socket = new PSocket('example.com', 80);
+        const onError = vi.fn();
+        const onClose = vi.fn();
+        socket.on('error', onError);
+        socket.on('close', onClose);
+
+        await flush();
+        handlerInstances[0].onError(new Error('relay unreachable'));
+        await flush();
+
+        expect(onError).toHaveBeenCalledTimes(1);
+        expect(onError.mock.calls[0][0]).toBeInstanceOf(Error);
+        expect(onError.mock.calls[0][0].message).toBe('relay unreachable');
+        expect(onClose).toHaveBeenCalledWith(true);
+        // The stream is never registered on a relay that failed to open.
+        expect(handlerInstances[0].register).not.toHaveBeenCalled();
+    });
+
+    it('drops the dead handler so a later socket redials', async () => {
+        const first = new PSocket('example.com', 80);
+        first.on('error', () => {});
+        await flush();
+        handlerInstances[0].onError(new Error('relay unreachable'));
+        await flush();
+
+        expect(wispInfo.handler).toBeUndefined();
+
+        const second = new PSocket('example.com', 80);
+        second.on('error', () => {});
+        await flush();
+
+        expect(handlerInstances).toHaveLength(2);
+    });
+});
+
+describe('PSocket write', () => {
+    const connected = async () => {
+        const socket = new PSocket('example.com', 80);
+        socket.on('error', () => {});
+        await flush();
+        handlerInstances[0].onReady();
+        await flush();
+        return socket;
+    };
+
+    it('writes a typed array through the relay handler', async () => {
+        const socket = await connected();
+        const payload = new Uint8Array([1, 2, 3]);
+
+        socket.write(payload);
+
+        expect(handlerInstances[0].write).toHaveBeenCalledWith(7, payload);
+    });
+
+    // This branch called `data.write(...)` — a method ArrayBuffers do not have
+    // — so every ArrayBuffer write threw instead of reaching the relay.
+    it('writes an ArrayBuffer through the relay handler', async () => {
+        const socket = await connected();
+        const buffer = new Uint8Array([4, 5, 6]).buffer;
+        const callback = vi.fn();
+
+        expect(() => socket.write(buffer, callback)).not.toThrow();
+
+        expect(handlerInstances[0].write).toHaveBeenCalledTimes(1);
+        const [streamID, sent] = handlerInstances[0].write.mock.calls[0];
+        expect(streamID).toBe(7);
+        expect(Array.from(sent)).toEqual([4, 5, 6]);
+        expect(callback).toHaveBeenCalledTimes(1);
+    });
+
+    it('encodes a string before writing', async () => {
+        const socket = await connected();
+
+        socket.write('hi');
+
+        const [, sent] = handlerInstances[0].write.mock.calls[0];
+        expect(Array.from(sent)).toEqual([104, 105]);
+    });
+
+    it('throws on an unsupported data type', async () => {
+        const socket = await connected();
+
+        expect(() => socket.write(42)).toThrow(/Invalid data type/);
+    });
+});

src/puter-js/src/modules/networking/PWispHandler.js

@@ -4,12 +4,29 @@ export class PWispHandler {
     _ws;
     _nextStreamID = 1;
     _bufferMax;
+    // Set once the relay answers the handshake with CONTINUE on stream 0.
+    // Decides whether a close is a dropped connection (reconnect) or a
+    // handshake that never completed (report failure).
+    _ready = false;
     onReady = undefined;
+    onError = undefined;
     streamMap = new Map();
     constructor (wispURL, puterAuth) {
         const setup = () => {
             this._ws = new WebSocket(wispURL);
             this._ws.binaryType = 'arraybuffer';
+            this._ws.onerror = () => {
+                this._fail(new Error(`Wisp relay connection failed: ${wispURL}`));
+            };
+            this._ws.onclose = () => {
+                if ( this._ready ) {
+                    // Pass the function itself: `setTimeout(setup(), 1000)`
+                    // would reconnect immediately and schedule nothing.
+                    setTimeout(setup, 1000);
+                    return;
+                }
+                this._fail(new Error('Wisp relay closed before the handshake completed'));
+            };
             this._ws.onmessage = (event) => {
                 const parsed = parseIncomingPacket(new Uint8Array(event.data));
                 switch ( parsed.packetType ) {
@@ -19,16 +36,14 @@ export class PWispHandler {
                 case CONTINUE:
                     if ( parsed.streamID === 0 ) {
                         this._bufferMax = parsed.remainingBuffer;
-                        this._ws.onclose = () => {
-                            setTimeout(setup(), 1000);
-                        };
+                        this._ready = true;
                         if ( this.onReady ) {
                             this.onReady();
                         }
                         return;
                     }
                     this.streamMap.get(parsed.streamID).buffer = parsed.remainingBuffer;
-                    this._continue();
+                    this._continue(parsed.streamID);
                     break;
                 case CLOSE:
                     if ( parsed.streamID !== 0 )
@@ -48,6 +63,13 @@ export class PWispHandler {
         };
         setup();
     }
+    // Reports a connection-level failure once, so a caller waiting on the
+    // handshake is rejected rather than left hanging.
+    _fail (error) {
+        const onError = this.onError;
+        this.onError = undefined;
+        if ( onError ) onError(error);
+    }
     _continue (streamID) {
         const queue = this.streamMap.get(streamID).queue;
         for ( let i = 0; i < queue.length; i++ ) {

src/puter-js/src/modules/networking/PWispHandler.test.js

@@ -0,0 +1,151 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { CONTINUE, DATA, createWispPacket, parseIncomingPacket } from './parsers.js';
+import { PWispHandler } from './PWispHandler.js';
+
+// Fake WebSocket standing in for the relay connection. Records every instance
+// so reconnects are observable, and every frame sent so writes are.
+class FakeWebSocket {
+    static instances = [];
+
+    constructor (url) {
+        this.url = url;
+        this.sent = [];
+        this.readyState = 0;
+        FakeWebSocket.instances.push(this);
+    }
+
+    send (data) {
+        this.sent.push(data);
+    }
+
+    close () {}
+}
+
+const latest = () => FakeWebSocket.instances.at(-1);
+
+// Hand a wisp packet to the handler the way the browser would.
+const deliver = (packet) => latest().onmessage({ data: packet.buffer });
+
+const handshake = (remainingBuffer = 4) =>
+    deliver(createWispPacket({ packetType: CONTINUE, streamID: 0, remainingBuffer }));
+
+const origWebSocket = globalThis.WebSocket;
+
+beforeEach(() => {
+    vi.useFakeTimers();
+    FakeWebSocket.instances = [];
+    globalThis.WebSocket = FakeWebSocket;
+});
+
+afterEach(() => {
+    vi.useRealTimers();
+    globalThis.WebSocket = origWebSocket;
+});
+
+describe('PWispHandler handshake', () => {
+    it('signals ready once the relay answers CONTINUE on stream 0', () => {
+        const handler = new PWispHandler('wss://relay.test/', 'token');
+        const onReady = vi.fn();
+        handler.onReady = onReady;
+
+        handshake(7);
+
+        expect(onReady).toHaveBeenCalledTimes(1);
+        expect(handler._bufferMax).toBe(7);
+    });
+
+    // Without an `onerror` handler a failed connection left every caller
+    // waiting on the handshake hanging forever.
+    it('reports a connection error to onError', () => {
+        const handler = new PWispHandler('wss://relay.test/', 'token');
+        const onError = vi.fn();
+        handler.onError = onError;
+
+        expect(typeof latest().onerror).toBe('function');
+        latest().onerror(new Event('error'));
+
+        expect(onError).toHaveBeenCalledTimes(1);
+        expect(onError.mock.calls[0][0]).toBeInstanceOf(Error);
+    });
+
+    it('reports a close that happens before the handshake completes', () => {
+        const handler = new PWispHandler('wss://relay.test/', 'token');
+        const onError = vi.fn();
+        handler.onError = onError;
+
+        latest().onclose();
+
+        expect(onError).toHaveBeenCalledTimes(1);
+        expect(onError.mock.calls[0][0].message).toMatch(/before the handshake/);
+    });
+
+    it('reports a connection failure only once', () => {
+        const handler = new PWispHandler('wss://relay.test/', 'token');
+        const onError = vi.fn();
+        handler.onError = onError;
+
+        latest().onerror(new Event('error'));
+        latest().onclose();
+
+        expect(onError).toHaveBeenCalledTimes(1);
+    });
+});
+
+describe('PWispHandler reconnect', () => {
+    // `setTimeout(setup(), 1000)` reconnected synchronously and scheduled
+    // `undefined`, so the 1s backoff never applied.
+    it('reconnects one second after an established connection drops', () => {
+        new PWispHandler('wss://relay.test/', 'token');
+        handshake();
+        expect(FakeWebSocket.instances).toHaveLength(1);
+
+        latest().onclose();
+
+        // Nothing yet: the reconnect is scheduled, not immediate.
+        expect(FakeWebSocket.instances).toHaveLength(1);
+
+        vi.advanceTimersByTime(1000);
+
+        expect(FakeWebSocket.instances).toHaveLength(2);
+        expect(latest().url).toBe('wss://relay.test/');
+    });
+
+    it('does not reconnect when the handshake never completed', () => {
+        new PWispHandler('wss://relay.test/', 'token');
+
+        latest().onclose();
+        vi.advanceTimersByTime(5000);
+
+        expect(FakeWebSocket.instances).toHaveLength(1);
+    });
+});
+
+describe('PWispHandler backpressure', () => {
+    // `this._continue()` was called with no argument, so the drain path did
+    // `streamMap.get(undefined).queue` and threw.
+    it('drains the queue for the stream the CONTINUE names', () => {
+        const handler = new PWispHandler('wss://relay.test/', 'token');
+        handshake(1);
+
+        const streamID = handler.register('example.com', 80, {
+            dataCallBack: vi.fn(), closeCallBack: vi.fn(),
+        });
+        const sentAfterRegister = latest().sent.length;
+
+        // Exhaust the credit, so the next write is queued rather than sent.
+        handler.write(streamID, new Uint8Array([1]));
+        handler.write(streamID, new Uint8Array([2]));
+        expect(handler.streamMap.get(streamID).queue).toHaveLength(1);
+        expect(latest().sent).toHaveLength(sentAfterRegister + 1);
+
+        expect(() => deliver(createWispPacket({
+            packetType: CONTINUE, streamID, remainingBuffer: 4,
+        }))).not.toThrow();
+
+        expect(handler.streamMap.get(streamID).queue).toHaveLength(0);
+        const flushed = parseIncomingPacket(latest().sent.at(-1));
+        expect(flushed.packetType).toBe(DATA);
+        expect(flushed.streamID).toBe(streamID);
+        expect(Array.from(flushed.payload)).toEqual([2]);
+    });
+});