dsh-remote/lib/proxy.mjs

// HTTP + WebSocket tunnelling between a browser-facing webserver and a remote DSH over
// the link mux.
//
//   viewer side (where the browser is):   /r/<device>/<rest>  ──stream "http"|"ws"──►
//   owner side (where the session runs):  forward to 127.0.0.1:<own port>/<rest>
//
// Owner side rewrites Host/Origin to its own loopback authority and attaches a freshly
// minted DSH auth cookie, so DSH's /api trust fence and cookie check pass. The viewer's
// own DSH cookie is stripped before forwarding (it would be meaningless remotely).
import http from 'node:http'

const HOP = new Set([
  'connection', 'keep-alive', 'proxy-connection', 'transfer-encoding', 'upgrade', 'te', 'trailer',
  'proxy-authorization', 'proxy-authenticate', 'host', 'content-length',
])

export function stripDshCookies(cookieHeader) {
  return String(cookieHeader || '')
    .split(';')
    .map((s) => s.trim())
    .filter((s) => s && !/^dsh-auth-/.test(s) && !/^dsh_auth_session=/.test(s))
    .join('; ')
}

function forwardHeaders(headers) {
  const out = {}
  for (const [k, v] of Object.entries(headers || {})) {
    if (HOP.has(k.toLowerCase())) continue
    out[k] = v
  }
  return out
}

// ---------------------------------------------------------------------------
// viewer side
// ---------------------------------------------------------------------------

/**
 * Pipe one browser HTTP request through `openStream()` (returns a mux Stream).
 * `rest` is the path+query to request on the remote side ("/api/remote/…").
 * `basePath` is the local mount ("/r/ai-pc") used to rewrite redirects/cookies.
 */
export function viewHttp(req, res, openStream, { rest, basePath, extra = {} }) {
  const headers = forwardHeaders(req.headers)
  headers.cookie = stripDshCookies(req.headers.cookie)
  if (!headers.cookie) delete headers.cookie
  const stream = openStream('http', { method: req.method, path: rest, headers, ...extra })
  let responded = false

  stream.on('head', (meta) => {
    responded = true
    const outHeaders = { ...meta.headers }
    // Keep relative redirects inside the mount (DSH answers "./" which is already relative).
    if (typeof outHeaders.location === 'string' && outHeaders.location.startsWith('/')) {
      outHeaders.location = basePath + outHeaders.location
    }
    // The remote never sets cookies the browser should keep under our authority.
    delete outHeaders['set-cookie']
    try { res.writeHead(meta.status || 502, outHeaders) } catch { /* headers already sent */ }
  })
  stream.on('data', (chunk) => {
    if (!res.write(chunk)) {
      stream.mux.sendControl({ t: 'pause', id: stream.id })
      res.once('drain', () => stream.mux.sendControl({ t: 'resume', id: stream.id }))
    }
  })
  stream.on('end', () => res.end())
  stream.on('error', (err) => {
    if (!responded) {
      res.writeHead(502, { 'content-type': 'text/plain; charset=utf-8', 'cache-control': 'no-store' })
      res.end(`dsh-remote: ${err.message}\n`)
    } else {
      res.destroy()
    }
  })
  res.on('close', () => { if (!stream.done) stream.fail('client closed') })

  req.on('data', (chunk) => {
    if (!stream.write(chunk)) {
      req.pause()
      stream.once('drain', () => req.resume())
    }
  })
  req.on('end', () => stream.end())
  req.on('error', () => stream.fail('request error'))
}

/** Tunnel a browser WebSocket upgrade (raw socket) through a mux stream. */
export function viewUpgrade(req, socket, head, openStream, { rest, extra = {} }) {
  const headers = forwardHeaders(req.headers)
  // Keep the WebSocket handshake headers (forwardHeaders dropped connection/upgrade).
  headers.cookie = stripDshCookies(req.headers.cookie)
  if (!headers.cookie) delete headers.cookie
  const stream = openStream('ws', { path: rest, headers, ...extra })
  let open = false
  stream.on('head', (meta) => {
    // The owner side already completed the upstream handshake; replay its 101 to the
    // browser with the browser's own Sec-WebSocket-Key accepted value.
    if (meta.status !== 101) {
      socket.end(`HTTP/1.1 ${meta.status || 502} Upgrade Failed\r\nConnection: close\r\nContent-Length: 0\r\n\r\n`)
      return
    }
    open = true
    const lines = ['HTTP/1.1 101 Switching Protocols']
    for (const [k, v] of Object.entries(meta.headers || {})) {
      for (const one of Array.isArray(v) ? v : [v]) lines.push(`${k}: ${one}`)
    }
    socket.write(lines.join('\r\n') + '\r\n\r\n')
    if (head && head.length) stream.write(head)
  })
  stream.on('data', (chunk) => { if (!socket.write(chunk)) { /* tcp backpressure handled by kernel */ } })
  stream.on('end', () => socket.end())
  stream.on('error', () => {
    if (!open) {
      try { socket.end('HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\nContent-Length: 0\r\n\r\n') } catch { socket.destroy() }
    } else socket.destroy()
  })
  socket.on('data', (chunk) => stream.write(chunk))
  socket.on('end', () => stream.end())
  socket.on('close', () => { if (!stream.done) stream.fail('client closed') })
  socket.on('error', () => { if (!stream.done) stream.fail('client error') })
}

// ---------------------------------------------------------------------------
// owner side
// ---------------------------------------------------------------------------

/**
 * Serve one incoming "http" stream against the local DSH.
 * @param target {host, port}; authority = `${host}:${port}`
 * @param cookie async () => "dsh-auth-…=v1…" for that authority
 */
export async function ownHttp(stream, { target, cookie }) {
  const { method = 'GET', path = '/', headers = {} } = stream.meta
  const authority = `${target.host}:${target.port}`
  const h = { ...headers, host: authority }
  if (h.origin) h.origin = `http://${authority}`
  if (h.referer) h.referer = `http://${authority}/`
  delete h['sec-fetch-site']
  const c = await cookie(authority)
  h.cookie = [h.cookie, c].filter(Boolean).join('; ')

  const req = http.request({ host: target.host, port: target.port, method, path, headers: h })
  req.on('socket', (s) => s.on('error', () => {}))
  req.on('response', (res) => {
    stream.head({ status: res.statusCode, headers: forwardHeaders(res.headers) })
    res.on('data', (chunk) => {
      if (!stream.write(chunk)) {
        res.pause()
        stream.once('drain', () => res.resume())
      }
    })
    res.on('end', () => stream.end())
    res.on('error', () => stream.fail('upstream response error'))
    stream.on('remote-pause', () => res.pause())
    stream.on('remote-resume', () => res.resume())
  })
  req.on('error', (err) => stream.fail(`upstream: ${err.message}`))
  stream.on('data', (chunk) => req.write(chunk))
  stream.on('end', () => req.end())
  stream.on('error', () => req.destroy())
  stream.start()
}

/** Serve one incoming "ws" stream: perform the upgrade against local DSH, then pipe. */
export async function ownUpgrade(stream, { target, cookie }) {
  const { path = '/', headers = {} } = stream.meta
  const authority = `${target.host}:${target.port}`
  const h = { ...headers, host: authority, connection: 'Upgrade', upgrade: headers.upgrade || 'websocket' }
  if (h.origin) h.origin = `http://${authority}`
  delete h['sec-fetch-site']
  const c = await cookie(authority)
  h.cookie = [h.cookie, c].filter(Boolean).join('; ')

  const req = http.request({ host: target.host, port: target.port, method: 'GET', path, headers: h })
  req.on('socket', (s) => s.on('error', () => {}))
  req.on('upgrade', (res, sock, head) => {
    const out = {}
    for (const [k, v] of Object.entries(res.headers)) out[k] = v
    stream.head({ status: 101, headers: out })
    if (head && head.length) stream.write(head)
    sock.on('data', (chunk) => stream.write(chunk))
    sock.on('end', () => stream.end())
    sock.on('error', () => stream.fail('upstream socket error'))
    stream.on('data', (chunk) => sock.write(chunk))
    stream.on('end', () => sock.end())
    stream.on('error', () => sock.destroy())
    stream.on('close', () => sock.destroy())
    // Browser frames sent before the upstream handshake finished were queued; flush now.
    stream.start()
  })
  req.on('response', (res) => {
    stream.head({ status: res.statusCode || 502, headers: {} })
    res.resume()
    stream.end()
  })
  req.on('error', (err) => stream.fail(`upstream: ${err.message}`))
  req.end()
}