sleepy-serv 0.6.2 → 0.7.0

package/src/{socket.js → socket.ts}

@@ -2,10 +2,17 @@ import Ajv from 'ajv'
 import addFormats from 'ajv-formats'
 import crypto from 'node:crypto'
 
-import { toSegments, formatError, executeMiddlewareChain } from './utils'
-import { TYPES, createMessage, validateMessage } from './messages'
+import { MessageType, createMessage, validateMessage } from './messages'
 
 import {
+  StatusCode,
+  toSegments,
+  formatError,
+  executeMiddlewareChain,
+} from './utils'
+
+import {
+  RequestError,
   NotFoundError,
   UnauthorizedError,
   MethodNotAllowedError,
@@ -14,6 +21,106 @@ import {
   ServiceUnavailableError,
 } from './errors'
 
+import type { WebSocketHandler } from 'bun'
+import type { ValidateFunction } from 'ajv'
+
+import type {
+  HttpMethod,
+  Request,
+  WebSocketRequest,
+  Middleware,
+  SocketData,
+  AppOptions,
+} from './utils'
+
+import type {
+  BaseMessage,
+  RawMessage,
+  RequestMessage,
+  ResponseMessage,
+} from './messages'
+
+export type Ticket = {
+  clientId: string
+  expiresAt: number
+}
+
+export type SocketConnection = {
+  data: SocketData
+  send: (data: string) => unknown
+  close: () => void
+}
+
+export type ActiveSession = {
+  token: string
+  ws: SocketConnection
+}
+
+export type InactiveSession = {
+  token: string
+  expiresAt: number
+}
+
+export type Session = ActiveSession | InactiveSession
+
+export type SocketState = {
+  disconnectThreshold: number
+  heartbeatInterval: number
+  maxTickets: number
+  reclaimTtl: number
+  ticketTtl: number
+  tickets: Map<string, Ticket>
+  activeSessions: Map<string, ActiveSession>
+  inactiveSessions: Map<string, InactiveSession>
+}
+
+export type SocketRoute = {
+  method: HttpMethod
+  path: string
+  segments: string[]
+  chain: Middleware[]
+}
+
+export type SocketEndpoint = {
+  method: HttpMethod
+  path: string
+  handler: (req: Request, res: unknown) => Response
+}
+
+export type SocketCommands = {
+  send: (clientId: string, event: string, body: unknown) => void
+  broadcast: (event: string, body: unknown) => void
+}
+
+type UpgradeData = {
+  clientId?: string
+  [key: string]: unknown
+}
+
+type UpgradeContext = {
+  data?: UpgradeData
+  [key: string]: unknown
+}
+
+type CreateSocketRequest = {
+  query: {
+    ticket: string
+  }
+  server: {
+    upgrade: (raw: unknown, ctx: UpgradeContext) => boolean
+  }
+  raw: unknown
+}
+
+type CreateTicketRequest = Record<string, unknown>
+
+type UpdateTicketRequest = {
+  headers: Headers
+  params: {
+    clientId: string
+  }
+}
+
 const ajv = new Ajv({
   allErrors: true,
 })
@@ -32,7 +139,7 @@ const validateNotMessage = ajv.compile({
   },
 })
 
-const createSocketValidator = ajv.compile({
+const createSocketValidator = ajv.compile<CreateSocketRequest>({
   type: 'object',
   properties: {
     clientId: {
@@ -68,7 +175,7 @@ const createSocketValidator = ajv.compile({
   },
 })
 
-const createTicketValidator = ajv.compile({
+const createTicketValidator = ajv.compile<CreateTicketRequest>({
   type: 'object',
   properties: {
     clientId: {
@@ -80,7 +187,7 @@ const createTicketValidator = ajv.compile({
   },
 })
 
-const updateTicketValidator = ajv.compile({
+const updateTicketValidator = ajv.compile<UpdateTicketRequest>({
   type: 'object',
   properties: {
     clientId: {
@@ -115,17 +222,21 @@ const updateTicketValidator = ajv.compile({
   },
 })
 
-function isSessionActive (session) {
+function isSessionActive (session: InactiveSession): boolean {
   return !session.expiresAt || session.expiresAt > Date.now()
 }
 
-function randomToken (count) {
-  return crypto.randomBytes(count).toString('base64url')
+function randomTicket (): string {
+  return crypto.randomBytes(24).toString('base64url')
 }
 
-function parseMessage (raw) {
+function randomToken (): string {
+  return crypto.randomBytes(32).toString('base64url')
+}
+
+function parseMessage (raw: string | Buffer): RawMessage | undefined {
   try {
-    return JSON.parse(raw)
+    return JSON.parse(String(raw))
   } catch (err) {
     console.error(err)
 
@@ -133,7 +244,7 @@ function parseMessage (raw) {
   }
 }
 
-function sweepInactiveSessions (state) {
+function sweepInactiveSessions (state: SocketState): void {
   for (const [key, session] of state.inactiveSessions) {
     if (!isSessionActive(session)) {
       state.inactiveSessions.delete(key)
@@ -141,8 +252,11 @@ function sweepInactiveSessions (state) {
   }
 }
 
-function validateSchema (obj, validator) {
-  const headers = obj.headers
+function validateSchema<T extends Record<string, unknown>> (
+  obj: Record<string, unknown>,
+  validator: ValidateFunction<T>,
+): T {
+  const headers = obj.headers instanceof Headers
     ? Object.fromEntries(obj.headers)
     : undefined
 
@@ -152,19 +266,24 @@ function validateSchema (obj, validator) {
   }
 
   if (!validateNotMessage(payload)) {
-    const errors = validateNotMessage.errors.map(item => formatError('', item))
+    const errors = validateNotMessage.errors!.map(item => formatError('', item))
 
     throw new UnprocessableContentError(errors)
   }
 
   if (!validator(payload)) {
-    const errors = validator.errors.map(item => formatError('', item))
+    const errors = validator.errors!.map(item => formatError('', item))
 
     throw new UnprocessableContentError(errors)
   }
+
+  return obj as T
 }
 
-function matchesSegments (patternSegments, requestSegments) {
+function matchesSegments (
+  patternSegments: string[],
+  requestSegments: string[],
+): boolean {
   if (patternSegments.length !== requestSegments.length) {
     return false
   }
@@ -175,7 +294,10 @@ function matchesSegments (patternSegments, requestSegments) {
   )
 }
 
-function matchRoute (routes, message) {
+function matchRoute (
+  routes: SocketRoute[],
+  message: RequestMessage,
+): SocketRoute {
   const requestSegments = toSegments(message.route)
 
   const matchingPaths = routes.filter(route =>
@@ -195,17 +317,24 @@ function matchRoute (routes, message) {
   return route
 }
 
-function buildParams (route, message) {
+function buildParams (
+  route: SocketRoute,
+  message: RequestMessage,
+): Record<string, string> {
   const requestSegments = toSegments(message.route)
 
-  return route.segments.reduce((accum, segment, index) =>
-    segment.startsWith(':') ? {
-      ...accum,
-      [segment.slice(1)]: requestSegments[index],
-    } : accum, {})
+  return route.segments.reduce<Record<string, string>>(
+    (accum, segment, index) =>
+      segment.startsWith(':') ? {
+        ...accum,
+        [segment.slice(1)]: requestSegments[index],
+      } : accum, {})
 }
 
-function buildRequest (params, message) {
+function buildRequest (
+  params: Record<string, string>,
+  message: RequestMessage,
+): WebSocketRequest {
   const { id, clientId, method, route } = message
   const headers = new Headers(message.headers ?? {})
   const query = message.query ?? {}
@@ -223,13 +352,17 @@ function buildRequest (params, message) {
   }
 }
 
-async function buildOutgoingMessage (id, clientId, response) {
+async function buildOutgoingMessage (
+  id: string,
+  clientId: string,
+  response: Response,
+): Promise<ResponseMessage> {
   const text = await response.text()
   const contentType = response.headers.get('content-type') ?? ''
   const usingJson = contentType.includes('application/json')
   const body = usingJson ? JSON.parse(text) : text
 
-  return createMessage(clientId, TYPES.RESPONSE, {
+  return createMessage(clientId, MessageType.Response, {
     id,
     status: response.status,
     headers: response.headers,
@@ -237,7 +370,34 @@ async function buildOutgoingMessage (id, clientId, response) {
   })
 }
 
-export function buildSocketState (opts = {}) {
+function buildErrorMessage (
+  message: RawMessage,
+  err: unknown,
+): ResponseMessage {
+  const { id, clientId } = message as Pick<BaseMessage, 'id' | 'clientId'>
+
+  if (err instanceof RequestError) {
+    const ctor = err.constructor as typeof RequestError
+
+    return createMessage(clientId, MessageType.Response, {
+      id,
+      status: ctor.status,
+      headers: new Headers({
+        'content-type': 'application/json;charset=utf-8',
+      }),
+      body: err.output,
+    })
+  }
+
+  return createMessage(clientId, MessageType.Response, {
+    id,
+    status: InternalServerError.status,
+    headers: new Headers(),
+    body: err instanceof Error ? err.message : undefined,
+  })
+}
+
+export function buildSocketState (opts: AppOptions = {}): SocketState {
   return {
     disconnectThreshold: opts.ws?.disconnectThreshold ?? 120_000,
     heartbeatInterval: opts.ws?.heartbeatInterval ?? 30_000,
@@ -250,7 +410,10 @@ export function buildSocketState (opts = {}) {
   }
 }
 
-export function buildSocketServer (routes, state) {
+export function buildSocketServer (
+  routes: SocketRoute[],
+  state: SocketState,
+): WebSocketHandler<SocketData> {
   const {
     disconnectThreshold,
     heartbeatInterval,
@@ -259,8 +422,10 @@ export function buildSocketServer (routes, state) {
     inactiveSessions,
   } = state
 
-  function armReaper (ws) {
-    clearTimeout(ws.data.reaperHandle)
+  function armReaper (ws: SocketConnection): void {
+    if (ws.data.reaperHandle) {
+      clearTimeout(ws.data.reaperHandle)
+    }
 
     ws.data.reaperHandle = setTimeout(() => {
       ws.data.reaped = true
@@ -270,10 +435,10 @@ export function buildSocketServer (routes, state) {
   }
 
   return {
-    open (ws) {
+    open (ws: SocketConnection): void {
       sweepInactiveSessions(state)
 
-      const token = randomToken(32)
+      const token = randomToken()
       const existingSession = activeSessions.get(ws.data.clientId)
 
       if (existingSession) {
@@ -293,9 +458,9 @@ export function buildSocketServer (routes, state) {
 
       const welcomeMessage = createMessage(
         ws.data.clientId,
-        TYPES.WELCOME,
+        MessageType.Welcome,
         {
-          headers: {},
+          headers: new Headers(),
           body: {
             heartbeatInterval,
             token,
@@ -305,8 +470,10 @@ export function buildSocketServer (routes, state) {
 
       ws.send(JSON.stringify(welcomeMessage))
     },
-    close (ws, code) {
-      clearTimeout(ws.data.reaperHandle)
+    close (ws: SocketConnection, code: number): void {
+      if (ws.data.reaperHandle) {
+        clearTimeout(ws.data.reaperHandle)
+      }
 
       if (ws.data.superseded) {
         return
@@ -327,7 +494,7 @@ export function buildSocketServer (routes, state) {
         })
       }
     },
-    async message (ws, raw) {
+    async message (ws: SocketConnection, raw: string | Buffer): Promise<void> {
       const incomingMsg = parseMessage(raw)
 
       if (incomingMsg === undefined) {
@@ -337,21 +504,21 @@ export function buildSocketServer (routes, state) {
       armReaper(ws)
 
       try {
-        validateMessage(incomingMsg)
+        const message = validateMessage(incomingMsg)
 
-        if (incomingMsg.type === TYPES.HEARTBEAT) {
-          const { id, clientId } = incomingMsg
-          const ack = createMessage(clientId, TYPES.HEARTBEAT, { id })
+        if (message.type === MessageType.Heartbeat) {
+          const { id, clientId } = message
+          const ack = createMessage(clientId, MessageType.Heartbeat, { id })
 
           ws.send(JSON.stringify(ack))
 
           return
         }
 
-        const { id, clientId } = incomingMsg
-        const route = matchRoute(routes, incomingMsg)
-        const params = buildParams(route, incomingMsg)
-        const req = buildRequest(params, incomingMsg)
+        const { id, clientId } = message
+        const route = matchRoute(routes, message)
+        const params = buildParams(route, message)
+        const req = buildRequest(params, message)
         const res = await executeMiddlewareChain(req, route.chain)
         const outgoingMsg = await buildOutgoingMessage(id, clientId, res)
 
@@ -359,20 +526,7 @@ export function buildSocketServer (routes, state) {
       } catch (err) {
         console.error(err)
 
-        const { id, clientId } = incomingMsg
-        const status = err.constructor.status ?? InternalServerError.status
-        const body = err.output !== undefined ? err.output : err.message
-
-        const headers = err.output !== undefined
-          ? { 'content-type': 'application/json;charset=utf-8' }
-          : {}
-
-        const res = createMessage(clientId, TYPES.RESPONSE, {
-          id,
-          status,
-          headers,
-          body,
-        })
+        const res = buildErrorMessage(incomingMsg, err)
 
         ws.send(JSON.stringify(res))
       }
@@ -380,7 +534,7 @@ export function buildSocketServer (routes, state) {
   }
 }
 
-export function buildSocketHandlers (state) {
+export function buildSocketHandlers (state: SocketState): SocketEndpoint[] {
   const {
     maxTickets,
     ticketTtl,
@@ -389,7 +543,7 @@ export function buildSocketHandlers (state) {
     inactiveSessions,
   } = state
 
-  function bindTicket (clientId) {
+  function bindTicket (clientId: string): string {
     for (const [key, entry] of tickets) {
       if (entry.expiresAt > Date.now()) {
         break
@@ -399,10 +553,10 @@ export function buildSocketHandlers (state) {
     }
 
     if (tickets.size >= maxTickets) {
-      throw new ServiceUnavailableError()
+      throw new ServiceUnavailableError('Unable to issue ticket')
     }
 
-    const ticket = randomToken(24)
+    const ticket = randomTicket()
     const expiresAt = Date.now() + ticketTtl
 
     tickets.set(ticket, {
@@ -413,7 +567,7 @@ export function buildSocketHandlers (state) {
     return ticket
   }
 
-  function redeemTicket (ticket) {
+  function redeemTicket (ticket: string): string | undefined {
     const entry = ticket ? tickets.get(ticket) : undefined
 
     if (!entry) {
@@ -433,23 +587,26 @@ export function buildSocketHandlers (state) {
     {
       method: 'GET',
       path: '/ws',
-      handler (req, res) {
-        validateSchema(req, createSocketValidator)
+      handler (req: Record<string, unknown>, res: unknown): Response {
+        const validReq = validateSchema(req, createSocketValidator)
 
         if (typeof res !== 'object') {
           throw new TypeError('Endpoint "res" must be an object')
         }
 
-        const ctx = res ? { ...res } : {}
+        const ctx: UpgradeContext = res ? { ...res } : {}
 
         ctx.data = ctx.data ?? {}
-        ctx.data.clientId = redeemTicket(req.query.ticket)
+        ctx.data.clientId = redeemTicket(validReq.query.ticket)
+        ctx.data.superseded = false
+        ctx.data.reaped = false
+        ctx.data.reaperHandle = null
 
         if (!ctx.data.clientId) {
           throw new NotFoundError()
         }
 
-        const useSocket = req.server.upgrade(req.raw, ctx)
+        const useSocket = validReq.server.upgrade(validReq.raw, ctx)
 
         if (!useSocket) {
           throw new NotFoundError()
@@ -461,7 +618,7 @@ export function buildSocketHandlers (state) {
     {
       method: 'POST',
       path: '/ws',
-      handler (req, res) {
+      handler (req: Record<string, unknown>, res: unknown): Response {
         validateSchema(req, createTicketValidator)
 
         const clientId = crypto.randomUUID()
@@ -470,25 +627,25 @@ export function buildSocketHandlers (state) {
           clientId,
           ticket: bindTicket(clientId),
           data: res,
-        }, { status: 201 })
+        }, { status: StatusCode.Created })
       },
     },
     {
       method: 'PUT',
       path: '/ws/:clientId',
-      handler (req, res) {
-        validateSchema(req, updateTicketValidator)
-
-        const authHeader = req.headers.get('authorization')
-        const token = authHeader.slice('Bearer '.length)
+      handler (req: Record<string, unknown>, res: unknown): Response {
+        const validReq = validateSchema(req, updateTicketValidator)
+        const authHeader = validReq.headers.get('authorization')
+        const token = authHeader!.slice('Bearer '.length)
 
-        let session = activeSessions.get(req.params.clientId)
+        let session: Session | undefined =
+          activeSessions.get(validReq.params.clientId)
 
         if (!session) {
-          const inactive = inactiveSessions.get(req.params.clientId)
+          const inactive = inactiveSessions.get(validReq.params.clientId)
 
           if (inactive && !isSessionActive(inactive)) {
-            inactiveSessions.delete(req.params.clientId)
+            inactiveSessions.delete(validReq.params.clientId)
           } else {
             session = inactive
           }
@@ -499,12 +656,12 @@ export function buildSocketHandlers (state) {
         }
 
         if (session.token !== token) {
-          throw new UnauthorizedError()
+          throw new UnauthorizedError('Invalid token')
         }
 
         return Response.json({
-          clientId: req.params.clientId,
-          ticket: bindTicket(req.params.clientId),
+          clientId: validReq.params.clientId,
+          ticket: bindTicket(validReq.params.clientId),
           data: res,
         })
       },
@@ -512,17 +669,21 @@ export function buildSocketHandlers (state) {
   ]
 }
 
-export function buildSocketCommands (state) {
-  function sendToClient (clientId, event, body) {
+export function buildSocketCommands (state: SocketState): SocketCommands {
+  function sendToClient (
+    clientId: string,
+    event: string,
+    body: unknown,
+  ): void {
     const session = state.activeSessions.get(clientId)
 
     if (!session) {
       throw new ReferenceError(`No live socket for client: ${clientId}`)
     }
 
-    const message = createMessage(clientId, TYPES.NOTIFICATION, {
+    const message = createMessage(clientId, MessageType.Notification, {
       event,
-      headers: {},
+      headers: new Headers(),
       body,
     })