kestrelos/server/plugins/websocket.js

import { WebSocketServer } from 'ws'
import { getDb } from '../utils/db.js'
import { handleWebSocketMessage } from '../utils/webrtcSignaling.js'

function parseCookie(cookieHeader) {
  const cookies = {}
  if (!cookieHeader) return cookies
  cookieHeader.split(';').forEach((cookie) => {
    const [name, ...valueParts] = cookie.trim().split('=')
    if (name && valueParts.length > 0) {
      cookies[name] = decodeURIComponent(valueParts.join('='))
    }
  })
  return cookies
}

let wss = null
const connections = new Map()

export function getWebSocketServer() {
  return wss
}

export function getSessionConnections(sessionId) {
  return connections.get(sessionId) || new Set()
}

export function addSessionConnection(sessionId, ws) {
  if (!connections.has(sessionId)) {
    connections.set(sessionId, new Set())
  }
  connections.get(sessionId).add(ws)
}

export function removeSessionConnection(sessionId, ws) {
  const conns = connections.get(sessionId)
  if (conns) {
    conns.delete(ws)
    if (conns.size === 0) {
      connections.delete(sessionId)
    }
  }
}

export function broadcastToSession(sessionId, message) {
  const conns = getSessionConnections(sessionId)
  const data = JSON.stringify(message)
  for (const ws of conns) {
    if (ws.readyState === 1) { // OPEN
      ws.send(data)
    }
  }
}

export default defineNitroPlugin((nitroApp) => {
  nitroApp.hooks.hook('ready', async () => {
    const server = nitroApp.h3App.server || nitroApp.h3App.nodeServer
    if (!server) {
      console.warn('[websocket] Could not attach to HTTP server')
      return
    }

    wss = new WebSocketServer({
      server,
      path: '/ws',
      verifyClient: async (info, callback) => {
        // Verify session cookie on upgrade request
        const cookies = parseCookie(info.req.headers.cookie || '')
        const sessionId = cookies.session_id
        if (!sessionId) {
          callback(false, 401, 'Unauthorized')
          return
        }

        try {
          const { get } = await getDb()
          const session = await get('SELECT user_id, expires_at FROM sessions WHERE id = ?', [sessionId])
          if (!session || new Date(session.expires_at) < new Date()) {
            callback(false, 401, 'Unauthorized')
            return
          }
          // Store user_id in request for later use
          info.req.userId = session.user_id
          callback(true)
        }
        catch (err) {
          console.error('[websocket] Auth error:', err)
          callback(false, 500, 'Internal Server Error')
        }
      },
    })

    wss.on('connection', (ws, req) => {
      const userId = req.userId
      if (!userId) {
        ws.close(1008, 'Unauthorized')
        return
      }

      let currentSessionId = null

      ws.on('message', async (data) => {
        try {
          const message = JSON.parse(data.toString())
          const { sessionId, type } = message

          if (!sessionId || !type) {
            ws.send(JSON.stringify({ error: 'Invalid message format' }))
            return
          }

          // Track session connection
          if (currentSessionId !== sessionId) {
            if (currentSessionId) {
              removeSessionConnection(currentSessionId, ws)
            }
            currentSessionId = sessionId
            addSessionConnection(sessionId, ws)
          }

          // Handle WebRTC signaling message
          const response = await handleWebSocketMessage(userId, sessionId, type, message.data || {})
          if (response) {
            ws.send(JSON.stringify(response))
          }
        }
        catch (err) {
          console.error('[websocket] Message error:', err)
          ws.send(JSON.stringify({ error: err.message || 'Internal error' }))
        }
      })

      ws.on('close', () => {
        if (currentSessionId) {
          removeSessionConnection(currentSessionId, ws)
        }
      })

      ws.on('error', (err) => {
        console.error('[websocket] Connection error:', err)
      })
    })

    console.log('[websocket] WebSocket server started on /ws')
  })

  nitroApp.hooks.hook('close', () => {
    if (wss) {
      wss.close()
      wss = null
    }
  })
})