better client management. Fix exporting and importing. Better key generation (not fucking dog shit).

2026-05-09 18:31:53 -05:00
parent 76154962f3
commit d62f6f2e34
8 changed files with 430 additions and 22 deletions
@@ -5,6 +5,7 @@ import { v4 as uuidv4 } from "uuid";
 import { saveBuild } from "../db";
 import { logger } from "../logger";
 import { getConfig } from "../config";
+import { signBuildToken } from "./build-signing";
 import { ensureDataDir } from "../paths";
 import * as buildManager from "../build/buildManager";
 import type { BuildStream } from "../build/types";
@@ -322,8 +323,13 @@ export async function startBuildProcess(
      sendToStream({ type: "output", text: "Mutex: disabled\n", level: "info" });
    }

-    const buildTag = uuidv4();
-    sendToStream({ type: "output", text: `Build tag: ${buildTag}\n`, level: "info" });
+    const buildTag = await signBuildToken({
+      v: 1,
+      bid: buildId,
+      uid: config.builtByUserId ?? null,
+      iat: Math.floor(Date.now() / 1000),
+    });
+    sendToStream({ type: "output", text: `Build tag: signed (bid=${buildId.substring(0, 8)})\n`, level: "info" });

    if (config.outputName) {
      sendToStream({ type: "output", text: `Custom output name: ${config.outputName}\n`, level: "info" });
@@ -0,0 +1,110 @@
+import { describe, expect, test, beforeAll, beforeEach } from "bun:test";
+import { mkdtemp } from "fs/promises";
+import { join } from "path";
+import { tmpdir } from "os";
+
+let signBuildToken: typeof import("./build-signing").signBuildToken;
+let verifyBuildToken: typeof import("./build-signing").verifyBuildToken;
+let isBuildBanned: typeof import("./build-signing").isBuildBanned;
+let addBuildToBanlist: typeof import("./build-signing").addBuildToBanlist;
+let removeBuildFromBanlist: typeof import("./build-signing").removeBuildFromBanlist;
+let resetCache: () => void;
+
+beforeAll(async () => {
+  process.env.DATA_DIR = await mkdtemp(join(tmpdir(), "overlord-build-signing-test-"));
+  const mod = await import("./build-signing");
+  signBuildToken = mod.signBuildToken;
+  verifyBuildToken = mod.verifyBuildToken;
+  isBuildBanned = mod.isBuildBanned;
+  addBuildToBanlist = mod.addBuildToBanlist;
+  removeBuildFromBanlist = mod.removeBuildFromBanlist;
+  resetCache = mod.resetBuildSigningCacheForTests;
+});
+
+describe("build-signing", () => {
+  test("sign/verify roundtrip preserves payload", async () => {
+    const payload = { v: 1 as const, bid: "build-abc-123", uid: 42, iat: 1700000000 };
+    const token = await signBuildToken(payload);
+    expect(typeof token).toBe("string");
+    expect(token).toContain(".");
+
+    const verified = await verifyBuildToken(token);
+    expect(verified).not.toBeNull();
+    expect(verified!.v).toBe(1);
+    expect(verified!.bid).toBe("build-abc-123");
+    expect(verified!.uid).toBe(42);
+    expect(verified!.iat).toBe(1700000000);
+  });
+
+  test("sign/verify works with null uid", async () => {
+    const payload = { v: 1 as const, bid: "build-no-owner", uid: null, iat: 1700000001 };
+    const token = await signBuildToken(payload);
+    const verified = await verifyBuildToken(token);
+    expect(verified).not.toBeNull();
+    expect(verified!.uid).toBe(null);
+  });
+
+  test("tampered payload fails verification", async () => {
+    const token = await signBuildToken({ v: 1, bid: "real-build", uid: 1, iat: 1700000000 });
+    const [payloadB64, sigB64] = token.split(".");
+    const tamperedPayload = Buffer.from(
+      JSON.stringify({ v: 1, bid: "evil-build", uid: 1, iat: 1700000000 }),
+    )
+      .toString("base64")
+      .replace(/\+/g, "-")
+      .replace(/\//g, "_")
+      .replace(/=+$/g, "");
+    expect(await verifyBuildToken(`${tamperedPayload}.${sigB64}`)).toBeNull();
+  });
+
+  test("tampered signature fails verification", async () => {
+    const token = await signBuildToken({ v: 1, bid: "real-build", uid: 1, iat: 1700000000 });
+    const [payloadB64] = token.split(".");
+    const fakeSig = "A".repeat(86);
+    expect(await verifyBuildToken(`${payloadB64}.${fakeSig}`)).toBeNull();
+  });
+
+  test("legacy UUID-style buildTag returns null (so DB fallback path triggers)", async () => {
+    expect(await verifyBuildToken("550e8400-e29b-41d4-a716-446655440000")).toBeNull();
+  });
+
+  test("garbage input returns null", async () => {
+    expect(await verifyBuildToken("")).toBeNull();
+    expect(await verifyBuildToken("not-a-token")).toBeNull();
+    expect(await verifyBuildToken(".")).toBeNull();
+    expect(await verifyBuildToken("a.")).toBeNull();
+    expect(await verifyBuildToken(".b")).toBeNull();
+  });
+
+  test("payload missing required fields returns null", async () => {
+    const badPayloadB64 = Buffer.from(JSON.stringify({ hello: "world" }))
+      .toString("base64")
+      .replace(/\+/g, "-")
+      .replace(/\//g, "_")
+      .replace(/=+$/g, "");
+    expect(await verifyBuildToken(`${badPayloadB64}.AAAA`)).toBeNull();
+  });
+
+  test("token signed with a different key fails verification", async () => {
+    const token = await signBuildToken({ v: 1, bid: "build-x", uid: 1, iat: 1700000000 });
+    const verifiedFirst = await verifyBuildToken(token);
+    expect(verifiedFirst).not.toBeNull();
+
+    process.env.DATA_DIR = await mkdtemp(join(tmpdir(), "overlord-build-signing-test-rotated-"));
+    resetCache();
+
+    expect(await verifyBuildToken(token)).toBeNull();
+  });
+
+  test("banlist add/remove/check", () => {
+    expect(isBuildBanned("nope")).toBe(false);
+    addBuildToBanlist("ban-me");
+    expect(isBuildBanned("ban-me")).toBe(true);
+    addBuildToBanlist("ban-me");
+    expect(isBuildBanned("ban-me")).toBe(true);
+    removeBuildFromBanlist("ban-me");
+    expect(isBuildBanned("ban-me")).toBe(false);
+    removeBuildFromBanlist("ban-me");
+    expect(isBuildBanned("ban-me")).toBe(false);
+  });
+});
@@ -0,0 +1,155 @@
+import logger from "../logger";
+import type { BuildTokenPayload } from "../types";
+import {
+  getBuildSigningSecrets,
+  setBuildSigningSecrets,
+  getBuildBanlist,
+  setBuildBanlist,
+} from "../config";
+
+let cachedPrivateKey: CryptoKey | null = null;
+let cachedPublicKey: CryptoKey | null = null;
+let cachedPublicKeyB64: string | null = null;
+
+function toBase64Url(bytes: Uint8Array): string {
+  return Buffer.from(bytes)
+    .toString("base64")
+    .replace(/\+/g, "-")
+    .replace(/\//g, "_")
+    .replace(/=+$/g, "");
+}
+
+function fromBase64Url(s: string): Uint8Array {
+  const pad = s.length % 4 === 0 ? "" : "=".repeat(4 - (s.length % 4));
+  const b64 = s.replace(/-/g, "+").replace(/_/g, "/") + pad;
+  return new Uint8Array(Buffer.from(b64, "base64"));
+}
+
+async function generateKeypair(): Promise<{ privateKeyB64: string; publicKeyB64: string }> {
+  const kp = (await crypto.subtle.generateKey("Ed25519", true, ["sign", "verify"])) as CryptoKeyPair;
+  const privRaw = await crypto.subtle.exportKey("pkcs8", kp.privateKey);
+  const pubRaw = await crypto.subtle.exportKey("raw", kp.publicKey);
+  return {
+    privateKeyB64: Buffer.from(privRaw).toString("base64"),
+    publicKeyB64: Buffer.from(pubRaw).toString("base64"),
+  };
+}
+
+async function importKeys(privateKeyB64: string, publicKeyB64: string) {
+  const privBytes = Buffer.from(privateKeyB64, "base64");
+  const pubBytes = Buffer.from(publicKeyB64, "base64");
+  cachedPrivateKey = await crypto.subtle.importKey(
+    "pkcs8",
+    privBytes,
+    { name: "Ed25519" },
+    false,
+    ["sign"],
+  );
+  cachedPublicKey = await crypto.subtle.importKey(
+    "raw",
+    pubBytes,
+    { name: "Ed25519" },
+    false,
+    ["verify"],
+  );
+  cachedPublicKeyB64 = publicKeyB64;
+}
+
+async function ensureKeysLoaded(): Promise<void> {
+  if (cachedPrivateKey && cachedPublicKey) return;
+
+  let secrets = getBuildSigningSecrets();
+  if (!secrets || !secrets.privateKey || !secrets.publicKey) {
+    const generated = await generateKeypair();
+    secrets = { privateKey: generated.privateKeyB64, publicKey: generated.publicKeyB64 };
+    setBuildSigningSecrets(secrets);
+    logger.info("[build-signing] generated new Ed25519 build-signing keypair");
+  }
+
+  await importKeys(secrets.privateKey, secrets.publicKey);
+}
+
+export function resetBuildSigningCacheForTests(): void {
+  cachedPrivateKey = null;
+  cachedPublicKey = null;
+  cachedPublicKeyB64 = null;
+}
+
+export async function getBuildSigningPublicKey(): Promise<string> {
+  await ensureKeysLoaded();
+  return cachedPublicKeyB64 || "";
+}
+
+export async function signBuildToken(payload: BuildTokenPayload): Promise<string> {
+  await ensureKeysLoaded();
+  const payloadJson = JSON.stringify(payload);
+  const payloadBytes = new TextEncoder().encode(payloadJson);
+  const payloadB64 = toBase64Url(payloadBytes);
+  const sigBuf = await crypto.subtle.sign("Ed25519", cachedPrivateKey!, payloadBytes as any);
+  const sigB64 = toBase64Url(new Uint8Array(sigBuf));
+  return `${payloadB64}.${sigB64}`;
+}
+
+export async function verifyBuildToken(token: string): Promise<BuildTokenPayload | null> {
+  if (!token || typeof token !== "string") return null;
+  const dot = token.indexOf(".");
+  if (dot <= 0 || dot === token.length - 1) return null;
+
+  try {
+    await ensureKeysLoaded();
+  } catch (err) {
+    logger.warn("[build-signing] failed to load signing keys", err);
+    return null;
+  }
+
+  try {
+    const payloadB64 = token.slice(0, dot);
+    const sigB64 = token.slice(dot + 1);
+    const payloadBytes = fromBase64Url(payloadB64);
+    const sigBytes = fromBase64Url(sigB64);
+    if (sigBytes.length !== 64) return null;
+
+    const ok = await crypto.subtle.verify(
+      "Ed25519",
+      cachedPublicKey!,
+      sigBytes as any,
+      payloadBytes as any,
+    );
+    if (!ok) return null;
+
+    const parsed = JSON.parse(new TextDecoder().decode(payloadBytes));
+    if (
+      !parsed ||
+      typeof parsed !== "object" ||
+      parsed.v !== 1 ||
+      typeof parsed.bid !== "string" ||
+      !parsed.bid ||
+      typeof parsed.iat !== "number" ||
+      (parsed.uid !== null && typeof parsed.uid !== "number")
+    ) {
+      return null;
+    }
+    return { v: 1, bid: parsed.bid, uid: parsed.uid ?? null, iat: parsed.iat };
+  } catch {
+    return null;
+  }
+}
+
+export function isBuildBanned(buildId: string): boolean {
+  if (!buildId) return false;
+  return getBuildBanlist().includes(buildId);
+}
+
+export function addBuildToBanlist(buildId: string): void {
+  if (!buildId) return;
+  const list = getBuildBanlist();
+  if (list.includes(buildId)) return;
+  setBuildBanlist([...list, buildId]);
+}
+
+export function removeBuildFromBanlist(buildId: string): void {
+  if (!buildId) return;
+  const list = getBuildBanlist();
+  const next = list.filter((id) => id !== buildId);
+  if (next.length !== list.length) setBuildBanlist(next);
+}
@@ -24,6 +24,7 @@ import { logger } from "../../logger";
 import { normalizeClientOs } from "../deploy-utils";
 import { canUserBuild, recordBuildStart, recordBuildEnd } from "../../build-rate-limit";
 import { getConfig } from "../../config";
+import { addBuildToBanlist, removeBuildFromBanlist } from "../build-signing";
 import path from "path";
 import fs from "fs";
 import { resolveRuntimeRoot } from "../runtime-paths";
@@ -552,6 +553,12 @@ export async function handleBuildRoutes(
        return Response.json({ error: "Build not found" }, { status: 404 });
      }

+      if (blocked) {
+        addBuildToBanlist(buildId);
+      } else {
+        removeBuildFromBanlist(buildId);
+      }
+
      let disconnected = 0;
      if (blocked && build.buildTag) {
        const clientIds = listClientIdsByBuildTag(build.buildTag);
@@ -19,6 +19,7 @@ import type { ClientInfo } from "../../types";
 import { clearClientSyncState, handleFrame, handleHello, handlePing, handlePong } from "../../wsHandlers";
 import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation";
 import { stopAllProxiesForClient } from "../socks5-proxy-manager";
+import { verifyBuildToken, isBuildBanned } from "../build-signing";

 const OFFLINE_GRACE_MS = (() => {
  const raw = process.env.OVERLORD_OFFLINE_GRACE_MS;
@@ -164,7 +165,12 @@ type WsLifecycleDeps = {
  clearPendingNotificationScreenshots: (clientId: string) => void;
  clearClientPluginState: (clientId: string) => void;
  notifyRemoteDesktopStatus: (clientId: string, status: string, reason?: string) => void;
-  handleBuildTagConnection: (clientId: string, buildTag: string, keyFingerprint: string) => void;
+  handleBuildTagConnection: (
+    clientId: string,
+    buildId: string | null,
+    builtByUserId: number | undefined,
+    keyFingerprint: string,
+  ) => void;
  notifyDashboard: () => void;
  notifyDashboardClientEvent: (
    event: "client_online" | "client_offline" | "client_purgatory",
@@ -360,11 +366,30 @@ export async function handleWebSocketMessage(
        const buildTag = typeof (payload as any).buildTag === "string"
          ? (payload as any).buildTag.trim()
          : "";
-        const build = buildTag ? getBuildByTag(buildTag) : null;
-        const builtByUserId = build?.builtByUserId;

-        if (buildTag && build?.blocked) {
-          logger.info(`[build-block] rejecting agent ${resolvedId} from blocked build ${build.id.substring(0, 8)}`);
+        let resolvedBuildId: string | null = null;
+        let builtByUserId: number | undefined;
+        let isRevoked = false;
+
+        if (buildTag) {
+          const verified = await verifyBuildToken(buildTag);
+          if (verified) {
+            resolvedBuildId = verified.bid;
+            builtByUserId = verified.uid ?? undefined;
+            isRevoked = isBuildBanned(verified.bid);
+          } else {
+            const build = getBuildByTag(buildTag);
+            if (build) {
+              resolvedBuildId = build.id;
+              builtByUserId = build.builtByUserId;
+              isRevoked = !!build.blocked || isBuildBanned(build.id);
+            }
+          }
+        }
+
+        if (buildTag && isRevoked) {
+          const shortId = resolvedBuildId ? resolvedBuildId.substring(0, 8) : "<unknown>";
+          logger.info(`[build-block] rejecting agent ${resolvedId} from revoked build ${shortId}`);
          try {
            ws.send(
              encodeMessage({
@@ -559,7 +584,7 @@ export async function handleWebSocketMessage(
          (ws as any).data.wasKnown = true;

          if (buildTag) {
-            deps.handleBuildTagConnection(infoObj.id, buildTag, keyFingerprint);
+            deps.handleBuildTagConnection(infoObj.id, resolvedBuildId, builtByUserId, keyFingerprint);
          }
        }
        break;