fix(ws): support parse 'none' and execute custom parsers on raw message (#1911)

src/ws/index.ts

@@ -209,11 +209,27 @@ export class ElysiaWS<Context = unknown, Route extends RouteSchema = {}>
 }
 
 export const createWSMessageParser = (
-	parse: MaybeArray<WSParseHandler<any>>
+	parse: MaybeArray<WSParseHandler<any>> | 'none' | undefined
 ) => {
+	if (parse === 'none')
+		return function parseMessage(ws: ServerWebSocket<any>, message: any) {
+			return message
+		}
+
 	const parsers = typeof parse === 'function' ? [parse] : parse
 
 	return async function parseMessage(ws: ServerWebSocket<any>, message: any) {
+		if (parsers) {
+			const elysiaWS = new ElysiaWS(ws as any, ws.data as any)
+
+			for (let i = 0; i < parsers.length; i++) {
+				let temp = parsers[i](elysiaWS as any, message)
+				if (temp instanceof Promise) temp = await temp
+
+				if (temp !== undefined) return temp
+			}
+		}
+
 		if (typeof message === 'string') {
 			const start = message?.charCodeAt(0)
 
@@ -229,14 +245,6 @@ export const createWSMessageParser = (
 			else if (message === 'null') message = null
 		}
 
-		if (parsers)
-			for (let i = 0; i < parsers.length; i++) {
-				let temp = parsers[i](ws as any, message)
-				if (temp instanceof Promise) temp = await temp
-
-				if (temp !== undefined) return temp
-			}
-
 		return message
 	}
 }

src/ws/types.ts

@@ -135,7 +135,7 @@ export type WSLocalHook<
 	 * Headers to register to websocket before `upgrade`
 	 */
 	upgrade?: Record<string, unknown> | ((context: Context) => unknown)
-	parse?: MaybeArray<WSParseHandler<Schema>>
+	parse?: MaybeArray<WSParseHandler<Schema>> | 'none'
 
 	/**
 	 * Transform context's value

test/ws/message.test.ts

@@ -609,4 +609,69 @@ describe('WebSocket message', () => {
         await wsClosed(ws)
         app.stop()
     })
+
+    it('should not parse message when parse is none', async () => {
+        const app = new Elysia()
+            .ws('/ws', {
+                parse: 'none',
+                message(ws, message) {
+                    ws.send(typeof message + ' ' + message)
+                }
+            })
+            .listen(0)
+
+        const ws = newWebsocket(app.server!)
+
+        await wsOpen(ws)
+
+        const message1 = wsMessage(ws)
+        ws.send('["ping"]')
+        const res1 = await message1
+
+        expect(res1.data).toBe('string ["ping"]')
+
+        const message2 = wsMessage(ws)
+        ws.send('{"hello":"world"}')
+        const res2 = await message2
+
+        expect(res2.data).toBe('string {"hello":"world"}')
+
+        await wsClosed(ws)
+        app.stop()
+    })
+
+    it('should pass raw message to custom parse handler', async () => {
+        const app = new Elysia()
+            .ws('/ws', {
+                parse(ws, raw) {
+                    if (typeof raw === 'string' && raw.startsWith('PREFIX:'))
+                        return raw.slice(7)
+                    if (typeof raw === 'string' && raw === '["ping"]')
+                        return raw
+                },
+                message(ws, message) {
+                    ws.send(typeof message + ' ' + message)
+                }
+            })
+            .listen(0)
+
+        const ws = newWebsocket(app.server!)
+
+        await wsOpen(ws)
+
+        const message1 = wsMessage(ws)
+        ws.send('PREFIX:["ping"]')
+        const res1 = await message1
+
+        expect(res1.data).toBe('string ["ping"]')
+
+        const message2 = wsMessage(ws)
+        ws.send('["ping"]')
+        const res2 = await message2
+
+        expect(res2.data).toBe('string ["ping"]')
+
+        await wsClosed(ws)
+        app.stop()
+    })
 })