hypermosh-shell / lib / relay.mjs

// relay — tunnel UDP datagrams over a hyperdht connection.
//
// The whole point: mosh's SSP already survives packet loss, latency, and
// roaming — it just needs *a* datagram path to the server. celltron (a NAT'd,
// Doze-prone phone) has no reachable UDP port, so we carry mosh's UDP over a
// hole-punched, identity-keyed hyperdht connection. mosh provides the
// resilience; this relay is a dumb UDP<->stream bridge with auto-reconnect
// (which becomes mosh's roaming when the phone's network flaps).
//
// Dual-runtime (node + bare) via the `#node:dgram` import map.

import DHT from 'hyperdht'
import dgram from '#node:dgram'
import b4a from 'b4a'

// Datagrams are length-prefixed ([uint16 BE len][payload]) so they reassemble
// cleanly out of the reliable UDX byte stream.
export function makeFramer (onDatagram) {
  let buf = b4a.alloc(0)
  return (chunk) => {
    buf = b4a.concat([buf, chunk])
    while (buf.length >= 2) {
      const len = buf.readUInt16BE(0)
      if (buf.length < 2 + len) break
      onDatagram(buf.subarray(2, 2 + len))
      buf = buf.subarray(2 + len)
    }
  }
}

export function frame (dg) {
  const h = b4a.alloc(2)
  h.writeUInt16BE(dg.length, 0)
  return b4a.concat([h, dg])
}

// SERVER side (runs next to mosh-server, e.g. on celltron): accept identity-
// firewalled hyperdht connections; bridge each to the local mosh-server UDP
// port. `allow` is `true` (any) or an array of allowed public keys.
export async function serve ({ keyPair, udpHost = '127.0.0.1', udpPort, allow = true, node = null } = {}) {
  if (!udpPort) throw new Error('serve requires udpPort (the local mosh-server UDP port)')
  const dht = node || new DHT()

  const firewall = (remotePublicKey) =>
    !(allow === true || allow.some((k) => b4a.equals(remotePublicKey, k))) // true = DENY

  const server = dht.createServer({ firewall }, (socket) => {
    const udp = dgram.createSocket('udp4')
    // mosh-server -> tunnel
    udp.on('message', (msg) => { try { socket.write(frame(msg)) } catch {} })
    // tunnel -> mosh-server
    socket.on('data', makeFramer((dg) => { try { udp.send(dg, udpPort, udpHost) } catch {} }))
    socket.on('error', () => {})
    socket.on('close', () => { try { udp.close() } catch {} })
  })

  await server.listen(keyPair)
  return {
    dht,
    server,
    publicKey: keyPair.publicKey,
    async close () { await server.close(); if (!node) await dht.destroy() }
  }
}

// CLIENT side (runs next to mosh-client): expose a local UDP port that
// mosh-client targets; bridge it to the remote server over hyperdht. The peer
// (mosh-client) UDP address is learned from its first packet. On stream close
// we redial — that reconnect IS mosh's roaming across network changes.
//
// NOTE: writes are NOT guarded on socket.writable — UDX buffers pre-connection
// writes, so mosh's earliest packets survive the ~seconds of DHT connect
// instead of being dropped (the bug that made an early prototype time out).
export async function connect ({ remoteKey, keyPair = null, localHost = '127.0.0.1', localPort = 0, node = null, reconnectMs = 500 } = {}) {
  const dht = node || new DHT()
  const udp = dgram.createSocket('udp4')
  let peer = null
  let socket = null
  let closed = false

  const dial = () => {
    if (closed) return
    socket = dht.connect(remoteKey, keyPair ? { keyPair } : undefined)
    socket.on('data', makeFramer((dg) => { if (peer) { try { udp.send(dg, peer.port, peer.address) } catch {} } }))
    socket.on('error', () => {})
    socket.on('close', () => { socket = null; if (!closed) setTimeout(dial, reconnectMs) })
  }

  udp.on('message', (msg, rinfo) => {
    peer = rinfo
    if (socket) { try { socket.write(frame(msg)) } catch {} } // buffers until connected
  })

  await new Promise((resolve) => udp.bind(localPort, localHost, resolve))
  dial()

  return {
    dht,
    udp,
    localPort: udp.address().port,
    async close () { closed = true; try { udp.close() } catch {}; if (socket) socket.destroy(); if (!node) await dht.destroy() }
  }
}