diff --git a/packages/agent-auth/package.json b/packages/agent-auth/package.json index 52457a78e..1d65dede1 100644 --- a/packages/agent-auth/package.json +++ b/packages/agent-auth/package.json @@ -1,6 +1,6 @@ { "name": "@better-auth/agent-auth", - "version": "0.4.5", + "version": "0.4.6", "description": "Agent Auth Protocol server plugin for Better Auth — agent identity, registration, and capability-based authorization", "license": "MIT", "repository": { diff --git a/packages/agent-auth/src/__tests__/register-host-lookup.test.ts b/packages/agent-auth/src/__tests__/register-host-lookup.test.ts new file mode 100644 index 000000000..3724dbe13 --- /dev/null +++ b/packages/agent-auth/src/__tests__/register-host-lookup.test.ts @@ -0,0 +1,181 @@ +import { describe, expect, it } from "vitest"; +import { getTestInstance } from "better-auth/test"; +import { + agentAuth, + agentAuthClientPlugin, + generateTestKeypair, + createHostJWT, + json, + createTestClient, + computeThumbprint, + BASE, +} from "./helpers"; + +/** + * Regression: registerAgent must look up hosts by both `id` and `kid`. + * + * The SDK signs host JWTs with `iss = JWK thumbprint` (§4.2). For hosts + * provisioned via the enrollment-token flow, the thumbprint lives in the + * `kid` column, not `id`. The middleware does a dual lookup + * (findHostById ?? findHostByKid); the register route previously did an + * id-only lookup, so pre-enrolled hosts fell through to the dynamic- + * registration branch and failed with DYNAMIC_HOST_REGISTRATION_DISABLED + * when the operator had disabled dynamic registration. + * + * Reported by .entomb on Discord, 2026-04-20. + */ +describe("registerAgent — pre-enrolled host lookup by kid (issue 1)", () => { + it("registers an agent when iss = JWK thumbprint and dynamic registration is disabled", async () => { + const t = await getTestInstance( + { + plugins: [ + agentAuth({ + providerName: "test-service", + allowDynamicHostRegistration: false, + capabilities: [{ name: "ping", description: "ping" }], + }), + ], + }, + { + clientOptions: { plugins: [agentAuthClientPlugin()] }, + }, + ); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const auth = t.auth as any; + const client = createTestClient((req: Request) => auth.handler(req)); + + const { headers } = await t.signInWithTestUser(); + const sessionCookie = headers.get("cookie") ?? ""; + + // 1. Provision a host in pending_enrollment (no public key supplied) + const provisionRes = await client.authedPost( + "/host/create", + { name: "Pre-enrolled host" }, + sessionCookie, + ); + expect(provisionRes.ok).toBe(true); + const { hostId, enrollmentToken, status } = await json<{ + hostId: string; + enrollmentToken: string; + status: string; + }>(provisionRes); + expect(status).toBe("pending_enrollment"); + + // 2. Device enrolls with its keypair — this populates the `kid` + // column with the JWK thumbprint and flips status to "active". + const hostKeypair = await generateTestKeypair(); + const thumbprint = await computeThumbprint(hostKeypair.publicKey); + const publicKeyWithKid = { ...hostKeypair.publicKey, kid: thumbprint }; + + const enrollRes = await client.api("/host/enroll", { + method: "POST", + body: JSON.stringify({ + token: enrollmentToken, + public_key: publicKeyWithKid, + }), + }); + const enrollBody = await json>(enrollRes); + expect(enrollRes.ok, JSON.stringify(enrollBody)).toBe(true); + expect(enrollBody.hostId).toBe(hostId); + expect(enrollBody.status).toBe("active"); + + // 3. Register an agent — SDK sends iss = thumbprint (per spec §4.2), + // not the host's UUID id. Without the kid lookup this hits the + // dynamic-registration branch and fails. + const agentKeypair = await generateTestKeypair(); + const hostJWT = await createHostJWT( + hostKeypair.privateKey, + publicKeyWithKid, + agentKeypair.publicKey, + thumbprint, + ); + + const res = await client.api("/agent/register", { + method: "POST", + headers: { authorization: `Bearer ${hostJWT}` }, + body: JSON.stringify({ name: "Test Agent", capabilities: ["ping"] }), + }); + const body = await json>(res); + expect(res.ok, `expected 2xx but got ${res.status}: ${JSON.stringify(body)}`).toBe(true); + expect(body).not.toHaveProperty("error", "dynamic_host_registration_disabled"); + expect(typeof body.agent_id).toBe("string"); + }); + + // Sanity check: BASE is referenced indirectly via createHostJWT's audience. + it("BASE is defined for the test client", () => { + expect(BASE).toMatch(/^http/); + }); +}); + +/** + * Regression: /agent/claim has the same single-field host lookup as + * /agent/register did. Same root cause, same fix — covered here so we + * don't regress when someone refactors one route without the other. + */ +describe("claimAgent — pre-enrolled host lookup by kid (issue 1, claim variant)", () => { + it("authenticates a pre-enrolled host with iss=thumbprint when dynamic registration is disabled", async () => { + const t = await getTestInstance( + { + plugins: [ + agentAuth({ + providerName: "test-service", + allowDynamicHostRegistration: false, + capabilities: [{ name: "ping", description: "ping" }], + }), + ], + }, + { + clientOptions: { plugins: [agentAuthClientPlugin()] }, + }, + ); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const auth = t.auth as any; + const client = createTestClient((req: Request) => auth.handler(req)); + + const { headers } = await t.signInWithTestUser(); + const sessionCookie = headers.get("cookie") ?? ""; + + const provisionRes = await client.authedPost( + "/host/create", + { name: "Pre-enrolled host" }, + sessionCookie, + ); + const { enrollmentToken } = await json<{ enrollmentToken: string }>(provisionRes); + + const hostKeypair = await generateTestKeypair(); + const thumbprint = await computeThumbprint(hostKeypair.publicKey); + const publicKeyWithKid = { ...hostKeypair.publicKey, kid: thumbprint }; + + await client.api("/host/enroll", { + method: "POST", + body: JSON.stringify({ + token: enrollmentToken, + public_key: publicKeyWithKid, + }), + }); + + const agentKeypair = await generateTestKeypair(); + const hostJWT = await createHostJWT( + hostKeypair.privateKey, + publicKeyWithKid, + agentKeypair.publicKey, + thumbprint, + ); + + // Target a non-existent agent. We only care that host authentication + // passes — not that the claim itself succeeds. Before the fix the + // host lookup missed and the call failed with + // DYNAMIC_HOST_REGISTRATION_DISABLED. After the fix it should reach + // the claim logic and fail with AGENT_NOT_FOUND instead. + const res = await client.api("/agent/claim", { + method: "POST", + headers: { authorization: `Bearer ${hostJWT}` }, + body: JSON.stringify({ agent_id: "nonexistent-agent-id" }), + }); + + expect(res.ok).toBe(false); + const body = await json>(res); + expect(body.error).not.toBe("dynamic_host_registration_disabled"); + expect(body.error).toBe("agent_not_found"); + }); +}); diff --git a/packages/agent-auth/src/routes/_helpers.ts b/packages/agent-auth/src/routes/_helpers.ts index 1cab90177..9997bdebf 100644 --- a/packages/agent-auth/src/routes/_helpers.ts +++ b/packages/agent-auth/src/routes/_helpers.ts @@ -82,6 +82,29 @@ export function resolveDeviceAuthPage(opts: ResolvedAgentAuthOptions, origin: st return `${origin}${path}`; } +/** + * Resolve a host by the `iss` claim of a host JWT, matching the middleware. + * + * SDK-signed JWTs use `iss = JWK thumbprint` (spec §4.2), stored in `kid`. + * The `id` lookup is a liberal fallback for hand-crafted callers that sign + * with the host's UUID instead — tightening it to kid-only is a follow-up. + */ +export async function findHostByIdOrKid( + adapter: AdapterFindOne, + iss: string, +): Promise { + return ( + (await adapter.findOne({ + model: TABLE.host, + where: [{ field: "id", value: iss }], + })) ?? + (await adapter.findOne({ + model: TABLE.host, + where: [{ field: "kid", value: iss }], + })) + ); +} + export async function findHostByKey( adapter: AdapterFindOne, publicKey: Record, diff --git a/packages/agent-auth/src/routes/claim.ts b/packages/agent-auth/src/routes/claim.ts index 3521b7a31..c77dddf6f 100644 --- a/packages/agent-auth/src/routes/claim.ts +++ b/packages/agent-auth/src/routes/claim.ts @@ -19,6 +19,7 @@ import type { } from "../types"; import { buildApprovalInfo, + findHostByIdOrKid, findHostByKey, formatGrantsResponse, isDynamicHostAllowed, @@ -112,10 +113,7 @@ export function claimAgent( let hostRecord: AgentHost | null = null; if (hostIdFromJwt) { - hostRecord = await ctx.context.adapter.findOne({ - model: TABLE.host, - where: [{ field: "id", value: hostIdFromJwt }], - }); + hostRecord = await findHostByIdOrKid(ctx.context.adapter, hostIdFromJwt); } if (hostRecord) { @@ -146,7 +144,9 @@ export function claimAgent( publicKey: hostPubKey, maxAge: opts.jwtMaxAge, }); - if (!payload || payload.iss !== hostRecord.id) { + // §4.2: iss identifies the host by either its registered id or + // its JWK thumbprint (stored in the `kid` column for pre-enrolled hosts). + if (!payload || (payload.iss !== hostRecord.id && payload.iss !== hostRecord.kid)) { throw agentError("UNAUTHORIZED", ERR.INVALID_JWT); } diff --git a/packages/agent-auth/src/routes/register.ts b/packages/agent-auth/src/routes/register.ts index cb07f6ab5..5b37e750d 100644 --- a/packages/agent-auth/src/routes/register.ts +++ b/packages/agent-auth/src/routes/register.ts @@ -24,6 +24,7 @@ import { buildApprovalInfo, capabilityItemZ, createGrantRows, + findHostByIdOrKid, findHostByKey, formatGrantsResponse, isDynamicHostAllowed, @@ -167,10 +168,7 @@ export function register( let hostRecord: AgentHost | null = null; if (hostIdFromJwt) { - hostRecord = await ctx.context.adapter.findOne({ - model: TABLE.host, - where: [{ field: "id", value: hostIdFromJwt }], - }); + hostRecord = await findHostByIdOrKid(ctx.context.adapter, hostIdFromJwt); } if (hostRecord) { @@ -212,8 +210,9 @@ export function register( maxAge: opts.jwtMaxAge, }); - // §4.2: iss identifies the host - if (!payload || payload.iss !== hostRecord.id) { + // §4.2: iss identifies the host by either its registered id or + // its JWK thumbprint (stored in the `kid` column for pre-enrolled hosts). + if (!payload || (payload.iss !== hostRecord.id && payload.iss !== hostRecord.kid)) { throw agentError("UNAUTHORIZED", ERR.INVALID_JWT); }