Skip to content

Fix WebSocket support #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix WebSocket support #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions src/handler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
/* eslint-disable sonarjs/no-duplicated-branches */
/* eslint-disable sonarjs/no-nested-switch */
/* eslint-disable sonarjs/no-duplicate-string */
import { serialize } from 'cookie'

import type { IncomingMessage, ServerResponse } from 'http'
import { Readable } from 'stream'

Expand All @@ -17,6 +15,8 @@ import type { HTTPHeaders, Prettify } from 'elysia/types'

import type { ReadStream } from 'fs'

import { ElysiaNodeWebsocketResponse } from '.'

type SetResponse = Prettify<
Omit<Context['set'], 'status'> & {
status: number
Expand Down Expand Up @@ -71,8 +71,9 @@ const handleFile = (
if (res)
return response.arrayBuffer().then((arrayBuffer) => {
set!.headers['content-type'] = response.type
set!.headers['content-range'] =
`bytes 0-${arrayBuffer.byteLength - 1}/${arrayBuffer.byteLength}`
set!.headers['content-range'] = `bytes 0-${
arrayBuffer.byteLength - 1
}/${arrayBuffer.byteLength}`

delete set?.headers['content-length']

Expand Down Expand Up @@ -259,6 +260,12 @@ export const mapResponse = (
}

switch (response?.constructor?.name) {
case 'Symbol': {
switch (response) {
case ElysiaNodeWebsocketResponse:
return [ElysiaNodeWebsocketResponse, set as any]
}
}
case 'String':
set.headers['content-type'] = 'text/plain;charset=utf8'

Expand Down
75 changes: 27 additions & 48 deletions src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,7 @@ import formidable from 'formidable'
import { Elysia } from 'elysia'
import type { Server } from 'elysia/universal'

import {
isNotEmpty,
isNumericString,
mergeLifeCycle,
randomId
} from 'elysia/utils'
import { isNotEmpty, isNumericString, randomId } from 'elysia/utils'
import { type ElysiaAdapter } from 'elysia/adapter'

import { mapResponse, mapEarlyResponse, mapCompactResponse } from './handler'
Expand All @@ -26,6 +21,7 @@ import {
} from './utils'

export const ElysiaNodeContext = Symbol('ElysiaNodeContext')
export const ElysiaNodeWebsocketResponse = Symbol('ElysiaNodeWebsocketResponse')

const getUrl = (req: IncomingMessage) => {
if (req.headers.host) return `http://${req.headers.host}${req.url}`
Expand All @@ -36,7 +32,7 @@ const getUrl = (req: IncomingMessage) => {
return `http://localhost${req.url}`
}

export const nodeRequestToWebstand = (
export const nodeRequestToWebStandard = (
req: IncomingMessage,
abortController?: AbortController
) => {
Expand Down Expand Up @@ -87,7 +83,7 @@ export const node = () => {
`Object.defineProperty(c,'request',{` +
`get(){` +
`if(_request)return _request\n` +
`return _request=nodeRequestToWebstand(c[ElysiaNodeContext].req)` +
`return _request=nodeRequestToWebStandard(c[ElysiaNodeContext].req)` +
`}` +
`})` +
'}\n'
Expand All @@ -97,7 +93,7 @@ export const node = () => {
headers: `c.headers=c[ElysiaNodeContext].req.headers\n`,
inject: {
ElysiaNodeContext,
nodeRequestToWebstand,
nodeRequestToWebStandard,
formidable,
readFileToWebStandardFile,
unwrapArrayIfSingle
Expand Down Expand Up @@ -191,7 +187,7 @@ export const node = () => {
composeGeneralHandler: {
parameters: 'r,res',
inject: {
nodeRequestToWebstand,
nodeRequestToWebStandard,
ElysiaNodeContext
},
createContext: (app) => {
Expand Down Expand Up @@ -219,7 +215,7 @@ export const node = () => {
fnLiteral +=
`get request(){` +
`if(_request)return _request\n` +
`return _request = nodeRequestToWebstand(r)` +
`return _request = nodeRequestToWebStandard(r)` +
`},`

fnLiteral +=
Expand Down Expand Up @@ -306,20 +302,15 @@ export const node = () => {
`return [error.message, context.set]`
},
ws(app, path, options) {
const key = Object.keys(app.router.static.ws).length
const key = Object.keys(app.router.history).length
app.router.static.ws[path] = key

const lifecycle = mergeLifeCycle(options, {})

app.router.history.push({
method: '$INTERNALWS',
path,
composed: undefined as any,
handler: undefined as any,
hooks: lifecycle,
websocket: options
const { parse, body, response, ...rest } = options
app.route('$INTERNALWS', path, () => ElysiaNodeWebsocketResponse, {
...rest,
websocket: options,
webSocket: options
})
app.router.http.history.push(['$INTERNALWS', path, options])
},
listen(app) {
return (options, callback) => {
Expand Down Expand Up @@ -347,29 +338,27 @@ export const node = () => {
// @ts-expect-error closest possible type
app.server = serverInfo

let server = createServer(
// @ts-expect-error private property
app._handle
).listen(
// @ts-expect-error private property
let handle = app._handle

const server = createServer((...a) => handle?.(...a)).listen(
typeof options === 'number'
? options
: {
...options,
// @ts-ignore
host: options?.hostname
},
},
() => {
const address = server.address()
const hostname =
typeof address === 'string'
? address
: address
? address.address
: 'localhost'
? address.address
: 'localhost'
const port =
typeof address === 'string'
? 0
: (address?.port ?? 0)
typeof address === 'string' ? 0 : address?.port ?? 0

const serverInfo: Server = {
id: randomId(),
Expand Down Expand Up @@ -405,20 +394,8 @@ export const node = () => {
server.unref()
},
reload() {
server.close(() => {
server = createServer(
// @ts-expect-error private property
app._handle
).listen(
typeof options === 'number'
? options
: {
...options,
// @ts-ignore
host: options?.hostname
}
)
})
// @ts-expect-error private property
handle = app._handle
},
requestIP() {
throw new Error(
Expand All @@ -434,7 +411,9 @@ export const node = () => {
)
},
url: new URL(
`http://${hostname === '::' ? 'localhost' : hostname}:${port}`
`http://${
hostname === '::' ? 'localhost' : hostname
}:${port}`
),
[Symbol.dispose]() {
server.close()
Expand All @@ -454,7 +433,7 @@ export const node = () => {
? (options as any)
: {
port: options
}
}
)
} catch {}
})
Expand Down
69 changes: 45 additions & 24 deletions src/ws.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,25 @@
import {
AnyElysia,
ELYSIA_REQUEST_ID,
redirect,
serializeCookie,
ValidationError,
type Context,
type TSchema
} from 'elysia'
import { parseQuery } from 'elysia/fast-querystring'
import type { TypeCheck } from 'elysia/type-system'
import { getSchemaValidator, isNotEmpty, randomId } from 'elysia/utils'

import { createServer, IncomingMessage, OutgoingMessage } from 'http'
import {
createServer,
IncomingMessage,
OutgoingMessage,
ServerResponse
} from 'http'
import { Socket } from 'net'

import { nodeRequestToWebstand, ElysiaNodeContext } from '.'
import { nodeRequestToWebStandard, ElysiaNodeContext, ElysiaNodeWebsocketResponse } from '.'
import { ElysiaNodeResponse } from './handler'

import { WebSocketServer } from 'ws'
import type { WebSocket as NodeWebSocket } from 'ws'
Expand Down Expand Up @@ -147,15 +154,16 @@ export const requestToContext = (
qi,
path,
url,
query: qi === -1 ? {} : parseQuery(url.substring(qi + 1)),
set,
redirect,
get request() {
if (_request) return _request
return (_request = nodeRequestToWebstand(request))
return (_request = nodeRequestToWebStandard(request))
},
[ElysiaNodeContext]: {
req: request,
res: undefined
res: response
},
headers: request.headers
}
Expand All @@ -171,28 +179,39 @@ export const attachWebSocket = (
})

const staticWsRouter = app.router.static.ws
const router = app.router.http
const history = app.router.history

server.on('upgrade', (request, socket, head) => {
wsServer.handleUpgrade(request, socket, head, async (ws) => {
const qi = request.url!.indexOf('?')
let path = request.url!
if (qi !== -1) path = request.url!.substring(0, qi)
server.on('upgrade', async (request, socket: Socket, head) => {
const qi = request.url!.indexOf('?')
let path = request.url!
if (qi !== -1) path = request.url!.substring(0, qi)

const index = staticWsRouter[path]
if (index === undefined) return
const index = staticWsRouter[path]
if (index === undefined) return

const route = history[index]
if (!route) {
router.find('$INTERNALWS', path)
return
}
const route = history[index]
if (!route) {
return
}

if (!route.websocket) return

const response = new ServerResponse(request)
response.assignSocket(socket)

if (!route.websocket) return
const context = requestToContext(app, request, response)

if (typeof route.composed === 'function') {
const [content] = (await route.composed(
context
)) as unknown as ElysiaNodeResponse
if (content !== ElysiaNodeWebsocketResponse) return
}

wsServer.handleUpgrade(request, socket, head, async (ws) => {
const websocket: AnyWSLocalHook = route.websocket

const validateMessage = getSchemaValidator(route.hooks.body, {
const validateMessage = getSchemaValidator(route.websocket.body, {
// @ts-expect-error private property
modules: app.definitions.typebox,
// @ts-expect-error private property
Expand All @@ -201,7 +220,7 @@ export const attachWebSocket = (
})

const validateResponse = getSchemaValidator(
route.hooks.response as any,
route.websocket.response,
{
// @ts-expect-error private property
modules: app.definitions.typebox,
Expand All @@ -214,7 +233,6 @@ export const attachWebSocket = (
const parseMessage = createWSMessageParser(route.hooks.parse)
const handleResponse = createHandleWSResponse(validateResponse)

const context = requestToContext(app, request, undefined as any)
const set = context.set

if (set.cookie && isNotEmpty(set.cookie)) {
Expand Down Expand Up @@ -295,8 +313,11 @@ export const attachWebSocket = (
)

if (websocket.message)
ws.on('message', async (_message) => {
const message = await parseMessage(elysiaWS, _message)
ws.on('message', async (_message, isBinary) => {
const message = await parseMessage(
elysiaWS,
isBinary ? _message : _message.toString()
)

if (validateMessage?.Check(message) === false)
return void ws.send(
Expand Down