Skip to content

Commit e9ac1db

Browse files
committed
address comments
1 parent af54aa6 commit e9ac1db

10 files changed

Lines changed: 430 additions & 41 deletions

apps/sim/lib/auth/auth.ts

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import * as schema from '@sim/db/schema'
88
import { createLogger } from '@sim/logger'
99
import { toError } from '@sim/utils/errors'
1010
import { generateId } from '@sim/utils/id'
11-
import { betterAuth, type User } from 'better-auth'
11+
import { type BetterAuthOptions, betterAuth, type User } from 'better-auth'
1212
import { drizzleAdapter } from 'better-auth/adapters/drizzle'
1313
import { APIError, createAuthMiddleware, getSessionFromCtx } from 'better-auth/api'
1414
import { nextCookies } from 'better-auth/next-js'
@@ -34,6 +34,7 @@ import {
3434
import { getAccessControlConfig, isEmailBlockedByAccessControl } from '@/lib/auth/access-control'
3535
import { createAnonymousSession, ensureAnonymousUserExists } from '@/lib/auth/anonymous'
3636
import { getRequestedSignInProviderId, isSignInProviderAllowed } from '@/lib/auth/constants'
37+
import { guardSubscriptionPlanWrites } from '@/lib/auth/stripe-adapter-guard'
3738
import { sendPlanWelcomeEmail } from '@/lib/billing'
3839
import {
3940
assertPersonalCheckoutAllowed,
@@ -197,10 +198,13 @@ export const auth = betterAuth({
197198
...(env.NEXT_PUBLIC_SOCKET_URL ? [env.NEXT_PUBLIC_SOCKET_URL] : []),
198199
...additionalTrustedOrigins,
199200
].filter(Boolean),
200-
database: drizzleAdapter(db, {
201-
provider: 'pg',
202-
schema,
203-
}),
201+
database: (options: BetterAuthOptions) =>
202+
guardSubscriptionPlanWrites(
203+
drizzleAdapter(db, {
204+
provider: 'pg',
205+
schema,
206+
})(options)
207+
),
204208
session: {
205209
cookieCache: {
206210
enabled: true,
@@ -3239,21 +3243,32 @@ export const auth = betterAuth({
32393243
throw orgError
32403244
}
32413245

3242-
await handleSubscriptionCreated(resolvedSubscription, event.id)
3243-
3244-
await syncSubscriptionUsageLimits(resolvedSubscription)
3245-
32463246
/**
32473247
* Transactional fence behind the personal-checkout admission
32483248
* guard: if the user joined a paid organization while their
32493249
* checkout was in flight, pause the fresh personal Pro at
32503250
* period end (same state a paid-org joiner's personal Pro
32513251
* enters; restored automatically if they leave the org).
3252+
*
3253+
* Runs BEFORE the free→paid transition handling: a personal
3254+
* subscription born covered is not a free→paid transition —
3255+
* the org plan keeps governing the user — so the usage reset
3256+
* (which would wipe org-attributed current-period usage) and
3257+
* its instrumentation must not run. Gated on `covered`, not
3258+
* `paused`, so event retries decide identically even when the
3259+
* join path already paused the subscription.
32523260
*/
3253-
if (isPro(resolvedSubscription.plan)) {
3254-
await pauseProSubscriptionForOrgCoverage(resolvedSubscription.referenceId)
3261+
const coveredByOrganization = isPro(resolvedSubscription.plan)
3262+
? (await pauseProSubscriptionForOrgCoverage(resolvedSubscription.referenceId))
3263+
.covered
3264+
: false
3265+
3266+
if (!coveredByOrganization) {
3267+
await handleSubscriptionCreated(resolvedSubscription, event.id)
32553268
}
32563269

3270+
await syncSubscriptionUsageLimits(resolvedSubscription)
3271+
32573272
await writeBillingInterval(resolvedSubscription.id, isAnnual ? 'year' : 'month')
32583273

32593274
await sendPlanWelcomeEmail(resolvedSubscription)
Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
5+
import { beforeEach, describe, expect, it, vi } from 'vitest'
6+
7+
vi.mock('@sim/db', () => dbChainMock)
8+
9+
import { guardSubscriptionPlanWrites } from '@/lib/auth/stripe-adapter-guard'
10+
11+
function createBaseAdapter() {
12+
return {
13+
create: vi.fn(async (data: { data: unknown }) => data.data),
14+
update: vi.fn(async (data: { update: unknown }) => data.update),
15+
updateMany: vi.fn(async () => 1),
16+
findOne: vi.fn(),
17+
}
18+
}
19+
20+
/**
21+
* The guard only relies on create/update/updateMany/findOne; the remaining
22+
* adapter surface passes through the spread untouched.
23+
*/
24+
// double-cast-allowed: test double implements only the adapter subset the guard touches
25+
const asAdapter = (base: ReturnType<typeof createBaseAdapter>) =>
26+
guardSubscriptionPlanWrites(base as unknown as Parameters<typeof guardSubscriptionPlanWrites>[0])
27+
28+
const ORG_ROW = { id: 'sub-1', referenceId: 'org-1', plan: 'team_6000' }
29+
const WHERE = [{ field: 'id', value: 'sub-1' }]
30+
31+
describe('guardSubscriptionPlanWrites', () => {
32+
beforeEach(() => {
33+
vi.clearAllMocks()
34+
resetDbChainMock()
35+
})
36+
37+
it('strips a non-org plan from an update targeting an org-referenced row but keeps the rest', async () => {
38+
const base = createBaseAdapter()
39+
base.findOne.mockResolvedValueOnce(ORG_ROW)
40+
// org existence lookup
41+
dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])
42+
43+
const guarded = asAdapter(base)
44+
await guarded.update({
45+
model: 'subscription',
46+
where: WHERE as never,
47+
update: { plan: 'pro_6000', status: 'active', seats: 2 },
48+
})
49+
50+
expect(base.update).toHaveBeenCalledWith(
51+
expect.objectContaining({ update: { status: 'active', seats: 2 } })
52+
)
53+
})
54+
55+
it('returns the current row without writing when stripping leaves an empty update', async () => {
56+
const base = createBaseAdapter()
57+
base.findOne.mockResolvedValueOnce(ORG_ROW).mockResolvedValueOnce(ORG_ROW)
58+
dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])
59+
60+
const guarded = asAdapter(base)
61+
const result = await guarded.update({
62+
model: 'subscription',
63+
where: WHERE as never,
64+
update: { plan: 'pro_6000' },
65+
})
66+
67+
expect(result).toEqual(ORG_ROW)
68+
expect(base.update).not.toHaveBeenCalled()
69+
})
70+
71+
it('passes through non-org plan updates for user-referenced rows', async () => {
72+
const base = createBaseAdapter()
73+
base.findOne.mockResolvedValueOnce({ id: 'sub-2', referenceId: 'user-1', plan: 'pro_6000' })
74+
// org existence lookup: user id is not an organization
75+
dbChainMockFns.limit.mockResolvedValueOnce([])
76+
77+
const guarded = asAdapter(base)
78+
await guarded.update({
79+
model: 'subscription',
80+
where: WHERE as never,
81+
update: { plan: 'pro_25000', status: 'active' },
82+
})
83+
84+
expect(base.update).toHaveBeenCalledWith(
85+
expect.objectContaining({ update: { plan: 'pro_25000', status: 'active' } })
86+
)
87+
})
88+
89+
it('passes through org-plan updates without any row lookup', async () => {
90+
const base = createBaseAdapter()
91+
92+
const guarded = asAdapter(base)
93+
await guarded.update({
94+
model: 'subscription',
95+
where: WHERE as never,
96+
update: { plan: 'team_25000', seats: 5 },
97+
})
98+
99+
expect(base.findOne).not.toHaveBeenCalled()
100+
expect(base.update).toHaveBeenCalledWith(
101+
expect.objectContaining({ update: { plan: 'team_25000', seats: 5 } })
102+
)
103+
})
104+
105+
it('passes through updates on other models untouched', async () => {
106+
const base = createBaseAdapter()
107+
108+
const guarded = asAdapter(base)
109+
await guarded.update({
110+
model: 'user',
111+
where: WHERE as never,
112+
update: { plan: 'pro_6000' },
113+
})
114+
115+
expect(base.findOne).not.toHaveBeenCalled()
116+
expect(base.update).toHaveBeenCalled()
117+
})
118+
119+
it('rejects creating an org-referenced subscription with a non-org plan', async () => {
120+
const base = createBaseAdapter()
121+
dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])
122+
123+
const guarded = asAdapter(base)
124+
await expect(
125+
guarded.create({
126+
model: 'subscription',
127+
data: { plan: 'pro_6000', referenceId: 'org-1' } as never,
128+
})
129+
).rejects.toThrow(/must hold a Team or Enterprise plan/)
130+
131+
expect(base.create).not.toHaveBeenCalled()
132+
})
133+
134+
it('allows creating a personal pro subscription', async () => {
135+
const base = createBaseAdapter()
136+
dbChainMockFns.limit.mockResolvedValueOnce([])
137+
138+
const guarded = asAdapter(base)
139+
await guarded.create({
140+
model: 'subscription',
141+
data: { plan: 'pro_6000', referenceId: 'user-1' } as never,
142+
})
143+
144+
expect(base.create).toHaveBeenCalled()
145+
})
146+
})
Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
import { db } from '@sim/db'
2+
import { organization } from '@sim/db/schema'
3+
import { createLogger } from '@sim/logger'
4+
import type { drizzleAdapter } from 'better-auth/adapters/drizzle'
5+
import { eq } from 'drizzle-orm'
6+
import { isOrgPlan } from '@/lib/billing/plan-helpers'
7+
8+
const logger = createLogger('StripeAdapterGuard')
9+
10+
type BetterAuthAdapter = ReturnType<ReturnType<typeof drizzleAdapter>>
11+
12+
type SubscriptionWriteSurface = Pick<
13+
BetterAuthAdapter,
14+
'create' | 'update' | 'updateMany' | 'findOne' | 'findMany'
15+
>
16+
17+
/**
18+
* The Better Auth Stripe plugin persists webhook state through the raw
19+
* database adapter BEFORE invoking our subscription callbacks — including the
20+
* `plan` column resolved from the Stripe price. That makes the adapter the
21+
* only in-process seam that can enforce the billing invariant that
22+
* organization-referenced subscriptions hold Team/Enterprise plans: by the
23+
* time `syncSubscriptionPlan` runs in a callback, the plugin's write has
24+
* already landed.
25+
*
26+
* Checkout admission blocks user-driven violations; this guard blocks the
27+
* remaining vector — an operator swapping an org subscription onto a personal
28+
* price in the Stripe dashboard. It never repairs state: an invalid `plan`
29+
* write is refused (update: the field is stripped so status/period/seat sync
30+
* still lands; create: the whole insert is rejected) and an error is logged
31+
* so operators fix the price in Stripe.
32+
*
33+
* Transactions are wrapped recursively so writes inside
34+
* `adapter.transaction(...)` callbacks go through the same guard.
35+
*/
36+
export function guardSubscriptionPlanWrites(adapter: BetterAuthAdapter): BetterAuthAdapter {
37+
const guarded: BetterAuthAdapter = {
38+
...adapter,
39+
...guardWriteSurface(adapter),
40+
}
41+
42+
const transaction = adapter.transaction
43+
if (typeof transaction === 'function') {
44+
guarded.transaction = (callback) =>
45+
transaction((trx) => callback({ ...trx, ...guardWriteSurface(trx) }))
46+
}
47+
48+
return guarded
49+
}
50+
51+
function guardWriteSurface<TAdapter extends SubscriptionWriteSurface>(
52+
adapter: TAdapter
53+
): SubscriptionWriteSurface {
54+
return {
55+
findOne: adapter.findOne,
56+
findMany: adapter.findMany,
57+
create: async (data) => {
58+
if (data.model === 'subscription') {
59+
const values = data.data as Record<string, unknown>
60+
const plan = values.plan
61+
const referenceId = values.referenceId
62+
if (
63+
typeof plan === 'string' &&
64+
!isOrgPlan(plan) &&
65+
typeof referenceId === 'string' &&
66+
(await isOrganizationReference(referenceId))
67+
) {
68+
logger.error(
69+
'Blocked creating an organization-referenced subscription with a non-org plan — fix the plan or reference in Stripe',
70+
{ referenceId, rejectedPlan: plan }
71+
)
72+
throw new Error(
73+
`Organization-referenced subscriptions must hold a Team or Enterprise plan (got '${plan}')`
74+
)
75+
}
76+
}
77+
return adapter.create(data)
78+
},
79+
update: async (data) => {
80+
if (data.model === 'subscription' && hasNonOrgPlanWrite(data.update)) {
81+
const row = await adapter.findOne<SubscriptionRowSlice>({
82+
model: 'subscription',
83+
where: data.where,
84+
})
85+
const sanitized = await stripPlanWhenOrgReferenced(
86+
row ? [row] : [],
87+
data.update as Record<string, unknown>
88+
)
89+
if (sanitized.blockedAll) return row as never
90+
return adapter.update({ ...data, update: sanitized.update as never })
91+
}
92+
return adapter.update(data)
93+
},
94+
updateMany: async (data) => {
95+
if (data.model === 'subscription' && hasNonOrgPlanWrite(data.update)) {
96+
const rows = await adapter.findMany<SubscriptionRowSlice>({
97+
model: 'subscription',
98+
where: data.where,
99+
})
100+
const sanitized = await stripPlanWhenOrgReferenced(rows, data.update)
101+
if (sanitized.blockedAll) return 0
102+
return adapter.updateMany({ ...data, update: sanitized.update })
103+
}
104+
return adapter.updateMany(data)
105+
},
106+
}
107+
}
108+
109+
interface SubscriptionRowSlice {
110+
id: string
111+
referenceId: string
112+
plan: string
113+
}
114+
115+
function hasNonOrgPlanWrite(update: unknown): boolean {
116+
if (!update || typeof update !== 'object' || !('plan' in update)) return false
117+
const plan = (update as { plan: unknown }).plan
118+
return typeof plan === 'string' && !isOrgPlan(plan)
119+
}
120+
121+
async function isOrganizationReference(referenceId: string): Promise<boolean> {
122+
const [referencedOrganization] = await db
123+
.select({ id: organization.id })
124+
.from(organization)
125+
.where(eq(organization.id, referenceId))
126+
.limit(1)
127+
return Boolean(referencedOrganization)
128+
}
129+
130+
interface SanitizedSubscriptionUpdate {
131+
update: Record<string, unknown>
132+
/** True when stripping the invalid plan left nothing to write. */
133+
blockedAll: boolean
134+
}
135+
136+
/**
137+
* Strip a non-org `plan` (and its derived `limits`) from an update when ANY
138+
* targeted row is organization-referenced. The remaining fields (status,
139+
* periods, seats, cancellation state) are legitimate Stripe state and still
140+
* sync. Evaluating every targeted row keeps multi-row updates safe: a mixed
141+
* personal/org target set must not leak the invalid plan onto the org rows.
142+
*/
143+
async function stripPlanWhenOrgReferenced(
144+
rows: SubscriptionRowSlice[],
145+
update: Record<string, unknown>
146+
): Promise<SanitizedSubscriptionUpdate> {
147+
let organizationRow: SubscriptionRowSlice | null = null
148+
for (const row of rows) {
149+
if (await isOrganizationReference(row.referenceId)) {
150+
organizationRow = row
151+
break
152+
}
153+
}
154+
if (!organizationRow) {
155+
return { update, blockedAll: false }
156+
}
157+
158+
logger.error(
159+
'Blocked writing a non-org plan onto an organization-referenced subscription — fix the price in Stripe',
160+
{
161+
subscriptionId: organizationRow.id,
162+
organizationId: organizationRow.referenceId,
163+
currentPlan: organizationRow.plan,
164+
rejectedPlan: update.plan,
165+
}
166+
)
167+
168+
const { plan: _plan, limits: _limits, ...rest } = update
169+
return { update: rest, blockedAll: Object.keys(rest).length === 0 }
170+
}

0 commit comments

Comments
 (0)