feat: route code queue egress through provider gateway

src/components/backend-core/src/index.ts

@@ -2,6 +2,7 @@ import type { Server, ServerWebSocket } from "bun";
 import { createHash } from "node:crypto";
 import { createReadStream } from "node:fs";
 import { mkdir, rm, stat } from "node:fs/promises";
+import { connect as connectTcp, type Socket } from "node:net";
 import { basename, dirname, relative, resolve, posix as pathPosix } from "node:path";
 import postgres from "postgres";
 import { createHourlyJsonlWriter, logRetentionBytesForService } from "../../shared/src/rotating-jsonl";
@@ -17,7 +18,13 @@ import {
   type CoreHostSshOpenMessage,
   type CoreHostSshResizeMessage,
   type CoreDispatchMessage,
+  type CoreEgressTcpCloseMessage,
+  type CoreEgressTcpDataMessage,
+  type CoreEgressTcpOpenedMessage,
   type JsonValue,
+  type ProviderEgressTcpCloseMessage,
+  type ProviderEgressTcpDataMessage,
+  type ProviderEgressTcpOpenMessage,
   type ProviderHostSshDataMessage,
   type ProviderHostSshErrorMessage,
   type ProviderHostSshExitMessage,
@@ -178,6 +185,13 @@ type ScheduleAction = DispatchScheduleAction | PgdataBackupScheduleAction;
 
 type TaskTerminalWaiter = (task: RawTaskRow | null) => void;
 
+interface EgressTcpConnection {
+  providerId: string;
+  connectionId: string;
+  socket: Socket;
+  provider: ProviderSocket;
+}
+
 interface MicroserviceProxyCacheEntry {
   expiresAt: number;
   staleExpiresAt: number;
@@ -201,6 +215,7 @@ interface MicroserviceAvailabilityEntry {
 const recentLogs: unknown[] = [];
 const activeProviders = new Map<string, ProviderSocket>();
 const activeSshClients = new Map<string, ProviderSocket>();
+const activeEgressTcpConnections = new Map<string, EgressTcpConnection>();
 const serviceStartedAt = new Date();
 const config = readConfig();
 const logger = createLogger("backend-core", config.logFile);
@@ -977,6 +992,89 @@ function wsSendJson(ws: ProviderSocket, value: unknown): void {
   if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(value));
 }
 
+function egressTcpKey(providerId: string, connectionId: string): string {
+  return `${providerId}:${connectionId}`;
+}
+
+function isValidEgressHost(host: string): boolean {
+  return host.length > 0 && host.length <= 253 && !/[\s/\\:@\0]/u.test(host);
+}
+
+function isValidEgressPort(port: number): boolean {
+  return Number.isInteger(port) && port > 0 && port <= 65_535;
+}
+
+function sendEgressClose(provider: ProviderSocket, connectionId: string, error?: string): void {
+  const message: CoreEgressTcpCloseMessage = error === undefined
+    ? { type: "egress_tcp_close", connectionId }
+    : { type: "egress_tcp_close", connectionId, error };
+  wsSendJson(provider, message);
+}
+
+function closeEgressTcpConnection(providerId: string, connectionId: string, error?: string): void {
+  const key = egressTcpKey(providerId, connectionId);
+  const connection = activeEgressTcpConnections.get(key);
+  if (connection === undefined) return;
+  activeEgressTcpConnections.delete(key);
+  connection.socket.destroy();
+  if (error !== undefined) sendEgressClose(connection.provider, connectionId, error);
+}
+
+function handleEgressTcpOpen(ws: ProviderSocket, message: ProviderEgressTcpOpenMessage): void {
+  const host = message.host.trim();
+  const port = Number(message.port);
+  if (!isValidEgressHost(host) || !isValidEgressPort(port)) {
+    sendEgressClose(ws, message.connectionId, "invalid egress target");
+    return;
+  }
+  const key = egressTcpKey(message.providerId, message.connectionId);
+  closeEgressTcpConnection(message.providerId, message.connectionId);
+  const socket = connectTcp({ host, port });
+  const connection: EgressTcpConnection = { providerId: message.providerId, connectionId: message.connectionId, socket, provider: ws };
+  activeEgressTcpConnections.set(key, connection);
+  socket.on("connect", () => {
+    const opened: CoreEgressTcpOpenedMessage = { type: "egress_tcp_opened", connectionId: message.connectionId };
+    wsSendJson(ws, opened);
+  });
+  socket.on("data", (chunk) => {
+    const data: CoreEgressTcpDataMessage = {
+      type: "egress_tcp_data",
+      connectionId: message.connectionId,
+      data: chunk.toString("base64"),
+      encoding: "base64",
+    };
+    wsSendJson(ws, data);
+  });
+  socket.on("close", () => {
+    if (activeEgressTcpConnections.get(key) !== connection) return;
+    activeEgressTcpConnections.delete(key);
+    sendEgressClose(ws, message.connectionId);
+  });
+  socket.on("error", (error) => {
+    if (activeEgressTcpConnections.get(key) !== connection) return;
+    activeEgressTcpConnections.delete(key);
+    sendEgressClose(ws, message.connectionId, error.message);
+  });
+}
+
+function handleEgressTcpData(message: ProviderEgressTcpDataMessage): void {
+  const connection = activeEgressTcpConnections.get(egressTcpKey(message.providerId, message.connectionId));
+  if (connection === undefined) return;
+  connection.socket.write(Buffer.from(message.data, message.encoding === "base64" ? "base64" : "utf8"));
+}
+
+function handleEgressTcpClose(message: ProviderEgressTcpCloseMessage): void {
+  closeEgressTcpConnection(message.providerId, message.connectionId);
+}
+
+function closeEgressTcpConnectionsForProvider(providerId: string): void {
+  for (const [key, connection] of activeEgressTcpConnections) {
+    if (connection.providerId !== providerId) continue;
+    activeEgressTcpConnections.delete(key);
+    connection.socket.destroy();
+  }
+}
+
 function sshClientFor(sessionId: string): ProviderSocket | null {
   return activeSshClients.get(sessionId) ?? null;
 }
@@ -1031,6 +1129,19 @@ async function handleProviderMessage(ws: ProviderSocket, raw: string | Buffer):
     return;
   }
 
+  if (message.type === "egress_tcp_open") {
+    handleEgressTcpOpen(ws, message);
+    return;
+  }
+  if (message.type === "egress_tcp_data") {
+    handleEgressTcpData(message);
+    return;
+  }
+  if (message.type === "egress_tcp_close") {
+    handleEgressTcpClose(message);
+    return;
+  }
+
   if (message.type === "register") {
     const labels = { ...message.labels, unideskCapabilities: message.capabilities };
     await upsertNodeOnline(message.providerId, message.name, labels);
@@ -3499,6 +3610,7 @@ const providerServer = Bun.serve<WsData>({
       const providerId = ws.data.providerId;
       logger("warn", "provider_socket_close", { providerId: providerId ?? null });
       if (providerId !== undefined) {
+        closeEgressTcpConnectionsForProvider(providerId);
         if (activeProviders.get(providerId) !== ws) {
           logger("info", "provider_socket_close_ignored_replaced", { providerId });
           return;