Merge pull request #1187 from carlosmfreitas2409/main

fix: Websocket `beforeLoad` not being executed

Author: SaltyAom

src/adapter/bun/index.ts

@@ -412,6 +412,14 @@ export const BunAdapter: ElysiaAdapter = {
 			normalize: app.config.normalize
 		})
 
+		const validateUpgradeData = getSchemaValidator(options.upgradeData, {
+			// @ts-expect-error private property
+			modules: app.definitions.typebox,
+			// @ts-expect-error private property
+			models: app.definitions.type as Record<string, TSchema>,
+			normalize: app.config.normalize
+		})
+
 		app.route(
 			'WS',
 			path as any,
@@ -457,6 +465,12 @@ export const BunAdapter: ElysiaAdapter = {
 
 				let _id: string | undefined
 
+				let _beforeHandleData: any
+				if (typeof options.beforeHandle === 'function') {
+					const result = options.beforeHandle(context)
+					_beforeHandleData = result instanceof Promise ? await result : result
+				}
+
 				const errorHandlers = [
 					...(Array.isArray(options.error)
 						? options.error
@@ -502,11 +516,22 @@ export const BunAdapter: ElysiaAdapter = {
 								options.pong?.(data)
 							},
 							open(ws: ServerWebSocket<any>) {
+								if (validateUpgradeData?.Check(_beforeHandleData) === false) {
+									return void ws.send(
+										new ValidationError(
+											'upgradeData',
+											validateUpgradeData,
+											_beforeHandleData
+										).message as string
+									)
+								}
+
 								try {
 									handleResponse(
 										ws,
 										options.open?.(
-											new ElysiaWS(ws, context as any)
+											new ElysiaWS(ws, context as any),
+											_beforeHandleData as any
 										)
 									)
 								} catch (error) {

src/index.ts

@@ -5424,7 +5424,8 @@ export default class Elysia<
 				MergeSchema<Ephemeral['schema'], Metadata['schema']>
 			>
 		>,
-		const Macro extends Metadata['macro']
+		const Macro extends Metadata['macro'],
+		const UpgradeDataSchema extends TSchema,
 	>(
 		path: Path,
 		options: WSLocalHook<
@@ -5436,7 +5437,8 @@ export default class Elysia<
 					Volatile['resolve'] &
 					MacroToContext<Metadata['macroFn'], Macro>
 			},
-			Macro
+			Macro,
+			UpgradeDataSchema
 		>
 	): Elysia<
 		BasePath,

src/ws/types.ts

@@ -1,3 +1,5 @@
+import { TSchema } from '@sinclair/typebox'
+
 import type { ElysiaWS } from './index'
 import { WebSocketHandler } from './bun'
 
@@ -17,7 +19,8 @@ import {
 	Prettify,
 	RouteSchema,
 	SingletonBase,
-	TransformHandler
+	TransformHandler,
+	UnwrapSchema
 } from '../types'
 
 type TypedWebSocketMethod =
@@ -33,10 +36,12 @@ export type FlattenResponse<Response extends RouteSchema['response']> =
 
 interface TypedWebSocketHandler<
 	in out Context,
-	in out Route extends RouteSchema = {}
+	in out Route extends RouteSchema = {},
+	in out UpgradeDataSchema extends unknown = unknown
 > extends Omit<WebSocketHandler<Context>, TypedWebSocketMethod> {
 	open?(
-		ws: ElysiaWS<Context, Omit<Route, 'body'> & { body: never }>
+		ws: ElysiaWS<Context, Omit<Route, 'body'> & { body: never }>,
+		data: UpgradeDataSchema
 	): MaybePromise<FlattenResponse<Route['response']> | void>
 	message?(
 		ws: ElysiaWS<Context, Route>,
@@ -118,13 +123,14 @@ export type WSParseHandler<Route extends RouteSchema, Context = {}> = (
 	message: unknown
 ) => MaybePromise<Route['body'] | void | undefined>
 
-export type AnyWSLocalHook = WSLocalHook<any, any, any, any>
+export type AnyWSLocalHook = WSLocalHook<any, any, any, any, any>
 
 export type WSLocalHook<
 	LocalSchema extends InputSchema,
 	Schema extends RouteSchema,
 	Singleton extends SingletonBase,
-	Macro extends MetadataBase['macro']
+	Macro extends MetadataBase['macro'],
+	UpgradeDataSchema extends TSchema,
 > = Prettify<Macro> &
 	(LocalSchema extends any ? LocalSchema : Prettify<LocalSchema>) & {
 		detail?: DocumentDecoration
@@ -134,6 +140,11 @@ export type WSLocalHook<
 		upgrade?: Record<string, unknown> | ((context: Context) => unknown)
 		parse?: MaybeArray<WSParseHandler<Schema>>
 
+		/**
+		 * Upgrade data's value
+		 */
+		upgradeData?: UpgradeDataSchema;
+
 		/**
 		 * Transform context's value
 		 */
@@ -163,5 +174,6 @@ export type WSLocalHook<
 		Omit<Context<Schema, Singleton>, 'body'> & {
 			body: never
 		},
-		Schema
+		Schema,
+		UnwrapSchema<UpgradeDataSchema>
 	>

test/ws/message.test.ts

@@ -487,4 +487,38 @@ describe('WebSocket message', () => {
 		await wsClosed(ws)
 		app.stop()
 	})
+
+	it('should call beforeHandle hook', async () => {
+		const app = new Elysia()
+			.ws('/ws', {
+				upgradeData: t.Object({
+					hello: t.String()
+				}),
+				beforeHandle() {
+					return {
+						hello: 'world'
+					}
+				},
+				open(ws, data) {
+					ws.send(data.hello)
+				}
+			})
+			.listen(0)
+
+		const ws = newWebsocket(app.server!)
+
+		await wsOpen(ws)
+
+		const message = wsMessage(ws)
+
+		ws.send('Hello!')
+
+		const { data } = await message
+		
+		expect(data).toBe('world')
+
+		await wsClosed(ws)
+
+		app.stop()
+	})
 })