Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 92 additions & 28 deletions apps/sim/lib/auth/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@ import * as schema from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { betterAuth, type User } from 'better-auth'
import { type BetterAuthOptions, betterAuth, type User } from 'better-auth'
import { drizzleAdapter } from 'better-auth/adapters/drizzle'
import { APIError, createAuthMiddleware } from 'better-auth/api'
import { APIError, createAuthMiddleware, getSessionFromCtx } from 'better-auth/api'
import { nextCookies } from 'better-auth/next-js'
import {
admin,
Expand All @@ -34,8 +34,13 @@ import {
import { getAccessControlConfig, isEmailBlockedByAccessControl } from '@/lib/auth/access-control'
import { createAnonymousSession, ensureAnonymousUserExists } from '@/lib/auth/anonymous'
import { getRequestedSignInProviderId, isSignInProviderAllowed } from '@/lib/auth/constants'
import { guardSubscriptionPlanWrites } from '@/lib/auth/stripe-adapter-guard'
import { sendPlanWelcomeEmail } from '@/lib/billing'
import { authorizeSubscriptionReference } from '@/lib/billing/authorization'
import {
assertPersonalCheckoutAllowed,
authorizeSubscriptionReference,
isPersonalCheckoutRequest,
} from '@/lib/billing/authorization'
import {
getOrganizationIdForSubscriptionReference,
syncSubscriptionPlan,
Expand All @@ -46,7 +51,8 @@ import {
ensureOrganizationForTeamSubscription,
syncSubscriptionUsageLimits,
} from '@/lib/billing/organization'
import { isTeam } from '@/lib/billing/plan-helpers'
import { pauseProSubscriptionForOrgCoverage } from '@/lib/billing/organizations/membership'
import { isPro, isTeam } from '@/lib/billing/plan-helpers'
import { getPlans, resolvePlanFromStripeSubscription } from '@/lib/billing/plans'
import { syncSeatsFromStripeQuantity } from '@/lib/billing/validation/seat-management'
import { handleAbandonedCheckout } from '@/lib/billing/webhooks/checkout'
Expand All @@ -58,7 +64,6 @@ import {
handleInvoicePaymentSucceeded,
} from '@/lib/billing/webhooks/invoices'
import {
handleOrganizationPlanDowngrade,
handleSubscriptionCreated,
handleSubscriptionDeleted,
} from '@/lib/billing/webhooks/subscription'
Expand Down Expand Up @@ -193,10 +198,13 @@ export const auth = betterAuth({
...(env.NEXT_PUBLIC_SOCKET_URL ? [env.NEXT_PUBLIC_SOCKET_URL] : []),
...additionalTrustedOrigins,
].filter(Boolean),
database: drizzleAdapter(db, {
provider: 'pg',
schema,
}),
database: (options: BetterAuthOptions) =>
guardSubscriptionPlanWrites(
drizzleAdapter(db, {
provider: 'pg',
schema,
})(options)
),
session: {
cookieCache: {
enabled: true,
Expand Down Expand Up @@ -903,6 +911,21 @@ export const auth = betterAuth({
}
}

/**
* Personal checkout guard. The Stripe plugin's `authorizeReference`
* only runs for organization references (it skips references equal to
* the session user), so duplicate-coverage enforcement for personal
* checkouts lives here: a member of an org with an entitled paid
* subscription must not buy a personal plan on top of it.
*/
if (isBillingEnabled && ctx.path === '/subscription/upgrade') {
const session = await getSessionFromCtx(ctx)
const sessionUserId = session?.user?.id
if (sessionUserId && isPersonalCheckoutRequest(ctx.body ?? {}, sessionUserId)) {
await assertPersonalCheckoutAllowed(sessionUserId)
}
}

return
}),
},
Expand Down Expand Up @@ -3140,8 +3163,21 @@ export const auth = betterAuth({
subscription: {
enabled: true,
plans: getPlans(),
authorizeReference: async ({ user, referenceId, action }) => {
return await authorizeSubscriptionReference(user.id, referenceId, action)
authorizeReference: async ({ user, referenceId, action }, ctx) => {
const body: unknown = ctx?.body
const requestedPlan =
typeof body === 'object' &&
body !== null &&
'plan' in body &&
typeof body.plan === 'string'
? body.plan
: undefined
return await authorizeSubscriptionReference(
user.id,
referenceId,
action,
requestedPlan
)
},
getCheckoutSessionParams: async () => ({
params: { allow_promotion_codes: true },
Expand Down Expand Up @@ -3175,11 +3211,16 @@ export const auth = betterAuth({
)
}

await syncSubscriptionPlan(subscription.id, subscription.plan, planFromStripe)
const syncedPlan = await syncSubscriptionPlan(
subscription.id,
subscription.plan,
planFromStripe,
subscription.referenceId
)

const subscriptionForOrg = {
...subscription,
plan: planFromStripe ?? subscription.plan,
plan: syncedPlan ?? subscription.plan,
enterpriseOperationId: stripeSubscription.metadata?.enterpriseOperationId ?? null,
}

Expand All @@ -3202,7 +3243,29 @@ export const auth = betterAuth({
throw orgError
}

await handleSubscriptionCreated(resolvedSubscription, event.id)
/**
* Transactional fence behind the personal-checkout admission
* guard: if the user joined a paid organization while their
* checkout was in flight, pause the fresh personal Pro at
* period end (same state a paid-org joiner's personal Pro
* enters; restored automatically if they leave the org).
*
* Runs BEFORE the free→paid transition handling: a personal
* subscription born covered is not a free→paid transition —
* the org plan keeps governing the user — so the usage reset
* (which would wipe org-attributed current-period usage) and
* its instrumentation must not run. Gated on `covered`, not
* `paused`, so event retries decide identically even when the
* join path already paused the subscription.
*/
const coveredByOrganization = isPro(resolvedSubscription.plan)
? (await pauseProSubscriptionForOrgCoverage(resolvedSubscription.referenceId))
.covered
: false

if (!coveredByOrganization) {
await handleSubscriptionCreated(resolvedSubscription, event.id)
}
Comment thread
icecrasher321 marked this conversation as resolved.

await syncSubscriptionUsageLimits(resolvedSubscription)

Expand Down Expand Up @@ -3238,8 +3301,6 @@ export const auth = betterAuth({
const isUpgradeToTeam =
isTeamPlan && !isTeam(subscription.plan) && referenceOrganizationId == null

const effectivePlanForTeamFeatures = planFromStripe ?? subscription.plan

logger.info('[onSubscriptionUpdate] Subscription updated', {
subscriptionId: subscription.id,
status: subscription.status,
Expand All @@ -3258,11 +3319,24 @@ export const auth = betterAuth({
)
}

await syncSubscriptionPlan(subscription.id, subscription.plan, planFromStripe)
const syncedPlan = await syncSubscriptionPlan(
subscription.id,
subscription.plan,
planFromStripe,
subscription.referenceId
)
Comment thread
icecrasher321 marked this conversation as resolved.

/**
* All downstream processing keys off the plan the DB actually
* holds after the sync — a plan write refused by the org/plan
* invariant must not leak the rejected Stripe plan into org
* resolution, seat sync, or usage limits.
*/
const effectivePlanForTeamFeatures = syncedPlan ?? subscription.plan

const subscriptionForOrg = {
...subscription,
plan: planFromStripe ?? subscription.plan,
plan: effectivePlanForTeamFeatures,
enterpriseOperationId: stripeSubscription.metadata?.enterpriseOperationId ?? null,
}

Expand Down Expand Up @@ -3333,16 +3407,6 @@ export const auth = betterAuth({
error,
})
}
} else {
await handleOrganizationPlanDowngrade(
{
id: resolvedSubscription.id,
plan: effectivePlanForTeamFeatures ?? null,
referenceId: resolvedSubscription.referenceId,
status: resolvedSubscription.status ?? null,
},
event.id
)
}

await writeBillingInterval(resolvedSubscription.id, isAnnual ? 'year' : 'month')
Expand Down
177 changes: 177 additions & 0 deletions apps/sim/lib/auth/stripe-adapter-guard.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
/**
* @vitest-environment node
*/
import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'

vi.mock('@sim/db', () => dbChainMock)

import { guardSubscriptionPlanWrites } from '@/lib/auth/stripe-adapter-guard'

function createBaseAdapter() {
return {
create: vi.fn(async (data: { data: unknown }) => data.data),
update: vi.fn(async (data: { update: unknown }) => data.update),
updateMany: vi.fn(async () => 1),
findOne: vi.fn(),
}
}

/**
* The guard only relies on create/update/updateMany/findOne; the remaining
* adapter surface passes through the spread untouched.
*/
// double-cast-allowed: test double implements only the adapter subset the guard touches
const asAdapter = (base: ReturnType<typeof createBaseAdapter>) =>
guardSubscriptionPlanWrites(base as unknown as Parameters<typeof guardSubscriptionPlanWrites>[0])

const ORG_ROW = { id: 'sub-1', referenceId: 'org-1', plan: 'team_6000' }
const WHERE = [{ field: 'id', value: 'sub-1' }]

describe('guardSubscriptionPlanWrites', () => {
beforeEach(() => {
vi.clearAllMocks()
resetDbChainMock()
})

it('strips a non-org plan from an update targeting an org-referenced row but keeps the rest', async () => {
const base = createBaseAdapter()
base.findOne.mockResolvedValueOnce(ORG_ROW)
// org existence lookup
dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])

const guarded = asAdapter(base)
await guarded.update({
model: 'subscription',
where: WHERE as never,
update: { plan: 'pro_6000', status: 'active', seats: 2 },
})

expect(base.update).toHaveBeenCalledWith(
expect.objectContaining({ update: { status: 'active', seats: 2 } })
)
})

it('returns the current row without writing when stripping leaves an empty update', async () => {
const base = createBaseAdapter()
base.findOne.mockResolvedValueOnce(ORG_ROW).mockResolvedValueOnce(ORG_ROW)
dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])

const guarded = asAdapter(base)
const result = await guarded.update({
model: 'subscription',
where: WHERE as never,
update: { plan: 'pro_6000' },
})

expect(result).toEqual(ORG_ROW)
expect(base.update).not.toHaveBeenCalled()
})

it('strips the plan when the pre-update lookup finds no row (fail closed)', async () => {
const base = createBaseAdapter()
base.findOne.mockResolvedValueOnce(null)

const guarded = asAdapter(base)
await guarded.update({
model: 'subscription',
where: WHERE as never,
update: { plan: 'pro_6000', status: 'active' },
})

expect(base.update).toHaveBeenCalledWith(
expect.objectContaining({ update: { status: 'active' } })
)
})

it('writes nothing when a plan-only update targets no readable row', async () => {
const base = createBaseAdapter()
base.findOne.mockResolvedValueOnce(null)

const guarded = asAdapter(base)
const result = await guarded.update({
model: 'subscription',
where: WHERE as never,
update: { plan: 'pro_6000' },
})

expect(result).toBeNull()
expect(base.update).not.toHaveBeenCalled()
})

it('passes through non-org plan updates for user-referenced rows', async () => {
const base = createBaseAdapter()
base.findOne.mockResolvedValueOnce({ id: 'sub-2', referenceId: 'user-1', plan: 'pro_6000' })
// org existence lookup: user id is not an organization
dbChainMockFns.limit.mockResolvedValueOnce([])

const guarded = asAdapter(base)
await guarded.update({
model: 'subscription',
where: WHERE as never,
update: { plan: 'pro_25000', status: 'active' },
})

expect(base.update).toHaveBeenCalledWith(
expect.objectContaining({ update: { plan: 'pro_25000', status: 'active' } })
)
})

it('passes through org-plan updates without any row lookup', async () => {
const base = createBaseAdapter()

const guarded = asAdapter(base)
await guarded.update({
model: 'subscription',
where: WHERE as never,
update: { plan: 'team_25000', seats: 5 },
})

expect(base.findOne).not.toHaveBeenCalled()
expect(base.update).toHaveBeenCalledWith(
expect.objectContaining({ update: { plan: 'team_25000', seats: 5 } })
)
})

it('passes through updates on other models untouched', async () => {
const base = createBaseAdapter()

const guarded = asAdapter(base)
await guarded.update({
model: 'user',
where: WHERE as never,
update: { plan: 'pro_6000' },
})

expect(base.findOne).not.toHaveBeenCalled()
expect(base.update).toHaveBeenCalled()
})

it('rejects creating an org-referenced subscription with a non-org plan', async () => {
const base = createBaseAdapter()
dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])

const guarded = asAdapter(base)
await expect(
guarded.create({
model: 'subscription',
data: { plan: 'pro_6000', referenceId: 'org-1' } as never,
})
).rejects.toThrow(/must hold a Team or Enterprise plan/)

expect(base.create).not.toHaveBeenCalled()
})

it('allows creating a personal pro subscription', async () => {
const base = createBaseAdapter()
dbChainMockFns.limit.mockResolvedValueOnce([])

const guarded = asAdapter(base)
await guarded.create({
model: 'subscription',
data: { plan: 'pro_6000', referenceId: 'user-1' } as never,
})

expect(base.create).toHaveBeenCalled()
})
})
Loading
Loading