diff --git a/src/puter-js/src/modules/networking/PSocket.js b/src/puter-js/src/modules/networking/PSocket.js index 99f62f720..92e74c081 100644 --- a/src/puter-js/src/modules/networking/PSocket.js +++ b/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)); diff --git a/src/puter-js/src/modules/networking/PSocket.test.js b/src/puter-js/src/modules/networking/PSocket.test.js new file mode 100644 index 000000000..d133ad443 --- /dev/null +++ b/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/); + }); +}); diff --git a/src/puter-js/src/modules/networking/PWispHandler.js b/src/puter-js/src/modules/networking/PWispHandler.js index 2d8d530d3..072ca73ba 100644 --- a/src/puter-js/src/modules/networking/PWispHandler.js +++ b/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++ ) { diff --git a/src/puter-js/src/modules/networking/PWispHandler.test.js b/src/puter-js/src/modules/networking/PWispHandler.test.js new file mode 100644 index 000000000..2b38901da --- /dev/null +++ b/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]); + }); +});