dev_arc_aws/backend/src/tunnels/manager.ts

import net from 'node:net'
import type { Client } from 'ssh2'
import { db } from '../db/index.js'
import { loadSshHost, connectTarget } from '../ssh/connect.js'
import { readSocks5Target, sendSocks5Success, sendSocks5Failure } from './socks5.js'

export type TunnelMode = 'local' | 'remote' | 'dynamic'
export type TunnelStatus = 'stopped' | 'connecting' | 'connected' | 'retrying' | 'error'

export interface TunnelRow {
  id: number
  name: string
  integration_id: number
  mode: TunnelMode
  source_port: number
  endpoint_host: string
  endpoint_port: number
  auto_start: number
  max_retries: number
  retry_interval_ms: number
  created_at: string
}

interface RuntimeState {
  status: TunnelStatus
  error: string | null
  retryCount: number
  client: Client | null
  jumpConn: Client | null
  server: net.Server | null
  retryTimer: NodeJS.Timeout | null
  stopRequested: boolean
}

const runtimes = new Map<number, RuntimeState>()

function emptyState(): RuntimeState {
  return {
    status: 'stopped',
    error: null,
    retryCount: 0,
    client: null,
    jumpConn: null,
    server: null,
    retryTimer: null,
    stopRequested: false,
  }
}

function getState(id: number): RuntimeState {
  let state = runtimes.get(id)
  if (!state) {
    state = emptyState()
    runtimes.set(id, state)
  }
  return state
}

export function getTunnelRow(id: number): TunnelRow | null {
  return (db.prepare('SELECT * FROM tunnels WHERE id = ?').get(id) as TunnelRow | undefined) ?? null
}

export function getStatus(id: number) {
  const state = getState(id)
  return { status: state.status, error: state.error, retryCount: state.retryCount }
}

function teardownNetwork(state: RuntimeState) {
  state.server?.close()
  state.client?.end()
  state.jumpConn?.end()
  state.server = null
  state.client = null
  state.jumpConn = null
}

function scheduleRetry(id: number, tunnel: TunnelRow, state: RuntimeState) {
  if (state.stopRequested) return
  if (state.retryCount >= tunnel.max_retries) {
    state.status = 'error'
    return
  }
  state.status = 'retrying'
  state.retryCount += 1
  state.retryTimer = setTimeout(() => startTunnel(id), tunnel.retry_interval_ms)
}

function bindLocalForward(client: Client, tunnel: TunnelRow, state: RuntimeState, onFail: (message: string) => void) {
  const server = net.createServer((socket) => {
    client.forwardOut(
      socket.remoteAddress ?? '127.0.0.1',
      socket.remotePort ?? 0,
      tunnel.endpoint_host,
      tunnel.endpoint_port,
      (err, stream) => {
        if (err) {
          socket.destroy()
          return
        }
        socket.pipe(stream).pipe(socket)
        stream.on('close', () => socket.destroy())
        socket.on('close', () => stream.end())
      },
    )
  })
  server.on('error', (err) => onFail(err.message))
  server.listen(tunnel.source_port, '127.0.0.1')
  state.server = server
}

function bindRemoteForward(client: Client, tunnel: TunnelRow, state: RuntimeState, onFail: (message: string) => void) {
  client.forwardIn('0.0.0.0', tunnel.source_port, (err) => {
    if (err) {
      onFail(err.message)
      return
    }
  })
  client.on('tcp connection', (info, accept, reject) => {
    if (info.destPort !== tunnel.source_port) {
      reject()
      return
    }
    const stream = accept()
    const sock = net.connect(tunnel.endpoint_port, tunnel.endpoint_host)
    sock.on('error', () => stream.end())
    stream.on('error', () => sock.destroy())
    sock.pipe(stream).pipe(sock)
  })
}

function bindDynamicForward(client: Client, tunnel: TunnelRow, state: RuntimeState, onFail: (message: string) => void) {
  const server = net.createServer((socket) => {
    readSocks5Target(socket)
      .then((target) => {
        client.forwardOut(socket.remoteAddress ?? '127.0.0.1', socket.remotePort ?? 0, target.host, target.port, (err, stream) => {
          if (err) {
            sendSocks5Failure(socket)
            socket.destroy()
            return
          }
          sendSocks5Success(socket)
          socket.pipe(stream).pipe(socket)
          stream.on('close', () => socket.destroy())
          socket.on('close', () => stream.end())
        })
      })
      .catch(() => socket.destroy())
  })
  server.on('error', (err) => onFail(err.message))
  server.listen(tunnel.source_port, '127.0.0.1')
  state.server = server
}

export function startTunnel(id: number) {
  const tunnel = getTunnelRow(id)
  if (!tunnel) return
  const state = getState(id)
  state.stopRequested = false
  state.status = 'connecting'
  state.error = null

  const target = loadSshHost(tunnel.integration_id)
  if (!target) {
    state.status = 'error'
    state.error = 'SSH integration not found'
    return
  }

  const onFail = (message: string) => {
    if (state.stopRequested) return
    state.error = message
    teardownNetwork(state)
    scheduleRetry(id, tunnel, state)
  }

  const result = connectTarget(
    target,
    (client) => {
      if (state.stopRequested) {
        client.end()
        return
      }
      state.client = client
      state.status = 'connected'
      state.error = null
      state.retryCount = 0
      client.on('error', (err) => onFail(err.message))
      client.on('close', () => onFail('SSH connection closed'))

      if (tunnel.mode === 'local') bindLocalForward(client, tunnel, state, onFail)
      else if (tunnel.mode === 'remote') bindRemoteForward(client, tunnel, state, onFail)
      else bindDynamicForward(client, tunnel, state, onFail)
    },
    onFail,
  )
  state.jumpConn = result.jumpConn
}

export function stopTunnel(id: number) {
  const state = getState(id)
  state.stopRequested = true
  if (state.retryTimer) clearTimeout(state.retryTimer)
  state.retryTimer = null
  state.retryCount = 0
  state.status = 'stopped'
  state.error = null
  teardownNetwork(state)
}

export function deleteTunnelRuntime(id: number) {
  stopTunnel(id)
  runtimes.delete(id)
}

export function startAutoStartTunnels() {
  const rows = db.prepare('SELECT * FROM tunnels WHERE auto_start = 1').all() as TunnelRow[]
  for (const row of rows) startTunnel(row.id)
}