diff --git a/src/adapter/bun/index.ts b/src/adapter/bun/index.ts index 47b794815..36dee52c1 100644 --- a/src/adapter/bun/index.ts +++ b/src/adapter/bun/index.ts @@ -502,6 +502,30 @@ export const BunAdapter: ElysiaAdapter = { let _id: string | undefined + // Cache ElysiaWS wrappers per raw ServerWebSocket to + // guarantee stable object identity across open / message / + // drain / close hooks (fixes #1716). + const elysiaWSCache = new WeakMap< + ServerWebSocket, + ElysiaWS + >() + + const getElysiaWS = ( + ws: ServerWebSocket, + body?: unknown + ): ElysiaWS => { + let cached = elysiaWSCache.get(ws) + + if (!cached) { + cached = new ElysiaWS(ws, context as any, body) + elysiaWSCache.set(ws, cached) + } else if (body !== undefined) { + cached.body = body + } + + return cached + } + if (typeof options.beforeHandle === 'function') { const result = options.beforeHandle(context) if (result instanceof Promise) await result @@ -559,9 +583,7 @@ export const BunAdapter: ElysiaAdapter = { try { await handleResponse( ws, - options.open?.( - new ElysiaWS(ws, context as any) - ) + options.open?.(getElysiaWS(ws) as any) ) } catch (error) { handleErrors(ws, error) @@ -592,14 +614,12 @@ export const BunAdapter: ElysiaAdapter = { } try { + const elysiaWs = getElysiaWS(ws, message) + await handleResponse( ws, options.message?.( - new ElysiaWS( - ws, - context as any, - message - ), + elysiaWs, message as any ) ) @@ -611,9 +631,7 @@ export const BunAdapter: ElysiaAdapter = { try { await handleResponse( ws, - options.drain?.( - new ElysiaWS(ws, context as any) - ) + options.drain?.(getElysiaWS(ws) as any) ) } catch (error) { handleErrors(ws, error) @@ -628,7 +646,7 @@ export const BunAdapter: ElysiaAdapter = { await handleResponse( ws, options.close?.( - new ElysiaWS(ws, context as any), + getElysiaWS(ws) as any, code, reason ) diff --git a/test/ws/identity.test.ts b/test/ws/identity.test.ts new file mode 100644 index 000000000..a1e0b416d --- /dev/null +++ b/test/ws/identity.test.ts @@ -0,0 +1,296 @@ +import { describe, it, expect } from 'bun:test' +import { Elysia } from '../../src' +import { newWebsocket, wsOpen, wsMessage, wsClosed } from './utils' + +describe('WebSocket object identity (#1716)', () => { + it('should return the same ElysiaWS instance across open and message hooks', async () => { + let openWs: any + let messageWs: any + + const app = new Elysia() + .ws('/ws', { + open(ws) { + openWs = ws + }, + message(ws, message) { + messageWs = ws + ws.send('ok') + } + }) + .listen(0) + + const ws = newWebsocket(app.server!) + + await wsOpen(ws) + + const reply = wsMessage(ws) + ws.send('hello') + + await reply + + expect(openWs).toBeDefined() + expect(messageWs).toBeDefined() + expect(openWs).toBe(messageWs) + + await wsClosed(ws) + app.stop() + }) + + it('should return the same ElysiaWS instance across open and close hooks', async () => { + let openWs: any + let closeWs: any + + const app = new Elysia() + .ws('/ws', { + open(ws) { + openWs = ws + }, + message() {}, + close(ws) { + closeWs = ws + } + }) + .listen(0) + + const ws = newWebsocket(app.server!) + + await wsOpen(ws) + + // Give the open hook time to fire + await Bun.sleep(10) + + await wsClosed(ws) + + // Give the close hook time to fire + await Bun.sleep(10) + + expect(openWs).toBeDefined() + expect(closeWs).toBeDefined() + expect(openWs).toBe(closeWs) + + app.stop() + }) + + it('should return the same instance across open, message, and close', async () => { + const instances: any[] = [] + + const app = new Elysia() + .ws('/ws', { + open(ws) { + instances.push(ws) + }, + message(ws) { + instances.push(ws) + ws.send('ack') + }, + close(ws) { + instances.push(ws) + } + }) + .listen(0) + + const ws = newWebsocket(app.server!) + + await wsOpen(ws) + + const reply = wsMessage(ws) + ws.send('test') + await reply + + await wsClosed(ws) + await Bun.sleep(10) + + expect(instances.length).toBe(3) + expect(instances[0]).toBe(instances[1]) + expect(instances[1]).toBe(instances[2]) + + app.stop() + }) + + it('should preserve custom properties set in open hook', async () => { + let customPropInMessage: string | undefined + let customPropInClose: string | undefined + + const app = new Elysia() + .ws('/ws', { + open(ws) { + ;(ws as any)._customId = 'test-id-123' + }, + message(ws) { + customPropInMessage = (ws as any)._customId + ws.send('ok') + }, + close(ws) { + customPropInClose = (ws as any)._customId + } + }) + .listen(0) + + const ws = newWebsocket(app.server!) + + await wsOpen(ws) + + const reply = wsMessage(ws) + ws.send('hi') + await reply + + await wsClosed(ws) + await Bun.sleep(10) + + expect(customPropInMessage).toBe('test-id-123') + expect(customPropInClose).toBe('test-id-123') + + app.stop() + }) + + it('should maintain separate identity for different connections', async () => { + const instances = new Set() + + const app = new Elysia() + .ws('/ws', { + open(ws) { + instances.add(ws) + }, + message() {} + }) + .listen(0) + + const ws1 = newWebsocket(app.server!) + await wsOpen(ws1) + + const ws2 = newWebsocket(app.server!) + await wsOpen(ws2) + + // Give open hooks time to fire + await Bun.sleep(10) + + expect(instances.size).toBe(2) + + await wsClosed(ws1) + await wsClosed(ws2) + + app.stop() + }) + + it('should allow tracking connections in a Set and cleaning up on close', async () => { + const activeConnections = new Set() + + const app = new Elysia() + .ws('/ws', { + open(ws) { + activeConnections.add(ws) + }, + message() {}, + close(ws) { + activeConnections.delete(ws) + } + }) + .listen(0) + + const ws1 = newWebsocket(app.server!) + await wsOpen(ws1) + + const ws2 = newWebsocket(app.server!) + await wsOpen(ws2) + + await Bun.sleep(10) + + expect(activeConnections.size).toBe(2) + + await wsClosed(ws1) + await Bun.sleep(10) + + // ws1 should be removed because close handler received the same instance + expect(activeConnections.size).toBe(1) + + await wsClosed(ws2) + await Bun.sleep(10) + + expect(activeConnections.size).toBe(0) + + app.stop() + }) + + it('should allow per-socket state storage in a Map', async () => { + const socketState = new Map() + + const app = new Elysia() + .ws('/ws', { + open(ws) { + socketState.set(ws, { messageCount: 0 }) + }, + message(ws) { + const state = socketState.get(ws) + if (state) { + state.messageCount++ + } + ws.send(`count:${state?.messageCount}`) + }, + close(ws) { + socketState.delete(ws) + } + }) + .listen(0) + + const ws = newWebsocket(app.server!) + await wsOpen(ws) + + // Send 3 messages and check state accumulates + let reply = wsMessage(ws) + ws.send('a') + let msg = await reply + expect(msg.data).toBe('count:1') + + reply = wsMessage(ws) + ws.send('b') + msg = await reply + expect(msg.data).toBe('count:2') + + reply = wsMessage(ws) + ws.send('c') + msg = await reply + expect(msg.data).toBe('count:3') + + await wsClosed(ws) + await Bun.sleep(10) + + expect(socketState.size).toBe(0) + + app.stop() + }) + + it('should update body on each message while preserving identity', async () => { + const bodies: unknown[] = [] + let stableWs: any + + const app = new Elysia() + .ws('/ws', { + open(ws) { + stableWs = ws + }, + message(ws, message) { + bodies.push(ws.body) + expect(ws).toBe(stableWs) + ws.send('ok') + } + }) + .listen(0) + + const ws = newWebsocket(app.server!) + await wsOpen(ws) + + let reply = wsMessage(ws) + ws.send('first') + await reply + + reply = wsMessage(ws) + ws.send('second') + await reply + + expect(bodies[0]).toBe('first') + expect(bodies[1]).toBe('second') + + await wsClosed(ws) + app.stop() + }) +})