Arcrun/mcp/tests/unit/partner-auth.test.ts

import { describe, it, expect } from "vitest";
import { Hono } from "hono";
import { partnerAuthMiddleware } from "../../src/middleware/partner-auth.js";
import { putAccessToken } from "../../src/oauth/store.js";
import type { Env } from "../../src/types.js";

// ── 記憶體 KV mock ────────────────────────────────────────────────────────────
function makeKV(): KVNamespace {
  const map = new Map<string, string>();
  return {
    async put(k: string, v: string) {
      map.set(k, v);
    },
    async get(k: string) {
      return map.get(k) ?? null;
    },
    async delete(k: string) {
      map.delete(k);
    },
  } as unknown as KVNamespace;
}

// 假 KBDB Fetcher：依 URL 回傳 partner 驗證結果。
function makeKBDB(valid: boolean, ns = "org-a"): Fetcher {
  return {
    async fetch() {
      return new Response(JSON.stringify({ valid, org_namespace: ns }), {
        status: valid ? 200 : 401,
      });
    },
  } as unknown as Fetcher;
}

function baseEnv(over: Partial<Env> = {}): Env {
  return {
    COMPONENT_REGISTRY: {} as Fetcher,
    CYPHER_EXECUTOR: {} as Fetcher,
    KBDB: makeKBDB(true),
    KBDB_INTERNAL_TOKEN: "internal",
    ...over,
  } as Env;
}

// 建一個掛了 partnerAuthMiddleware 的最小 app，protected handler 回傳解出的 namespace。
function buildApp(env: Env) {
  const app = new Hono<{ Bindings: Env; Variables: { org_namespace: string; partner_token: string } }>();
  app.get("/mcp", partnerAuthMiddleware, (c) =>
    c.json({ org_namespace: c.get("org_namespace"), partner_token: c.get("partner_token") }),
  );
  return (auth?: string) =>
    app.request("https://mcp.arcrun.dev/mcp", auth ? { headers: { Authorization: auth } } : {}, env);
}

describe("partner-auth: 無 / 壞 Authorization → 401 + WWW-Authenticate（RFC 9728）", () => {
  it("缺 header → 401 帶 WWW-Authenticate resource_metadata", async () => {
    const r = await buildApp(baseEnv({ MULTI_TENANT: "false" }))();
    expect(r.status).toBe(401);
    expect(r.headers.get("WWW-Authenticate")).toBe(
      'Bearer resource_metadata="https://mcp.arcrun.dev/.well-known/oauth-protected-resource"',
    );
  });

  it("非 Bearer → 401", async () => {
    const r = await buildApp(baseEnv())("Basic abc");
    expect(r.status).toBe(401);
    expect(r.headers.get("WWW-Authenticate")).toContain("resource_metadata=");
  });
});

describe("partner-auth: OAuth access token 路徑（遠端 claude.ai）", () => {
  it("有效 access_token → 解出綁定 namespace", async () => {
    const kv = makeKV();
    await putAccessToken(
      kv,
      "good-token",
      { namespace: "leo", client_id: "c1", scope: "mcp", aud: "x", exp: Math.floor(Date.now() / 1000) + 100 },
      100,
    );
    const r = await buildApp(baseEnv({ MULTI_TENANT: "false", OAUTH_KV: kv }))("Bearer good-token");
    expect(r.status).toBe(200);
    expect((await r.json()).org_namespace).toBe("leo");
  });

  it("未知 token（非 OAuth、非 static）在 self-hosted → 401（明碼 namespace 不再放行）", async () => {
    const kv = makeKV();
    const r = await buildApp(baseEnv({ MULTI_TENANT: "false", OAUTH_KV: kv }))("Bearer leo");
    expect(r.status).toBe(401);
    expect(r.headers.get("WWW-Authenticate")).toContain("resource_metadata=");
  });
});

describe("partner-auth: MCP_STATIC_TOKEN 相容路徑（本機 CLI/GUI 真祕密）", () => {
  it("Bearer == static token → owner namespace", async () => {
    const r = await buildApp(
      baseEnv({ MULTI_TENANT: "false", MCP_STATIC_TOKEN: "real-secret-xyz", MCP_OWNER_NAMESPACE: "leo" }),
    )("Bearer real-secret-xyz");
    expect(r.status).toBe(200);
    expect((await r.json()).org_namespace).toBe("leo");
  });

  it("Bearer != static token（明碼 namespace）→ 401", async () => {
    const r = await buildApp(
      baseEnv({ MULTI_TENANT: "false", MCP_STATIC_TOKEN: "real-secret-xyz" }),
    )("Bearer leo");
    expect(r.status).toBe(401);
  });
});

describe("partner-auth: 官方 SaaS partner-key（行為不變）", () => {
  it("合法 partner key → org_namespace", async () => {
    const r = await buildApp(baseEnv({ KBDB: makeKBDB(true, "org-a") }))("Bearer pk_live_x");
    expect(r.status).toBe(200);
    expect((await r.json()).org_namespace).toBe("org-a");
  });

  it("非法 partner key → 401", async () => {
    const r = await buildApp(baseEnv({ KBDB: makeKBDB(false) }))("Bearer pk_bad");
    expect(r.status).toBe(401);
  });
});

describe("partner-auth: 明碼逃生門（預設關）", () => {
  it("ALLOW_PLAINTEXT_NAMESPACE=true → 恢復舊明碼行為（遷移期）", async () => {
    const r = await buildApp(
      baseEnv({ MULTI_TENANT: "false", ALLOW_PLAINTEXT_NAMESPACE: "true" }),
    )("Bearer leo");
    expect(r.status).toBe(200);
    expect((await r.json()).org_namespace).toBe("leo");
  });

  it("預設（未設逃生門）明碼 namespace 被擋", async () => {
    const r = await buildApp(baseEnv({ MULTI_TENANT: "false" }))("Bearer leo");
    expect(r.status).toBe(401);
  });
});