rnxsim 0.0.0 → 0.1.312

package/src/host/websocket-proxy.ts

@@ -0,0 +1,210 @@
+import { WebSocket, WebSocketServer } from 'ws'
+import type { IncomingMessage } from 'http'
+import type { Duplex } from 'stream'
+
+export const WEBSOCKET_PROXY_PATH = '/__websocket-proxy'
+
+const STRIP_UPSTREAM_HEADERS = new Set([
+  'host',
+  'connection',
+  'upgrade',
+  'transfer-encoding',
+  'content-length',
+  'sec-websocket-accept',
+  'sec-websocket-extensions',
+  'sec-websocket-key',
+  'sec-websocket-protocol',
+  'sec-websocket-version',
+])
+
+function rejectUpgrade(socket: Duplex, status: number, message: string) {
+  try {
+    socket.write(
+      `HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Type: text/plain\r\nContent-Length: ${message.length}\r\n\r\n${message}`,
+    )
+  } catch {}
+  socket.destroy()
+}
+
+function isSameOriginUpgrade(req: IncomingMessage): boolean {
+  const origin = req.headers.origin
+  const host = req.headers.host
+  if (!origin || !host) return false
+  try {
+    return new URL(origin).host === host
+  } catch {
+    return false
+  }
+}
+
+function decodeProxyHeaders(encoded: string | null): Record<string, string> {
+  if (!encoded) return {}
+  const base64 = encoded.replace(/-/g, '+').replace(/_/g, '/')
+  const padded = base64 + '='.repeat((4 - (base64.length % 4)) % 4)
+  const parsed = JSON.parse(Buffer.from(padded, 'base64').toString('utf8'))
+  if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return {}
+  const headers: Record<string, string> = {}
+  for (const [key, value] of Object.entries(parsed)) {
+    if (value == null) continue
+    if (STRIP_UPSTREAM_HEADERS.has(key.toLowerCase())) continue
+    headers[key] = Array.isArray(value) ? value.join(', ') : String(value)
+  }
+  return headers
+}
+
+function getDefaultWebSocketOrigin(targetUrl: URL): string {
+  const origin = new URL(targetUrl.href)
+  origin.protocol = targetUrl.protocol === 'wss:' ? 'https:' : 'http:'
+  return origin.origin
+}
+
+function getRequestedProtocols(req: IncomingMessage): string[] {
+  const header = req.headers['sec-websocket-protocol']
+  const value = Array.isArray(header) ? header.join(',') : header || ''
+  return value
+    .split(',')
+    .map((part) => part.trim())
+    .filter(Boolean)
+}
+
+function safeClose(ws: WebSocket, code: number, reason: string) {
+  if (ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) {
+    return
+  }
+  try {
+    ws.close(code, reason)
+  } catch {
+    ws.terminate()
+  }
+}
+
+function connectProxyPair(clientWs: WebSocket, upstream: WebSocket) {
+  let closing = false
+
+  const closeBoth = (
+    source: WebSocket,
+    target: WebSocket,
+    code: number,
+    reason: Buffer,
+  ) => {
+    if (closing) return
+    closing = true
+    safeClose(target, code, reason.toString())
+    if (source.readyState === WebSocket.OPEN) {
+      safeClose(source, code, reason.toString())
+    }
+  }
+
+  clientWs.on('message', (data, isBinary) => {
+    if (upstream.readyState === WebSocket.OPEN) {
+      upstream.send(data, { binary: isBinary })
+    }
+  })
+  upstream.on('message', (data, isBinary) => {
+    if (clientWs.readyState === WebSocket.OPEN) {
+      clientWs.send(data, { binary: isBinary })
+    }
+  })
+
+  clientWs.on('close', (code, reason) => closeBoth(clientWs, upstream, code, reason))
+  upstream.on('close', (code, reason) => closeBoth(upstream, clientWs, code, reason))
+  clientWs.on('error', () => safeClose(upstream, 1011, 'proxy client error'))
+  upstream.on('error', () => safeClose(clientWs, 1011, 'upstream websocket error'))
+}
+
+function createUpstreamWebSocket(
+  targetUrl: URL,
+  protocols: string[],
+  headers: Record<string, string>,
+  forwardOrigin: boolean,
+): WebSocket {
+  const upstreamHeaders = { ...headers }
+  const originKeys = Object.keys(upstreamHeaders).filter(
+    (key) => key.toLowerCase() === 'origin',
+  )
+  if (forwardOrigin && originKeys.length === 0) {
+    upstreamHeaders.origin = getDefaultWebSocketOrigin(targetUrl)
+  } else if (!forwardOrigin) {
+    for (const key of originKeys) delete upstreamHeaders[key]
+  }
+  if (Object.keys(upstreamHeaders).length === 0)
+    return new WebSocket(targetUrl.href, protocols)
+  return new WebSocket(targetUrl.href, protocols, {
+    headers: upstreamHeaders,
+  })
+}
+
+export function isWebSocketProxyRequestUrl(rawUrl: string | undefined): boolean {
+  if (!rawUrl) return false
+  try {
+    return new URL(rawUrl, 'http://localhost').pathname === WEBSOCKET_PROXY_PATH
+  } catch {
+    return false
+  }
+}
+
+export function handleWebSocketProxyUpgrade(
+  req: IncomingMessage,
+  socket: Duplex,
+  head: Buffer,
+  forwardOrigin = true,
+): boolean {
+  if (!isWebSocketProxyRequestUrl(req.url)) return false
+  if (!isSameOriginUpgrade(req)) {
+    rejectUpgrade(socket, 403, 'forbidden websocket proxy origin')
+    return true
+  }
+
+  let targetUrl: URL
+  let headers: Record<string, string>
+  try {
+    const requestUrl = new URL(req.url || '/', 'http://localhost')
+    const target = requestUrl.searchParams.get('url')
+    if (!target) {
+      rejectUpgrade(socket, 400, 'missing websocket proxy url')
+      return true
+    }
+    targetUrl = new URL(target)
+    if (targetUrl.protocol !== 'ws:' && targetUrl.protocol !== 'wss:') {
+      rejectUpgrade(socket, 400, 'invalid websocket proxy protocol')
+      return true
+    }
+    headers = decodeProxyHeaders(requestUrl.searchParams.get('headers'))
+  } catch {
+    rejectUpgrade(socket, 400, 'invalid websocket proxy request')
+    return true
+  }
+
+  const protocols = getRequestedProtocols(req)
+  const upstream = createUpstreamWebSocket(targetUrl, protocols, headers, forwardOrigin)
+  let completed = false
+  socket.once('close', () => {
+    if (!completed) upstream.terminate()
+  })
+  upstream.once('open', () => {
+    if (completed) return
+    completed = true
+    const selectedProtocol = upstream.protocol
+    const proxyServer = new WebSocketServer({
+      noServer: true,
+      clientTracking: false,
+      handleProtocols(requestedProtocols) {
+        return selectedProtocol || requestedProtocols.values().next().value || false
+      },
+    })
+    proxyServer.handleUpgrade(req, socket, head, (clientWs) => {
+      connectProxyPair(clientWs, upstream)
+    })
+  })
+  upstream.once('error', () => {
+    if (completed) return
+    completed = true
+    rejectUpgrade(socket, 502, 'upstream websocket error')
+  })
+  upstream.once('close', () => {
+    if (completed) return
+    completed = true
+    rejectUpgrade(socket, 502, 'upstream websocket closed')
+  })
+  return true
+}