diff --git a/packages/access-policy/src/store.ts b/packages/access-policy/src/store.ts index 8f35bf09b..926534ee2 100644 --- a/packages/access-policy/src/store.ts +++ b/packages/access-policy/src/store.ts @@ -110,41 +110,63 @@ export function createDrizzleAccessPolicyStore< return row !== undefined; }, + // Read-modify-write under a transaction with a row lock rather than + // an optimistic version-stamp check: two transactions starting in + // the same wall-clock tick could still both pass a version check, + // which is exactly the lost-update bug this closes. + // + // Because this method's row may not exist yet (it is an upsert), + // `SELECT ... FOR UPDATE` alone has nothing to lock for a brand-new + // tenant. The `INSERT ... ON CONFLICT DO NOTHING` below guarantees + // a row first: two concurrent first-writes for the same tenant + // serialize on that insert's unique-index conflict — the loser + // blocks until the winner's transaction commits, then no-ops and + // its own subsequent `SELECT ... FOR UPDATE` sees the winner's + // committed row rather than racing it. async upsertPolicy(tenantId, patch) { - const current = await db - .select() - .from(policy) - .where(eq(policy.tenantId, tenantId)); - const existing = - current[0] === undefined - ? DEFAULT_ACCESS_POLICY - : resolveAccessPolicy(current[0]); - const next: AccessPolicy = { - selfSignup: patch.selfSignup ?? existing.selfSignup, - allowedDomains: patch.allowedDomains ?? existing.allowedDomains, - tenancyCreation: patch.tenancyCreation ?? existing.tenancyCreation, - }; - const now = new Date(); - await db - .insert(policy) - .values({ - tenantId, - selfSignup: next.selfSignup, - allowedDomains: serializeAllowedDomains(next.allowedDomains), - tenancyCreation: next.tenancyCreation, - createdAt: now, - updatedAt: now, - }) - .onConflictDoUpdate({ - target: policy.tenantId, - set: { + return db.transaction(async (tx) => { + const now = new Date(); + await tx + .insert(policy) + .values({ + tenantId, + selfSignup: DEFAULT_ACCESS_POLICY.selfSignup, + allowedDomains: serializeAllowedDomains( + DEFAULT_ACCESS_POLICY.allowedDomains, + ), + tenancyCreation: DEFAULT_ACCESS_POLICY.tenancyCreation, + createdAt: now, + updatedAt: now, + }) + .onConflictDoNothing({ target: policy.tenantId }); + + const [current] = await tx + .select() + .from(policy) + .where(eq(policy.tenantId, tenantId)) + .for("update"); + if (current === undefined) { + throw new Error( + `upsertPolicy: no access_policy.policy row for tenant ${tenantId} after ensuring one exists`, + ); + } + const existing = resolveAccessPolicy(current); + const next: AccessPolicy = { + selfSignup: patch.selfSignup ?? existing.selfSignup, + allowedDomains: patch.allowedDomains ?? existing.allowedDomains, + tenancyCreation: patch.tenancyCreation ?? existing.tenancyCreation, + }; + await tx + .update(policy) + .set({ selfSignup: next.selfSignup, allowedDomains: serializeAllowedDomains(next.allowedDomains), tenancyCreation: next.tenancyCreation, - updatedAt: now, - }, - }); - return next; + updatedAt: new Date(), + }) + .where(eq(policy.tenantId, tenantId)); + return next; + }); }, async createPendingInvite(tenantId, input) { diff --git a/packages/access-policy/test/store.drizzle.test.ts b/packages/access-policy/test/store.drizzle.test.ts index 54d25d962..3a24ee877 100644 --- a/packages/access-policy/test/store.drizzle.test.ts +++ b/packages/access-policy/test/store.drizzle.test.ts @@ -180,6 +180,104 @@ describeIfDb("createDrizzleAccessPolicyStore", () => { } }); + test("upsertPolicy is atomic: two concurrent patches to different fields on an existing row both land, neither reverts the other", async () => { + const setupSql = postgres(scratchUrl, { max: 1 }); + try { + const store = createDrizzleAccessPolicyStore(drizzle(setupSql)); + await store.upsertPolicy("tnt_race_existing", { + selfSignup: "off", + allowedDomains: [], + tenancyCreation: "owners", + }); + } finally { + await setupSql.end(); + } + + // A plain `Promise.all` of two real calls does not reliably force + // the worst-case interleaving on a fast local connection: one + // call's whole read-modify-write often finishes before the other's + // read even starts, so the two never actually overlap. Instead, a + // third connection takes the row's lock first and holds it open + // while both real `upsertPolicy` calls start and queue up behind + // it — releasing it then guarantees both calls' reads had to + // happen without seeing the other's write yet, exactly the + // interleaving that silently reverted one admin's change. + const holderSql = postgres(scratchUrl, { max: 1 }); + const sqlA = postgres(scratchUrl, { max: 1 }); + const sqlB = postgres(scratchUrl, { max: 1 }); + try { + const storeA = createDrizzleAccessPolicyStore(drizzle(sqlA)); + const storeB = createDrizzleAccessPolicyStore(drizzle(sqlB)); + + let holderReady: () => void = () => undefined; + const holderHasLock = new Promise((resolve) => { + holderReady = resolve; + }); + let releaseHolder: () => void = () => undefined; + const releaseSignal = new Promise((resolve) => { + releaseHolder = resolve; + }); + const holderTx = holderSql.begin(async (tx) => { + await tx`select * from access_policy.policy where tenant_id = 'tnt_race_existing' for update`; + holderReady(); + await releaseSignal; + }); + + await holderHasLock; + + const racers = Promise.all([ + storeA.upsertPolicy("tnt_race_existing", { selfSignup: "open" }), + storeB.upsertPolicy("tnt_race_existing", { + allowedDomains: ["acme.example"], + }), + ]); + // Give both calls time to actually issue their row-locking read + // and start queuing behind the holder before it releases. + await new Promise((resolve) => setTimeout(resolve, 100)); + + releaseHolder(); + await holderTx; + await racers; + + const policy = await storeA.getPolicy("tnt_race_existing"); + expect(policy.selfSignup).toBe("open"); + expect(policy.allowedDomains).toEqual(["acme.example"]); + expect(policy.tenancyCreation).toBe("owners"); + } finally { + await holderSql.end(); + await sqlA.end(); + await sqlB.end(); + } + }); + + test("upsertPolicy is atomic: two concurrent first-writes for a brand-new tenant both land, neither reverts the other", async () => { + const sql = postgres(scratchUrl, { max: 5 }); + try { + const store = createDrizzleAccessPolicyStore(drizzle(sql)); + + // No row exists yet for this tenant, so both calls race the + // create path too — the ensure-row-then-lock step inside + // `upsertPolicy` has to serialize this case as well, not only + // the existing-row case above. + await Promise.all([ + store.upsertPolicy("tnt_race_new", { selfSignup: "open" }), + store.upsertPolicy("tnt_race_new", { + allowedDomains: ["acme.example"], + }), + ]); + + const policy = await store.getPolicy("tnt_race_new"); + expect(policy.selfSignup).toBe("open"); + expect(policy.allowedDomains).toEqual(["acme.example"]); + + const rows = + await sql`select count(*)::int as count from access_policy.policy where tenant_id = 'tnt_race_new'`; + expect(rows[0]?.["count"]).toBe(1); + } finally { + await sql.end(); + } + }); + test("pending invites: a domain match is found for any email on that domain", async () => { const sql = postgres(scratchUrl, { max: 1 }); try {