Arcrun/mcp/tests/unit/oauth.test.ts

import { describe, it, expect } from "vitest";
import { Hono } from "hono";
import {
  randomToken,
  sha256Base64Url,
  sha256Hex,
  constantTimeEqual,
  verifyPkceS256,
} from "../../src/oauth/crypto.js";
import {
  putAuthCode,
  consumeAuthCode,
  putAccessToken,
  getAccessToken,
} from "../../src/oauth/store.js";
import {
  originOf,
  resourceUri,
  protectedResourceMetadata,
  authorizationServerMetadata,
  wwwAuthenticateHeader,
} from "../../src/oauth/metadata.js";
import { esc, consentPage } from "../../src/oauth/consent.js";
import { registerOAuthRoutes } from "../../src/oauth/routes.js";
import type { Env } from "../../src/types.js";

// ── 記憶體 KV mock（支援 expirationTtl → 用 exp 過期；delete）─────────────────────
function makeKV(): KVNamespace {
  const map = new Map<string, { value: string; exp?: number }>();
  const kv = {
    async put(key: string, value: string, opts?: { expirationTtl?: number }) {
      map.set(key, {
        value,
        exp: opts?.expirationTtl ? Date.now() + opts.expirationTtl * 1000 : undefined,
      });
    },
    async get(key: string) {
      const e = map.get(key);
      if (!e) return null;
      if (e.exp && e.exp < Date.now()) {
        map.delete(key);
        return null;
      }
      return e.value;
    },
    async delete(key: string) {
      map.delete(key);
    },
  };
  return kv as unknown as KVNamespace;
}

// PKCE：verifier "1234...43+ chars" → 直接算 challenge。
async function pkcePair() {
  const verifier = "a".repeat(64);
  const challenge = await sha256Base64Url(verifier);
  return { verifier, challenge };
}

function baseEnv(over: Partial<Env> = {}): Env {
  return {
    COMPONENT_REGISTRY: {} as Fetcher,
    CYPHER_EXECUTOR: {} as Fetcher,
    KBDB: {} as Fetcher,
    KBDB_INTERNAL_TOKEN: "internal",
    OAUTH_KV: makeKV(),
    MCP_OWNER_SECRET: "s3cr3t-owner",
    MCP_OWNER_NAMESPACE: "leo",
    ...over,
  } as Env;
}

// ── crypto ──────────────────────────────────────────────────────────────────────
describe("oauth/crypto", () => {
  it("randomToken 長度足夠且每次不同", () => {
    const a = randomToken();
    const b = randomToken();
    expect(a).not.toBe(b);
    expect(a.length).toBeGreaterThanOrEqual(40);
    expect(a).toMatch(/^[A-Za-z0-9\-_]+$/); // base64url 無 padding
  });

  it("sha256Hex/base64url 為已知值", async () => {
    // echo -n "abc" | sha256sum → ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad
    expect(await sha256Hex("abc")).toBe(
      "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
    );
    // RFC 7636 附錄範例：verifier → challenge
    expect(await sha256Base64Url("dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk")).toBe(
      "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM",
    );
  });

  it("constantTimeEqual 正確", () => {
    expect(constantTimeEqual("abc", "abc")).toBe(true);
    expect(constantTimeEqual("abc", "abd")).toBe(false);
    expect(constantTimeEqual("abc", "abcd")).toBe(false);
  });

  it("verifyPkceS256：S256 正確、拒 plain/短 verifier/竄改", async () => {
    const { verifier, challenge } = await pkcePair();
    expect(await verifyPkceS256(verifier, challenge, "S256")).toBe(true);
    expect(await verifyPkceS256(verifier, challenge, "plain")).toBe(false);
    expect(await verifyPkceS256(verifier, challenge, undefined)).toBe(false);
    expect(await verifyPkceS256("short", challenge, "S256")).toBe(false);
    expect(await verifyPkceS256("b".repeat(64), challenge, "S256")).toBe(false);
  });
});

// ── store（短效 KV）────────────────────────────────────────────────────────────
describe("oauth/store", () => {
  it("authorization code 一次性：consume 後即失效（防重放）", async () => {
    const kv = makeKV();
    await putAuthCode(kv, "code-1", {
      client_id: "c1",
      redirect_uri: "https://claude.ai/cb",
      code_challenge: "cc",
      code_challenge_method: "S256",
      scope: "mcp",
      resource: "https://mcp/mcp",
      namespace: "leo",
    });
    const first = await consumeAuthCode(kv, "code-1");
    expect(first?.namespace).toBe("leo");
    const second = await consumeAuthCode(kv, "code-1");
    expect(second).toBeNull();
  });

  it("access token 存取 + exp 過期回 null", async () => {
    const kv = makeKV();
    await putAccessToken(
      kv,
      "tok-1",
      { namespace: "leo", client_id: "c1", scope: "mcp", aud: "https://mcp/mcp", exp: Math.floor(Date.now() / 1000) + 100 },
      100,
    );
    expect((await getAccessToken(kv, "tok-1"))?.namespace).toBe("leo");
    // 已過期
    await putAccessToken(
      kv,
      "tok-2",
      { namespace: "leo", client_id: "c1", scope: "mcp", aud: "x", exp: Math.floor(Date.now() / 1000) - 10 },
      100,
    );
    expect(await getAccessToken(kv, "tok-2")).toBeNull();
  });

  it("KV key 為 hash（不把 raw token 當 key）", async () => {
    // 白盒：store 內部用 sha256Hex，這裡驗 hash 與 raw 不同即可
    expect(await sha256Hex("tok-1")).not.toContain("tok-1");
  });
});

// ── metadata ──────────────────────────────────────────────────────────────────
describe("oauth/metadata", () => {
  it("originOf / resourceUri", () => {
    expect(originOf("https://Mcp.Arcrun.dev/mcp")).toBe("https://mcp.arcrun.dev");
    expect(resourceUri("https://mcp.arcrun.dev")).toBe("https://mcp.arcrun.dev/mcp");
  });
  it("protectedResourceMetadata 必要欄位", () => {
    const m = protectedResourceMetadata("https://mcp.arcrun.dev");
    expect(m.resource).toBe("https://mcp.arcrun.dev/mcp");
    expect(m.authorization_servers).toEqual(["https://mcp.arcrun.dev"]);
  });
  it("authorizationServerMetadata 必要欄位（S256/code/none）", () => {
    const m = authorizationServerMetadata("https://mcp.arcrun.dev");
    expect(m.authorization_endpoint).toBe("https://mcp.arcrun.dev/authorize");
    expect(m.token_endpoint).toBe("https://mcp.arcrun.dev/token");
    expect(m.registration_endpoint).toBe("https://mcp.arcrun.dev/register");
    expect(m.response_types_supported).toEqual(["code"]);
    expect(m.code_challenge_methods_supported).toEqual(["S256"]);
    expect(m.token_endpoint_auth_methods_supported).toEqual(["none"]);
  });
  it("wwwAuthenticateHeader 指向 protected-resource metadata", () => {
    expect(wwwAuthenticateHeader("https://mcp.arcrun.dev")).toBe(
      'Bearer resource_metadata="https://mcp.arcrun.dev/.well-known/oauth-protected-resource"',
    );
  });
});

// ── consent escaping（XSS）────────────────────────────────────────────────────
describe("oauth/consent", () => {
  it("esc 跳脫 HTML 特殊字元", () => {
    expect(esc(`<script>"&'`)).toBe("&lt;script&gt;&quot;&amp;&#39;");
  });
  it("consentPage 把惡意 state 跳脫掉（不注入）", () => {
    const html = consentPage({
      client_id: "c1",
      redirect_uri: "https://claude.ai/cb",
      state: '"><img src=x onerror=alert(1)>',
      code_challenge: "cc",
      code_challenge_method: "S256",
      scope: "mcp",
      resource: "r",
    });
    expect(html).not.toContain("<img src=x");
    expect(html).toContain("&lt;img src=x");
  });
});

// ── 完整 OAuth 流程（route 整合）──────────────────────────────────────────────
function buildApp(env: Env) {
  const app = new Hono<{ Bindings: Env }>();
  registerOAuthRoutes(app);
  return {
    req: (path: string, init?: RequestInit) =>
      app.request("https://mcp.arcrun.dev" + path, init, env),
  };
}

describe("oauth flow (整合)", () => {
  it("well-known metadata 端點以 request origin 動態生成", async () => {
    const app = buildApp(baseEnv());
    const r = await app.req("/.well-known/oauth-protected-resource");
    expect(r.status).toBe(200);
    const j = await r.json();
    expect(j.resource).toBe("https://mcp.arcrun.dev/mcp");
    const r2 = await app.req("/.well-known/oauth-authorization-server");
    expect((await r2.json()).issuer).toBe("https://mcp.arcrun.dev");
  });

  it("/register 回 client_id（public client，無 secret）", async () => {
    const app = buildApp(baseEnv());
    const r = await app.req("/register", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({ redirect_uris: ["https://claude.ai/cb"], client_name: "Claude" }),
    });
    expect(r.status).toBe(201);
    const j = await r.json();
    expect(j.client_id).toMatch(/^mcp_/);
    expect(j.token_endpoint_auth_method).toBe("none");
    expect(j).not.toHaveProperty("client_secret");
  });

  it("/register 擋不允許的 redirect_uri host", async () => {
    const app = buildApp(baseEnv());
    const r = await app.req("/register", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({ redirect_uris: ["https://evil.example.com/cb"] }),
    });
    expect(r.status).toBe(400);
    expect((await r.json()).error).toBe("invalid_redirect_uri");
  });

  it("GET /authorize 要求 PKCE S256 且顯示同意頁", async () => {
    const app = buildApp(baseEnv());
    const { challenge } = await pkcePair();
    const ok = await app.req(
      `/authorize?response_type=code&client_id=c1&redirect_uri=${encodeURIComponent(
        "https://claude.ai/cb",
      )}&code_challenge=${challenge}&code_challenge_method=S256&state=xyz&scope=mcp`,
    );
    expect(ok.status).toBe(200);
    expect(await ok.text()).toContain("Owner 祕密");
    // 缺 PKCE → 400
    const bad = await app.req(
      `/authorize?response_type=code&client_id=c1&redirect_uri=${encodeURIComponent(
        "https://claude.ai/cb",
      )}`,
    );
    expect(bad.status).toBe(400);
  });

  it("GET /authorize：MCP_OWNER_SECRET 未設 → 503（不留不安全預設）", async () => {
    const app = buildApp(baseEnv({ MCP_OWNER_SECRET: undefined }));
    const { challenge } = await pkcePair();
    const r = await app.req(
      `/authorize?response_type=code&client_id=c1&redirect_uri=${encodeURIComponent(
        "https://claude.ai/cb",
      )}&code_challenge=${challenge}&code_challenge_method=S256`,
    );
    expect(r.status).toBe(503);
  });

  it("完整 code→token：正確 owner 祕密 + 正確 verifier → access_token", async () => {
    const env = baseEnv();
    const app = buildApp(env);
    const { verifier, challenge } = await pkcePair();
    const redirect = "https://claude.ai/cb";

    // POST /authorize 正確祕密 → 302 帶 code
    const authRes = await app.req("/authorize", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        client_id: "c1",
        redirect_uri: redirect,
        state: "st-1",
        code_challenge: challenge,
        code_challenge_method: "S256",
        scope: "mcp",
        resource: "https://mcp.arcrun.dev/mcp",
        owner_secret: "s3cr3t-owner",
      }).toString(),
      redirect: "manual",
    });
    expect(authRes.status).toBe(302);
    const loc = new URL(authRes.headers.get("location")!);
    expect(loc.searchParams.get("state")).toBe("st-1");
    const code = loc.searchParams.get("code")!;
    expect(code).toBeTruthy();

    // POST /token 正確 verifier → access_token
    const tokRes = await app.req("/token", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        grant_type: "authorization_code",
        code,
        code_verifier: verifier,
        redirect_uri: redirect,
        client_id: "c1",
      }).toString(),
    });
    expect(tokRes.status).toBe(200);
    const tok = await tokRes.json();
    expect(tok.token_type).toBe("Bearer");
    expect(tok.access_token).toBeTruthy();
    expect(tok.expires_in).toBeGreaterThan(0);

    // 換發的 token 綁定 owner namespace
    const at = await getAccessToken(env.OAUTH_KV!, tok.access_token);
    expect(at?.namespace).toBe("leo");
    expect(at?.aud).toBe("https://mcp.arcrun.dev/mcp");
  });

  it("錯誤 owner 祕密 → 401、不發 code", async () => {
    const app = buildApp(baseEnv());
    const { challenge } = await pkcePair();
    const r = await app.req("/authorize", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        client_id: "c1",
        redirect_uri: "https://claude.ai/cb",
        code_challenge: challenge,
        code_challenge_method: "S256",
        owner_secret: "WRONG",
      }).toString(),
      redirect: "manual",
    });
    expect(r.status).toBe(401);
    expect(r.headers.get("location")).toBeNull();
    expect(await r.text()).toContain("不正確");
  });

  it("/token 錯誤 verifier → invalid_grant；重用 code → invalid_grant", async () => {
    const env = baseEnv();
    const app = buildApp(env);
    const { challenge } = await pkcePair();
    const redirect = "https://claude.ai/cb";
    const authRes = await app.req("/authorize", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        client_id: "c1",
        redirect_uri: redirect,
        code_challenge: challenge,
        code_challenge_method: "S256",
        owner_secret: "s3cr3t-owner",
      }).toString(),
      redirect: "manual",
    });
    const code = new URL(authRes.headers.get("location")!).searchParams.get("code")!;

    // 錯誤 verifier
    const bad = await app.req("/token", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        grant_type: "authorization_code",
        code,
        code_verifier: "z".repeat(64),
        redirect_uri: redirect,
      }).toString(),
    });
    expect(bad.status).toBe(400);
    expect((await bad.json()).error).toBe("invalid_grant");

    // 該 code 已被 consume（即使失敗也一次性）→ 再用正確 verifier 也 invalid_grant
    const reuse = await app.req("/token", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        grant_type: "authorization_code",
        code,
        code_verifier: "a".repeat(64),
        redirect_uri: redirect,
      }).toString(),
    });
    expect((await reuse.json()).error).toBe("invalid_grant");
  });

  it("未設 OAUTH_KV → /token 回 503（誠實，不假綠）", async () => {
    const app = buildApp(baseEnv({ OAUTH_KV: undefined }));
    const r = await app.req("/token", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        grant_type: "authorization_code",
        code: "x",
        code_verifier: "a".repeat(64),
        redirect_uri: "https://claude.ai/cb",
      }).toString(),
    });
    expect(r.status).toBe(503);
  });
});

// ── 防 drift：對 OAUTH_KV 的每一次 put 都必須帶 expirationTtl（守儲存鐵律）──────────
// spy KV 記錄所有 put(key,value,opts)；跑完整流程後斷言沒有任何一次「無 TTL」的 put，
// 防未來有人往這顆短效 KV 塞長效資料（access_token / code 以外的東西）。
describe("oauth store drift guard：OAUTH_KV 的 put 一律帶 TTL", () => {
  function spyKV(): { kv: KVNamespace; puts: Array<{ key: string; opts?: { expirationTtl?: number } }> } {
    const map = new Map<string, string>();
    const puts: Array<{ key: string; opts?: { expirationTtl?: number } }> = [];
    const kv = {
      async put(key: string, value: string, opts?: { expirationTtl?: number }) {
        puts.push({ key, opts });
        map.set(key, value);
      },
      async get(key: string) {
        return map.get(key) ?? null;
      },
      async delete(key: string) {
        map.delete(key);
      },
    } as unknown as KVNamespace;
    return { kv, puts };
  }

  it("完整 authorize→token 流程中，OAUTH_KV 的每次 put 都有 expirationTtl>0", async () => {
    const { kv, puts } = spyKV();
    const env = baseEnv({ OAUTH_KV: kv });
    const app = buildApp(env);
    const { verifier, challenge } = await pkcePair();
    const redirect = "https://claude.ai/cb";

    const authRes = await app.req("/authorize", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        client_id: "c1",
        redirect_uri: redirect,
        code_challenge: challenge,
        code_challenge_method: "S256",
        owner_secret: "s3cr3t-owner",
      }).toString(),
      redirect: "manual",
    });
    const code = new URL(authRes.headers.get("location")!).searchParams.get("code")!;
    await app.req("/token", {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded" },
      body: new URLSearchParams({
        grant_type: "authorization_code",
        code,
        code_verifier: verifier,
        redirect_uri: redirect,
        client_id: "c1",
      }).toString(),
    });

    // 至少發生了 put（code + token 各一），且每一次都帶 TTL。
    expect(puts.length).toBeGreaterThanOrEqual(2);
    for (const p of puts) {
      expect(p.opts?.expirationTtl, `put ${p.key} 缺 expirationTtl`).toBeTypeOf("number");
      expect(p.opts!.expirationTtl!).toBeGreaterThan(0);
    }
  });

  it("直接呼叫 store 層 putAuthCode / putAccessToken 也一律帶 TTL", async () => {
    const { kv, puts } = spyKV();
    await putAuthCode(kv, "c", {
      client_id: "c1",
      redirect_uri: "https://claude.ai/cb",
      code_challenge: "cc",
      code_challenge_method: "S256",
      scope: "mcp",
      resource: "https://mcp/mcp",
      namespace: "leo",
    });
    await putAccessToken(
      kv,
      "t",
      { namespace: "leo", client_id: "c1", scope: "mcp", aud: "https://mcp/mcp", exp: 1 },
      100,
    );
    expect(puts).toHaveLength(2);
    for (const p of puts) {
      expect(p.opts?.expirationTtl).toBeGreaterThan(0);
    }
  });
});