ContainerFlow/src/server/index.ts

import { Hono } from "hono";
import { serveStatic } from "hono/bun";
import { cors } from "hono/cors";
import { compress } from "hono/compress";
import path from "path";
import fs from "fs";
import { docker, discoverServices, discoverConnections, getContainerLogs, streamContainerLogs } from "./docker";
import { pollStats, watchDockerEvents } from "./watcher";
import type { Service, WSMessage } from "../shared/types";

/** Env-file overrides per compose file (persisted to file) */
const ENV_FILES_FILE = path.join(process.cwd(), ".dockerflow-env-files.json");

function loadEnvFiles(): Record<string, string> {
  try {
    if (fs.existsSync(ENV_FILES_FILE)) {
      return JSON.parse(fs.readFileSync(ENV_FILES_FILE, "utf-8"));
    }
  } catch {}
  return {};
}

/** Build env-file args for docker compose by detecting .env files next to the compose file */
function findEnvFileArgs(composeFile: string): string[] {
  // Check for user override first
  const overrides = loadEnvFiles();
  const override = overrides[composeFile];
  if (override) {
    const resolved = path.isAbsolute(override) ? override : path.join(path.dirname(composeFile), override);
    if (fs.existsSync(resolved)) {
      return ["--env-file", resolved];
    }
  }

  // Auto-detect heuristic
  const dir = path.dirname(composeFile);
  const baseName = path.basename(composeFile, path.extname(composeFile)); // e.g. "docker-compose.prod"
  const candidates: string[] = [];

  // Prefer specific env file matching compose name (e.g. .env.prod for docker-compose.prod.yml)
  const suffix = baseName.replace(/^docker-compose\.?/, ""); // "prod" or ""
  if (suffix) {
    candidates.push(path.join(dir, `.env.${suffix}`));
  }
  candidates.push(path.join(dir, ".env"));

  for (const envFile of candidates) {
    if (fs.existsSync(envFile)) {
      return ["--env-file", envFile];
    }
  }
  return [];
}

const app = new Hono();

// ── Compression ──
app.use("*", compress());

// ── CORS ──
app.use("/api/*", cors());

// ── CLI args ──
const args = process.argv.slice(2);
const ALL = args.includes("--all");
const projectsFlag = args.find((a) => a.startsWith("--projects="));
const PROJECTS = projectsFlag
  ? projectsFlag.split("=")[1]!.split(",")
  : ALL
    ? []
    : [path.basename(process.cwd())];

// ── Config ──
const PORT = parseInt(process.env.PORT || "9470");
const AUTH_TOKEN = process.env.AUTH_TOKEN || "";
const HOST = AUTH_TOKEN ? "0.0.0.0" : "127.0.0.1";
const POLL_INTERVAL_MS = 5000;
const WS_RECONNECT_MS = 3000;

// ── Rate limiting (in-memory, per IP) ──
const RATE_LIMIT_MAX = 5;
const RATE_LIMIT_WINDOW_MS = 60_000;
const failedAttempts = new Map<string, { count: number; resetAt: number }>();

function isRateLimited(ip: string): boolean {
  const entry = failedAttempts.get(ip);
  if (!entry) return false;
  if (Date.now() > entry.resetAt) {
    failedAttempts.delete(ip);
    return false;
  }
  return entry.count >= RATE_LIMIT_MAX;
}

function recordFailedAttempt(ip: string): void {
  const now = Date.now();
  const entry = failedAttempts.get(ip);
  if (!entry || now > entry.resetAt) {
    failedAttempts.set(ip, { count: 1, resetAt: now + RATE_LIMIT_WINDOW_MS });
  } else {
    entry.count++;
  }
}

// Cleanup stale entries every 5 minutes
setInterval(() => {
  const now = Date.now();
  for (const [ip, entry] of failedAttempts) {
    if (now > entry.resetAt) failedAttempts.delete(ip);
  }
}, 5 * 60_000);

// ── Auth middleware ──
if (AUTH_TOKEN) {
  app.use("*", async (c, next) => {
    // Skip static assets and auth page
    if (c.req.path === "/" || c.req.path.startsWith("/assets") || c.req.path.endsWith(".png") || c.req.path.endsWith(".webp") || c.req.path.endsWith(".ico")) return next();
    if (c.req.path === "/api/auth") return next();

    const ip = c.req.header("x-forwarded-for")?.split(",")[0]?.trim() || "unknown";
    if (isRateLimited(ip)) return c.json({ error: "Too many failed attempts. Try again later." }, 429);

    const token = c.req.header("Authorization")?.replace("Bearer ", "");
    if (token !== AUTH_TOKEN) {
      recordFailedAttempt(ip);
      return c.json({ error: "Unauthorized" }, 401);
    }
    return next();
  });
}

// ── API ──
app.get("/api/services", async (c) => {
  const services = await discoverServices(ALL, PROJECTS);
  return c.json(services);
});

app.get("/api/connections", async (c) => {
  const services = await discoverServices(ALL, PROJECTS);
  const connections = await discoverConnections(services);
  return c.json(connections);
});

app.get("/api/health", (c) => c.json({ ok: true, mode: ALL ? "all" : "filtered", projects: PROJECTS }));

// ── Combined init endpoint (services + connections + positions in one call) ──
app.get("/api/init", async (c) => {
  const services = await discoverServices(ALL, PROJECTS);
  const connections = await discoverConnections(services);
  let positions: Record<string, any> = {};
  try {
    if (fs.existsSync(POSITIONS_FILE)) {
      positions = JSON.parse(fs.readFileSync(POSITIONS_FILE, "utf-8"));
    }
  } catch {}
  return c.json({ services, connections, positions });
});

// ── Container actions ──
app.post("/api/containers/:id/stop", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const container = docker.getContainer(id);
    await container.stop();
    immediateRefresh();
    return c.json({ ok: true });
  } catch (err: any) {
    if (err?.statusCode === 304) return c.json({ ok: true, message: "Already stopped" });
    return c.json({ error: err?.message || "Failed to stop container" }, 500);
  }
});

app.post("/api/containers/:id/start", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const container = docker.getContainer(id);
    await container.start();
    immediateRefresh();
    return c.json({ ok: true });
  } catch (err: any) {
    if (err?.statusCode === 304) return c.json({ ok: true, message: "Already running" });
    return c.json({ error: err?.message || "Failed to start container" }, 500);
  }
});

app.post("/api/containers/:id/restart", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const container = docker.getContainer(id);
    await container.restart();
    immediateRefresh();
    return c.json({ ok: true });
  } catch (err: any) {
    return c.json({ error: err?.message || "Failed to restart container" }, 500);
  }
});

app.post("/api/containers/:id/rebuild", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const container = docker.getContainer(id);
    const info = await container.inspect();
    const composeFile = info.Config?.Labels?.["com.docker.compose.project.config_files"];
    const serviceName = info.Config?.Labels?.["com.docker.compose.service"];
    const project = info.Config?.Labels?.["com.docker.compose.project"] || "standalone";
    if (!composeFile || !serviceName) {
      return c.json({ error: "Not a Compose service — rebuild requires docker-compose" }, 400);
    }
    const uid = `${project}/${serviceName}`;
    const envArgs = findEnvFileArgs(composeFile);
    // Run rebuild in background — respond immediately
    const proc = Bun.spawn(["docker", "compose", "-f", composeFile, ...envArgs, "up", "--build", "-d", serviceName], {
      stdout: "pipe",
      stderr: "pipe",
    });
    proc.exited.then(async (exitCode) => {
      if (exitCode !== 0) {
        const stderr = await new Response(proc.stderr).text();
        broadcast({ type: "action_error", data: { uid, action: "rebuild", error: stderr || `Rebuild failed with exit code ${exitCode}` } });
      }
      scheduleRefresh();
    }).catch((err) => {
      broadcast({ type: "action_error", data: { uid, action: "rebuild", error: err?.message || "Rebuild failed" } });
    });
    return c.json({ ok: true });
  } catch (err: any) {
    return c.json({ error: err?.message || "Failed to rebuild container" }, 500);
  }
});

app.post("/api/containers/:id/remove", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const container = docker.getContainer(id);
    const info = await container.inspect();
    const composeFile = info.Config?.Labels?.["com.docker.compose.project.config_files"];
    const serviceName = info.Config?.Labels?.["com.docker.compose.service"];
    if (!composeFile || !serviceName) {
      // Not a compose service — just stop and remove the container
      try { await container.stop(); } catch {}
      await container.remove({ force: true });
      return c.json({ ok: true });
    }
    const envArgs = findEnvFileArgs(composeFile);
    const proc = Bun.spawn(["docker", "compose", "-f", composeFile, ...envArgs, "rm", "-sf", serviceName], {
      stdout: "pipe",
      stderr: "pipe",
    });
    const exitCode = await proc.exited;
    if (exitCode !== 0) {
      const stderr = await new Response(proc.stderr).text();
      return c.json({ error: stderr || `Remove failed with exit code ${exitCode}` }, 500);
    }
    return c.json({ ok: true });
  } catch (err: any) {
    return c.json({ error: err?.message || "Failed to remove container" }, 500);
  }
});

app.post("/api/containers/:id/exec", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const body = await c.req.json();
    const cmd = body?.cmd;
    if (!cmd || typeof cmd !== "string") return c.json({ error: "Missing cmd" }, 400);

    // Parse command respecting quotes
    const parts: string[] = [];
    let current = "";
    let inQuote: string | null = null;
    for (const ch of cmd) {
      if (inQuote) {
        if (ch === inQuote) { inQuote = null; }
        else { current += ch; }
      } else if (ch === '"' || ch === "'") {
        inQuote = ch;
      } else if (ch === " ") {
        if (current) { parts.push(current); current = ""; }
      } else {
        current += ch;
      }
    }
    if (current) parts.push(current);
    if (parts.length === 0) return c.json({ error: "Empty command" }, 400);

    const container = docker.getContainer(id);
    const exec = await container.exec({ Cmd: parts, AttachStdout: true, AttachStderr: true });
    const stream = await exec.start({});

    // Collect output using dockerode's demuxStream
    const stdoutChunks: Buffer[] = [];
    const stderrChunks: Buffer[] = [];
    await new Promise<void>((resolve) => {
      const passStdout = new (require("stream").PassThrough)();
      const passStderr = new (require("stream").PassThrough)();
      passStdout.on("data", (chunk: Buffer) => stdoutChunks.push(chunk));
      passStderr.on("data", (chunk: Buffer) => stderrChunks.push(chunk));
      docker.modem.demuxStream(stream, passStdout, passStderr);
      stream.on("end", resolve);
      stream.on("error", resolve);
      setTimeout(resolve, 30000);
    });

    const stdout = Buffer.concat(stdoutChunks).toString("utf-8");
    const stderr = Buffer.concat(stderrChunks).toString("utf-8");
    const output = (stdout + stderr).trim();

    const inspect = await exec.inspect();
    return c.json({ ok: true, output, exitCode: inspect.ExitCode ?? -1 });
  } catch (err: any) {
    return c.json({ error: err?.message || "Failed to exec" }, 500);
  }
});

app.get("/api/logs/:id", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) {
    return c.json({ error: "Invalid container ID" }, 400);
  }
  const tail = Math.min(Math.max(parseInt(c.req.query("tail") || "200") || 200, 1), 5000);
  const since = c.req.query("since") ? parseInt(c.req.query("since")!) : undefined;
  try {
    const lines = await getContainerLogs(id, tail, since);
    return c.json(lines);
  } catch (err) {
    return c.json({ error: "Failed to fetch logs" }, 500);
  }
});

// ── Node positions (persisted to file) ──
const POSITIONS_FILE = path.join(process.cwd(), ".dockerflow-positions.json");

app.get("/api/positions", (c) => {
  try {
    if (fs.existsSync(POSITIONS_FILE)) {
      const data = JSON.parse(fs.readFileSync(POSITIONS_FILE, "utf-8"));
      return c.json(data);
    }
  } catch (err) {
    console.error("Failed to read positions file:", err);
  }
  return c.json({});
});

app.put("/api/positions", async (c) => {
  try {
    const body = await c.req.json();
    fs.writeFileSync(POSITIONS_FILE, JSON.stringify(body, null, 2));
    return c.json({ ok: true });
  } catch {
    return c.json({ error: "Failed to save" }, 500);
  }
});

// ── Env-file overrides (persisted to file) ──
app.get("/api/env-files", (c) => {
  return c.json(loadEnvFiles());
});

app.put("/api/env-files", async (c) => {
  try {
    const body = await c.req.json();
    fs.writeFileSync(ENV_FILES_FILE, JSON.stringify(body, null, 2));
    return c.json({ ok: true });
  } catch {
    return c.json({ error: "Failed to save" }, 500);
  }
});

app.get("/api/env-files/detect/:id", async (c) => {
  const id = c.req.param("id");
  if (!/^[a-f0-9]{12,64}$/.test(id)) return c.json({ error: "Invalid container ID" }, 400);
  try {
    const container = docker.getContainer(id);
    const info = await container.inspect();
    const composeFile = info.Config?.Labels?.["com.docker.compose.project.config_files"];
    if (!composeFile) return c.json({ files: [], composeFile: null });
    const dir = path.dirname(composeFile);
    const entries = fs.readdirSync(dir);
    const envFiles = entries.filter((e: string) => e.startsWith(".env") && !e.endsWith(".example") && !e.endsWith(".sample"))
      .map((e: string) => e)
      .sort();
    return c.json({ files: envFiles, composeFile });
  } catch (err: any) {
    return c.json({ error: err?.message || "Failed to detect env files" }, 500);
  }
});

// ── Cache headers for static assets ──
app.use("/*", async (c, next) => {
  await next();
  const p = c.req.path;
  if (p.startsWith("/assets/")) {
    // Hashed filenames — cache forever
    c.header("Cache-Control", "public, max-age=31536000, immutable");
  } else if (p.endsWith(".webp") || p.endsWith(".png") || p.endsWith(".ico")) {
    c.header("Cache-Control", "public, max-age=86400");
  }
});

// ── Serve frontend build ──
app.use("/*", serveStatic({ root: "./dist" }));
app.get("/*", serveStatic({ root: "./dist", path: "index.html" }));

// ── WebSocket ──
const clients = new Set<WebSocket>();
const authenticatedClients = new Set<WebSocket>();
const logStreams = new Map<WebSocket, { destroy: () => void }>();

function isAuthenticated(ws: WebSocket): boolean {
  return !AUTH_TOKEN || authenticatedClients.has(ws);
}

function broadcast(msg: WSMessage) {
  const data = JSON.stringify(msg);
  for (const ws of clients) {
    if (!isAuthenticated(ws)) continue;
    try {
      ws.send(data);
    } catch {}
  }
}

function cleanupLogStream(ws: WebSocket) {
  const stream = logStreams.get(ws);
  if (stream) {
    logStreams.delete(ws);
    try {
      stream.destroy();
    } catch (err) {
      console.error("Failed to destroy log stream:", err);
    }
  }
}

// ── Docker events ──
let servicesLock = false;
let statsLock = false;

async function refreshServices() {
  if (servicesLock) return;
  servicesLock = true;
  try {
    const services = await discoverServices(ALL, PROJECTS);

    const svcHash = services.map((s) => `${s.uid}:${s.state}`).join("|");
    if (svcHash !== lastServicesHash) {
      lastServicesHash = svcHash;
      broadcast({ type: "services", data: services });
    }

    const connections = await discoverConnections(services);
    const connHash = connections.map((c) => `${c.from}:${c.to}`).join("|");
    if (connHash !== lastConnectionsHash) {
      lastConnectionsHash = connHash;
      broadcast({ type: "connections", data: connections });
    }

    // Stats polling is separate — don't block services refresh
    refreshStats(services);
  } catch (err) {
    console.error("Refresh error:", err);
  } finally {
    servicesLock = false;
  }
}

let statsLockTimer: ReturnType<typeof setTimeout> | undefined;
async function refreshStats(services: Service[]) {
  if (statsLock) return;
  statsLock = true;
  // Safety: force-unlock after 30s in case pollStats hangs
  clearTimeout(statsLockTimer);
  statsLockTimer = setTimeout(() => { statsLock = false; }, 30000);
  try {
    const stats = await pollStats(services);
    broadcast({ type: "stats", data: stats });
  } catch (err) {
    console.error("Stats error:", err);
  } finally {
    clearTimeout(statsLockTimer);
    statsLock = false;
  }
}

// Debounced refresh for Docker events
let refreshTimer: ReturnType<typeof setTimeout> | undefined;
let retryTimer: ReturnType<typeof setTimeout> | undefined;
function scheduleRefresh() {
  // Invalidate hash so next refresh always broadcasts (restart: same final state but clients need the update)
  lastServicesHash = "";
  clearTimeout(refreshTimer);
  clearTimeout(retryTimer);
  refreshTimer = setTimeout(() => {
    refreshServices();
    retryTimer = setTimeout(refreshServices, 1500);
  }, 500);
}

// Immediate refresh after action endpoints (container already changed state)
function immediateRefresh() {
  lastServicesHash = "";
  clearTimeout(refreshTimer);
  clearTimeout(retryTimer);
  refreshServices();
}

watchDockerEvents((event) => {
  broadcast({ type: "docker_event", data: event });
  scheduleRefresh();
});

// ── Stats polling ──
let lastServicesHash = "";
let lastConnectionsHash = "";

setInterval(refreshServices, POLL_INTERVAL_MS);

// ── Start ──
const server = Bun.serve({
  hostname: HOST,
  port: PORT,
  fetch(req, server) {
    const url = new URL(req.url);

    // WebSocket upgrade (auth handled via first message)
    if (url.pathname === "/ws") {
      if (server.upgrade(req)) return undefined;
      return new Response("WebSocket upgrade failed", { status: 400 });
    }

    return app.fetch(req, server);
  },
  websocket: {
    open(ws) {
      const native = ws as unknown as WebSocket;
      clients.add(native);

      if (!AUTH_TOKEN) {
        // No auth required — send data immediately
        discoverServices(ALL, PROJECTS).then(async (services) => {
          const connections = await discoverConnections(services);
          const stats = await pollStats(services);
          try {
            native.send(JSON.stringify({ type: "services", data: services }));
            native.send(JSON.stringify({ type: "connections", data: connections }));
            native.send(JSON.stringify({ type: "stats", data: stats }));
          } catch {}
        }).catch(() => {});
      }
    },
    close(ws) {
      const native = ws as unknown as WebSocket;
      cleanupLogStream(native);
      clients.delete(native);
      authenticatedClients.delete(native);
    },
    message(ws, message) {
      try {
        const msg = JSON.parse(typeof message === "string" ? message : new TextDecoder().decode(message as unknown as ArrayBuffer));
        const native = ws as unknown as WebSocket;

        // Handle authentication via first message
        if (msg.type === "auth") {
          const wsIp = (ws as any).remoteAddress || "unknown";
          if (AUTH_TOKEN && isRateLimited(wsIp)) {
            native.send(JSON.stringify({ type: "auth_error", reason: "rate_limited" }));
            native.close();
            return;
          }
          if (msg.token === AUTH_TOKEN) {
            authenticatedClients.add(native);
            native.send(JSON.stringify({ type: "auth_ok" }));
            // Send current services/connections/stats immediately
            discoverServices(ALL, PROJECTS).then(async (services) => {
              const connections = await discoverConnections(services);
              const stats = await pollStats(services);
              try {
                native.send(JSON.stringify({ type: "services", data: services }));
                native.send(JSON.stringify({ type: "connections", data: connections }));
                native.send(JSON.stringify({ type: "stats", data: stats }));
              } catch {}
            }).catch(() => {});
          } else {
            if (AUTH_TOKEN) recordFailedAttempt(wsIp);
            native.send(JSON.stringify({ type: "auth_error" }));
            native.close();
          }
          return;
        }

        // Reject messages from unauthenticated clients
        if (!isAuthenticated(native)) return;

        if (msg.type === "subscribe_logs" && msg.container) {
          cleanupLogStream(native);

          const stream = streamContainerLogs(msg.container, (line) => {
            try {
              native.send(JSON.stringify({ type: "log_line", data: line }));
            } catch {}
          });
          logStreams.set(native, stream);
        } else if (msg.type === "unsubscribe_logs") {
          cleanupLogStream(native);
        }
      } catch (err) {
        console.error("Failed to handle WS message:", err);
      }
    },
  },
});

const mode = ALL ? "all projects" : `project(s): ${PROJECTS.join(", ")}`;
console.log(`\n  ContainerFlow`);
console.log(`  → http://${HOST}:${PORT}`);
console.log(`  → Mode: ${mode}`);
console.log(`  → Auth: ${AUTH_TOKEN ? "enabled" : "disabled (localhost only)"}\n`);