From 9f29f019faf8fa525d02b029e56605da7b0cffce Mon Sep 17 00:00:00 2001 From: Dotta Date: Wed, 22 Jul 2026 16:00:58 +0000 Subject: [PATCH] fix(connections): make broker claims retryable Co-Authored-By: Paperclip --- packages/connect-protocol/src/index.test.ts | 4 +- packages/connect-protocol/src/index.ts | 5 ++- server/src/services/connection-broker.test.ts | 42 +++++++++++++++---- server/src/services/connection-broker.ts | 4 +- 4 files changed, 43 insertions(+), 12 deletions(-) diff --git a/packages/connect-protocol/src/index.test.ts b/packages/connect-protocol/src/index.test.ts index e9c0a92d1d8..4a3d0675673 100644 --- a/packages/connect-protocol/src/index.test.ts +++ b/packages/connect-protocol/src/index.test.ts @@ -114,7 +114,7 @@ describe("connect protocol", () => { const { privateKey, publicKey } = generateKeyPairSync("ed25519"); const store = replayStore(); const compactJws = signRequestEnvelope({ - body: { claimCode: "claim-code-with-enough-entropy" }, + body: { claimCode: "claim-code-with-enough-entropy", idempotencyKey: "a".repeat(64) }, instanceId: "in_devinstance", keyId: "key_1", privateKey, @@ -131,7 +131,7 @@ describe("connect protocol", () => { replayStore: store, now: new Date("2026-07-21T18:00:01.000Z"), }; - await expect(verifyRequestEnvelope(input)).resolves.toMatchObject({ body: { claimCode: "claim-code-with-enough-entropy" } }); + await expect(verifyRequestEnvelope(input)).resolves.toMatchObject({ body: { claimCode: "claim-code-with-enough-entropy", idempotencyKey: "a".repeat(64) } }); await expect(verifyRequestEnvelope(input)).rejects.toMatchObject({ code: "replayed_jti", status: 401 }); }); diff --git a/packages/connect-protocol/src/index.ts b/packages/connect-protocol/src/index.ts index 5b9f85dec89..4ff73747cd0 100644 --- a/packages/connect-protocol/src/index.ts +++ b/packages/connect-protocol/src/index.ts @@ -52,7 +52,10 @@ export const handshakeResponseSchema = z.object({ expiresAt: z.string().datetime(), }); -export const claimRequestSchema = z.object({ claimCode: z.string().min(22) }); +export const claimRequestSchema = z.object({ + claimCode: z.string().min(22), + idempotencyKey: z.string().regex(/^[a-f0-9]{64}$/), +}); export const tokenMetadataSchema = z.object({ scopes: z.array(z.string()), diff --git a/server/src/services/connection-broker.test.ts b/server/src/services/connection-broker.test.ts index 5052ad874e0..6e18ba62536 100644 --- a/server/src/services/connection-broker.test.ts +++ b/server/src/services/connection-broker.test.ts @@ -17,7 +17,7 @@ const CONNECTION_ID = "44444444-4444-4444-8444-444444444444"; function fixture(mode: "A" | "B2" = "A") { const { privateKey, publicKey } = generateKeyPairSync("ed25519"); const replay = new Set(); - const claims = new Map([["claim-code-at-least-22-chars", { owner: INSTANCE_ID, attempts: 0, consumed: false }]]); + const claims = new Map([["claim-code-at-least-22-chars", { owner: INSTANCE_ID, attempts: 0, consumed: false, idempotencyKey: "", response: null as Record | null }]]); const saved: Parameters[0][] = []; const secrets: Array<{ purpose: string; value: string }> = []; const replayStore: ReplayStore = { consume: (_instance, jti) => !replay.has(jti) && Boolean(replay.add(jti)) }; @@ -28,16 +28,22 @@ function fixture(mode: "A" | "B2" = "A") { try { const envelope = await verifyRequestEnvelope({ compactJws, publicKey, expectedInstanceId: INSTANCE_ID, expectedPath: path, replayStore }); if (path === "/v1/claims") { - const code = (envelope.body as { claimCode: string }).claimCode; + const { claimCode: code, idempotencyKey } = envelope.body as { claimCode: string; idempotencyKey: string }; const claim = claims.get(code); if (!claim) return { status: 404, body: { error: "unknown_claim" } }; if (claim.attempts >= 5) return { status: 404, body: { error: "unknown_claim" } }; if (claim.owner !== envelope.iss) { claim.attempts += 1; return { status: 403, body: { error: "claim_instance_mismatch" } }; } - if (claim.consumed) return { status: 409, body: { error: "claim_already_consumed" } }; + if (claim.consumed) { + return claim.idempotencyKey === idempotencyKey && claim.response + ? { status: 200, body: claim.response } + : { status: 409, body: { error: "claim_already_consumed" } }; + } claim.consumed = true; - return mode === "A" - ? { status: 200, body: { custodyMode: "A", grantRef: "gr_testgrant", tokenMeta: { scopes: ["read"], expiresAt: null } } } - : { status: 200, body: { custodyMode: "B2", grantSealed: "sealed-refresh", tokenMeta: { scopes: ["read"], expiresAt: null } } }; + claim.idempotencyKey = idempotencyKey; + claim.response = mode === "A" + ? { custodyMode: "A", grantRef: "gr_testgrant", tokenMeta: { scopes: ["read"], expiresAt: null } } + : { custodyMode: "B2", grantSealed: "sealed-refresh", tokenMeta: { scopes: ["read"], expiresAt: null } }; + return { status: 200, body: claim.response }; } if (path.endsWith("/token")) return { status: 200, body: { grantSealed: "sealed-access", tokenMeta: { scopes: ["read"], expiresAt: null } } }; if (path === "/v1/handshakes") return { status: 201, body: { handshakeId: "hs_test", authorizeUrl: "https://connect.example/authorize", expiresAt: new Date(Date.now() + 60_000).toISOString() } }; @@ -74,7 +80,7 @@ describe("connection broker", () => { it("surfaces a 409 redemption as an interception alert", async () => { const f = fixture(); await f.service.claimOAuth({ companyId: "company", connectionId: CONNECTION_ID, claimCode: "claim-code-at-least-22-chars", kind: "workspace" }); - await expect(f.service.claimOAuth({ companyId: "company", connectionId: CONNECTION_ID, claimCode: "claim-code-at-least-22-chars", kind: "workspace" })).rejects.toBeInstanceOf(ConnectClaimInterceptionError); + await expect(f.client.claim("claim-code-at-least-22-chars", "b".repeat(64))).rejects.toBeInstanceOf(ConnectClaimInterceptionError); }); it("does not consume mismatched claims and locks the service-side double after five attempts", async () => { @@ -122,6 +128,28 @@ describe("connection broker", () => { expect(f.claimStates.size).toBe(0); }); + it("retries the remote claim idempotently when local recovery persistence fails", async () => { + const f = fixture("B2"); + let failSave = true; + const service = connectionBrokerService({ + client: f.client, + store: { + ...f.store, + saveClaimState: async (input) => { + if (failSave) { failSave = false; throw new Error("database unavailable"); } + await f.store.saveClaimState(input); + }, + }, + enabled: true, + vault: { put: async (input) => ({ secretId: "secret-grant", versionSelector: "latest", configPath: "broker.grant", required: true, label: input.purpose }) }, + }); + const params = { companyId: "company", connectionId: CONNECTION_ID, claimCode: "claim-code-at-least-22-chars", kind: "workspace" as const }; + + await expect(service.claimOAuth(params)).rejects.toThrow("database unavailable"); + expect(f.claims.get(params.claimCode)?.consumed).toBe(true); + await expect(service.claimOAuth(params)).resolves.toMatchObject({ custodyMode: "B2" }); + }); + it("rejects user grants before consuming a claim when membership is inactive", async () => { const f = fixture(); const service = connectionBrokerService({ diff --git a/server/src/services/connection-broker.ts b/server/src/services/connection-broker.ts index f44971b454a..946d376e15a 100644 --- a/server/src/services/connection-broker.ts +++ b/server/src/services/connection-broker.ts @@ -28,7 +28,7 @@ export function createSignedConnectBrokerClient(input: { instanceId: string; key } return { createHandshake: async (body: unknown) => handshakeResponseSchema.parse(await request("/v1/handshakes", handshakeRequestSchema.parse(body))), - claim: async (claimCode: string) => claimResponseSchema.parse(await request("/v1/claims", claimRequestSchema.parse({ claimCode }))), + claim: async (claimCode: string, idempotencyKey: string) => claimResponseSchema.parse(await request("/v1/claims", claimRequestSchema.parse({ claimCode, idempotencyKey }))), mintGrantToken: async (grantRef: string, grantWrapKey: string) => grantTokenResponseSchema.parse(await request(`/v1/grants/${encodeURIComponent(grantRef)}/token`, grantTokenRequestSchema.parse({ grantWrapKey }))), registerRelay: async (body: unknown) => relayRegistrationResponseSchema.parse(await request("/v1/connections/register", relayRegistrationRequestSchema.parse(body))), }; @@ -112,7 +112,7 @@ export function connectionBrokerService(input: { client: ConnectBrokerClient; st } const claimCodeHash = createHash("sha256").update(params.claimCode).digest("hex"); const existingState = await input.store.loadClaimState({ companyId: params.companyId, connectionId: params.connectionId, claimCodeHash }); - const claim = existingState?.claim ?? await input.client.claim(params.claimCode); + const claim = existingState?.claim ?? await input.client.claim(params.claimCode, claimCodeHash); if (!existingState) await input.store.saveClaimState({ companyId: params.companyId, connectionId: params.connectionId, claimCodeHash, claim }); if (claim.custodyMode !== "A" && claim.custodyMode !== "B2") throw new ConnectServiceResponseError(403, "mode_not_allowed"); const serviceGrantRef = claim.custodyMode === "A" ? claim.grantRef : null;