diff --git a/test/integration/worker/setup-bootstrap-lock.test.ts b/test/integration/worker/setup-bootstrap-lock.test.ts new file mode 100644 index 00000000..f82dc25f --- /dev/null +++ b/test/integration/worker/setup-bootstrap-lock.test.ts @@ -0,0 +1,71 @@ +import { env } from "cloudflare:test"; +import { beforeAll, beforeEach, describe, expect, it } from "vitest"; + +import { + claimBootstrapLock, + releaseBootstrapLock, + renewBootstrapLock +} from "../../../worker/features/setup/bootstrap-lock"; +import { applyCurrentMigrations } from "./current-migrations"; + +describe("setup bootstrap lock", () => { + beforeAll(async () => { + await applyCurrentMigrations(); + }); + + beforeEach(async () => { + await env.DB.prepare("DELETE FROM app_settings WHERE key = 'setup_bootstrap_lock'").run(); + }); + + it("lets only one concurrent setup claim the fresh workspace", async () => { + const claims = await Promise.allSettled([ + claimBootstrapLock(env.DB), + claimBootstrapLock(env.DB) + ]); + const fulfilled = claims.filter( + (claim): claim is PromiseFulfilledResult>> => + claim.status === "fulfilled" + ); + const rejected = claims.filter( + (claim): claim is PromiseRejectedResult => claim.status === "rejected" + ); + + expect(fulfilled).toHaveLength(1); + expect(rejected).toHaveLength(1); + expect(rejected[0]?.reason).toMatchObject({ code: "SETUP_IN_PROGRESS", status: 409 }); + + const winner = fulfilled[0]; + if (!winner) throw new Error("Expected one bootstrap lock claim to succeed."); + await releaseBootstrapLock(env.DB, winner.value); + await expect(claimBootstrapLock(env.DB)).resolves.toBeDefined(); + }); + + it("reclaims a lock at the lease boundary without letting its old owner release the new claim", async () => { + const first = await claimBootstrapLock(env.DB, new Date("2026-08-22T12:00:00.000Z")); + const replacement = await claimBootstrapLock(env.DB, new Date("2026-08-22T12:05:00.000Z")); + + await releaseBootstrapLock(env.DB, first); + await expect( + claimBootstrapLock(env.DB, new Date("2026-08-22T12:06:01.000Z")) + ).rejects.toMatchObject({ code: "SETUP_IN_PROGRESS", status: 409 }); + + await releaseBootstrapLock(env.DB, replacement); + await expect( + claimBootstrapLock(env.DB, new Date("2026-08-22T12:06:02.000Z")) + ).resolves.toBeDefined(); + }); + + it("keeps an active bootstrap claim beyond one lease through renewal", async () => { + const first = await claimBootstrapLock(env.DB, new Date("2026-08-22T12:00:00.000Z")); + await renewBootstrapLock(env.DB, first, new Date("2026-08-22T12:04:30.000Z")); + + await expect( + claimBootstrapLock(env.DB, new Date("2026-08-22T12:05:01.000Z")) + ).rejects.toMatchObject({ code: "SETUP_IN_PROGRESS", status: 409 }); + + await renewBootstrapLock(env.DB, first, new Date("2026-08-22T12:09:00.000Z")); + await expect( + claimBootstrapLock(env.DB, new Date("2026-08-22T12:10:00.000Z")) + ).rejects.toMatchObject({ code: "SETUP_IN_PROGRESS", status: 409 }); + }); +}); diff --git a/test/unit/worker/features/setup/bootstrap-security.test.ts b/test/unit/worker/features/setup/bootstrap-security.test.ts new file mode 100644 index 00000000..6e22fd84 --- /dev/null +++ b/test/unit/worker/features/setup/bootstrap-security.test.ts @@ -0,0 +1,54 @@ +import { + type BootstrapLock, + startBootstrapLockHeartbeat +} from "@worker/features/setup/bootstrap-lock"; +import { requireDirectBootstrapClientIp } from "@worker/features/setup/routes"; +import { describe, expect, it, vi } from "vitest"; + +describe("setup bootstrap security", () => { + it("rejects Worker-originated and unidentified bootstrap requests", () => { + expect(() => + requireDirectBootstrapClientIp( + new Request("https://hqbase.test/api/setup/bootstrap", { + headers: { "cf-connecting-ip": "192.0.2.10", "cf-worker": "example.com" } + }) + ) + ).toThrowError(expect.objectContaining({ code: "SETUP_DIRECT_REQUEST_REQUIRED", status: 403 })); + + expect(() => + requireDirectBootstrapClientIp( + new Request("https://hqbase.test/api/setup/bootstrap", { + headers: { "cf-connecting-ip": " " } + }) + ) + ).toThrowError(expect.objectContaining({ code: "SETUP_CLIENT_IP_REQUIRED", status: 403 })); + }); + + it("returns the Cloudflare client IP for a direct request", () => { + expect( + requireDirectBootstrapClientIp( + new Request("https://hqbase.test/api/setup/bootstrap", { + headers: { "cf-connecting-ip": " 192.0.2.10 " } + }) + ) + ).toBe("192.0.2.10"); + }); + + it("renews an active lock on the heartbeat interval", async () => { + vi.useFakeTimers(); + try { + const lock: BootstrapLock = { value: '{"token":"test"}' }; + const first = vi.fn(async () => ({ value_json: lock.value })); + const bind = vi.fn(() => ({ first })); + const db = { prepare: vi.fn(() => ({ bind })) } as unknown as D1Database; + const heartbeat = startBootstrapLockHeartbeat(db, lock, 100); + + await vi.advanceTimersByTimeAsync(250); + await heartbeat.stop(); + + expect(first).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/worker/features/setup/bootstrap-lock.ts b/worker/features/setup/bootstrap-lock.ts new file mode 100644 index 00000000..7447ba47 --- /dev/null +++ b/worker/features/setup/bootstrap-lock.ts @@ -0,0 +1,96 @@ +import { AppError } from "../../lib/errors"; + +const bootstrapLockKey = "setup_bootstrap_lock"; +const bootstrapLockTtlMs = 5 * 60 * 1000; +const bootstrapLockHeartbeatMs = 30 * 1000; + +export type BootstrapLock = { + value: string; +}; + +export type BootstrapLockHeartbeat = { + renew: () => Promise; + stop: () => Promise; +}; + +export async function claimBootstrapLock(db: D1Database, now = new Date()): Promise { + const value = JSON.stringify({ token: crypto.randomUUID() }); + const timestamp = now.toISOString(); + const staleBefore = new Date(now.getTime() - bootstrapLockTtlMs).toISOString(); + const claimed = await db + .prepare( + `INSERT INTO app_settings (key, value_json, created_at, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + created_at = excluded.created_at, + updated_at = excluded.updated_at + WHERE app_settings.updated_at <= ? + RETURNING value_json` + ) + .bind(bootstrapLockKey, value, timestamp, timestamp, staleBefore) + .first<{ value_json: string }>(); + + if (claimed?.value_json !== value) { + throw new AppError("SETUP_IN_PROGRESS", "Setup is already being completed.", 409); + } + return { value }; +} + +export async function renewBootstrapLock( + db: D1Database, + lock: BootstrapLock, + now = new Date() +): Promise { + const renewed = await db + .prepare( + `UPDATE app_settings + SET updated_at = ? + WHERE key = ? AND value_json = ? + RETURNING value_json` + ) + .bind(now.toISOString(), bootstrapLockKey, lock.value) + .first<{ value_json: string }>(); + if (renewed?.value_json !== lock.value) { + throw new AppError("SETUP_LOCK_LOST", "Setup lost its exclusive bootstrap claim.", 409); + } +} + +export function startBootstrapLockHeartbeat( + db: D1Database, + lock: BootstrapLock, + intervalMs = bootstrapLockHeartbeatMs +): BootstrapLockHeartbeat { + let pending = Promise.resolve(); + let failure: unknown; + const enqueueRenewal = () => { + const renewal = pending.then(() => renewBootstrapLock(db, lock)); + pending = renewal.catch((error: unknown) => { + failure ??= error; + }); + return renewal; + }; + const timer = setInterval(() => { + void enqueueRenewal().catch(() => undefined); + }, intervalMs); + + return { + async renew() { + if (failure) throw failure; + await enqueueRenewal(); + if (failure) throw failure; + }, + async stop() { + clearInterval(timer); + await pending; + if (failure) throw failure; + } + }; +} + +export async function releaseBootstrapLock(db: D1Database, lock: BootstrapLock): Promise { + await db + .prepare("DELETE FROM app_settings WHERE key = ? AND value_json = ?") + .bind(bootstrapLockKey, lock.value) + .run(); +} diff --git a/worker/features/setup/routes.ts b/worker/features/setup/routes.ts index e4536975..10557dd2 100644 --- a/worker/features/setup/routes.ts +++ b/worker/features/setup/routes.ts @@ -1,7 +1,9 @@ import { Hono } from "hono"; import type { HonoApp } from "../../lib/env"; +import { AppError } from "../../lib/errors"; import { readJson } from "../../lib/json"; import { parseWith } from "../../lib/validation"; +import { enforceRateLimit } from "../../security/rate-limit"; import { clearRuntimeCloudflareGrantCookie, finishRuntimeCloudflareOAuth, @@ -88,6 +90,13 @@ setupRoutes.post("/cloudflare/configure", async (c) => { }); setupRoutes.post("/bootstrap", async (c) => { + const ip = requireDirectBootstrapClientIp(c.req.raw); + await enforceRateLimit(c.env.DB, c.env.BETTER_AUTH_SECRET, { + scope: "setup.bootstrap.ip", + subject: ip, + limit: 5, + windowSeconds: 15 * 60 + }); const input = parseWith(bootstrapSetupSchema, await readJson(c.req.raw)); const grant = await resolveRuntimeCloudflareGrant(c.req.raw, c.env); const result = await bootstrapSetup(c.env, c.req.raw, input); @@ -95,3 +104,22 @@ setupRoutes.post("/bootstrap", async (c) => { c.header("set-cookie", clearRuntimeCloudflareGrantCookie()); return c.json(result, 201); }); + +export function requireDirectBootstrapClientIp(request: Request): string { + if (request.headers.get("cf-worker")?.trim()) { + throw new AppError( + "SETUP_DIRECT_REQUEST_REQUIRED", + "Complete setup directly in a browser.", + 403 + ); + } + const ip = request.headers.get("cf-connecting-ip")?.trim(); + if (!ip) { + throw new AppError( + "SETUP_CLIENT_IP_REQUIRED", + "Cloudflare client IP information is required to complete setup.", + 403 + ); + } + return ip; +} diff --git a/worker/features/setup/service.ts b/worker/features/setup/service.ts index b8d72b13..3c4a2a51 100644 --- a/worker/features/setup/service.ts +++ b/worker/features/setup/service.ts @@ -7,6 +7,11 @@ import { createMailbox } from "../mailboxes/service"; import type { Mailbox } from "../mailboxes/types"; import { setDefaultFromMailboxId } from "../preferences/queries"; +import { + claimBootstrapLock, + releaseBootstrapLock, + startBootstrapLockHeartbeat +} from "./bootstrap-lock"; import { getSetupStatus, setChecklistAcknowledged, @@ -42,72 +47,83 @@ export async function bootstrapSetup( request: Request, input: BootstrapInput ): Promise { - const existing = await getSetupStatus(env.DB); - if (existing.isComplete) { - throw new AppError("SETUP_ALREADY_COMPLETE", "Setup is already complete.", 409); - } - if (existing.userCount > 0) { - throw new AppError("SETUP_OWNER_EXISTS", "An owner user already exists.", 409); - } - - const domains = input.emailDomains ?? [{ name: input.primaryDomain ?? "" }]; - if (!domains[0]?.name) throw new AppError("DOMAIN_REQUIRED", "Choose an email domain.", 400); - assertLoginEmailOutsideDomains( - input.ownerEmail, - domains.map((domain) => domain.name) - ); + const lock = await claimBootstrapLock(env.DB); + const heartbeat = startBootstrapLockHeartbeat(env.DB, lock); + try { + const existing = await getSetupStatus(env.DB); + if (existing.isComplete) { + throw new AppError("SETUP_ALREADY_COMPLETE", "Setup is already complete.", 409); + } + if (existing.userCount > 0) { + throw new AppError("SETUP_OWNER_EXISTS", "An owner user already exists.", 409); + } - for (const domain of domains) { - await upsertMailDomain(env.DB, { - ...domain, - receivingStatus: "ready", - sendingStatus: "ready", - dnsStatus: "ready" - }); - } + const domains = input.emailDomains ?? [{ name: input.primaryDomain ?? "" }]; + if (!domains[0]?.name) throw new AppError("DOMAIN_REQUIRED", "Choose an email domain.", 400); + assertLoginEmailOutsideDomains( + input.ownerEmail, + domains.map((domain) => domain.name) + ); - const owner = await signUpOwnerUser(env, request, { - email: input.ownerEmail, - name: input.ownerName, - password: input.ownerPassword, - role: "owner" - }); + for (const domain of domains) { + await upsertMailDomain(env.DB, { + ...domain, + receivingStatus: "ready", + sendingStatus: "ready", + dnsStatus: "ready" + }); + } - await setPrimaryDomain(env.DB, domains[0].name); - if (input.portalHostname) { - await upsertWorkspaceHost(env.DB, { - hostname: input.portalHostname, - zoneId: - domains.find((domain) => input.portalHostname?.endsWith(`.${domain.name}`))?.zoneId ?? null, - kind: "portal", - canonical: true + await heartbeat.renew(); + const owner = await signUpOwnerUser(env, request, { + email: input.ownerEmail, + name: input.ownerName, + password: input.ownerPassword, + role: "owner" }); - } - await setChecklistAcknowledged(env.DB, input.checklistAcknowledged); + await heartbeat.renew(); - const mailboxes: Mailbox[] = []; - for (const mailbox of input.mailboxes) { - mailboxes.push(await createMailbox(env.DB, mailbox)); - } - const defaultFromMailbox = mailboxes.find( - (mailbox) => mailbox.address === input.defaultFromMailboxAddress - ); - if (!defaultFromMailbox) { - throw new AppError( - "DEFAULT_FROM_MAILBOX_REQUIRED", - "Choose one of the setup mailboxes as the default From mailbox.", - 400 + await setPrimaryDomain(env.DB, domains[0].name); + if (input.portalHostname) { + await upsertWorkspaceHost(env.DB, { + hostname: input.portalHostname, + zoneId: + domains.find((domain) => input.portalHostname?.endsWith(`.${domain.name}`))?.zoneId ?? + null, + kind: "portal", + canonical: true + }); + } + await setChecklistAcknowledged(env.DB, input.checklistAcknowledged); + + const mailboxes: Mailbox[] = []; + for (const mailbox of input.mailboxes) { + mailboxes.push(await createMailbox(env.DB, mailbox)); + } + const defaultFromMailbox = mailboxes.find( + (mailbox) => mailbox.address === input.defaultFromMailboxAddress ); - } - await setDefaultFromMailboxId(env.DB, owner.id, defaultFromMailbox.id); + if (!defaultFromMailbox) { + throw new AppError( + "DEFAULT_FROM_MAILBOX_REQUIRED", + "Choose one of the setup mailboxes as the default From mailbox.", + 400 + ); + } + await setDefaultFromMailboxId(env.DB, owner.id, defaultFromMailbox.id); - await completeSetupIfReady(env.DB); + await heartbeat.renew(); + await completeSetupIfReady(env.DB); - return { - owner, - mailboxes, - setup: await getSetupStatus(env.DB) - }; + return { + owner, + mailboxes, + setup: await getSetupStatus(env.DB) + }; + } finally { + await heartbeat.stop().catch(() => undefined); + await releaseBootstrapLock(env.DB, lock).catch(() => undefined); + } } export async function completeSetupIfReady(db: D1Database): Promise {