tronbyt/src/proxy.ts

import type { ServerWebSocket } from 'bun'

export interface WsData {
  path: string
  protocols: string[]
}

const upstreams = new Map<ServerWebSocket<WsData>, WebSocket>()

export function createProxy(socketPath: string, isHealthy: () => boolean, isRunning: () => boolean) {
  function proxyFetch(req: Request): Promise<Response> | Response {
    const url = new URL(req.url)

    if (url.pathname === '/ok') {
      if (!isHealthy()) return new Response('starting', { status: isRunning() ? 200 : 503 })
      return fetch('http://localhost/health', { unix: socketPath })
        .then((r) => (r.ok ? new Response('ok') : new Response('unhealthy', { status: 503 })))
        .catch(() => new Response('unhealthy', { status: 503 }))
    }

    const hasBody = req.method !== 'GET' && req.method !== 'HEAD'
    const headers = new Headers(req.headers)
    headers.delete('accept-encoding')

    return fetch(`http://localhost${url.pathname}${url.search}`, {
      method: req.method,
      headers,
      body: hasBody ? req.body : undefined,
      unix: socketPath,
    }).then((r) => {
      const respHeaders = new Headers(r.headers)
      respHeaders.delete('content-encoding')
      respHeaders.delete('content-length')
      return new Response(r.body, { status: r.status, headers: respHeaders })
    }).catch((e) => {
      console.error('Proxy error:', e)
      return new Response('Tronbyt server is not responding', { status: 502 })
    })
  }

  const websocket = {
    open(ws: ServerWebSocket<WsData>) {
      const upstream = new WebSocket(`ws+unix://${socketPath}:${ws.data.path}`, ws.data.protocols)
      upstream.binaryType = 'arraybuffer'
      upstreams.set(ws, upstream)

      upstream.addEventListener('message', (e) => ws.send(e.data as string | ArrayBuffer))
      upstream.addEventListener('close', () => { upstreams.delete(ws); ws.close() })
      upstream.addEventListener('error', () => { upstreams.delete(ws); ws.close() })
    },

    message(ws: ServerWebSocket<WsData>, msg: string | ArrayBuffer | Uint8Array) {
      const upstream = upstreams.get(ws)
      if (upstream?.readyState === WebSocket.OPEN) upstream.send(msg)
    },

    close(ws: ServerWebSocket<WsData>) {
      upstreams.get(ws)?.close()
      upstreams.delete(ws)
    },
  }

  return { proxyFetch, websocket }
}