Skip to content

Commit bc828fb

Browse files
committed
fix(billing): withhold the free allowance from stats rows created outside signup
A storage write or an ownership transfer that creates a missing user_stats row gave it the default free allowance without the deleted-identity check ensureUserStatsExists makes. When either insert creates the row, the same email-only check now runs on the caller's transaction, in a savepoint that fails open; an existing row costs no extra query.
1 parent 0455b02 commit bc828fb

5 files changed

Lines changed: 108 additions & 22 deletions

File tree

‎apps/sim/lib/billing/core/usage.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -702,6 +702,36 @@ async function forfeitFreeTierAllowance(userId: string): Promise<void> {
702702
logger.info('Withheld the free allowance from a returning deleted-account identity', { userId })
703703
}
704704

705+
/**
706+
* Withholds the free allowance from a `user_stats` row the caller just created
707+
* in `tx`, when the user's email belongs to an account deleted within the
708+
* retention window: the email-only check {@link ensureUserStatsExists} makes for
709+
* its fallback row, on the caller's transaction so the row never commits
710+
* holding an allowance it should not. Call it only when the insert created the
711+
* row, so an existing row costs nothing. Runs in a savepoint and fails open: a
712+
* failure leaves the row eligible and the caller's transaction intact.
713+
*/
714+
export async function withholdFreeAllowanceFromCreatedStatsRow(
715+
tx: DbOrTx,
716+
userId: string
717+
): Promise<void> {
718+
if (!(await userHasDeletedAccountIdentity(userId, { includeSignIns: false }, tx))) return
719+
try {
720+
await tx.transaction(async (savepoint) => {
721+
await savepoint
722+
.update(userStats)
723+
.set({ freeTierEligible: false })
724+
.where(eq(userStats.userId, userId))
725+
await syncUsageLimitsFromSubscription(userId, savepoint)
726+
})
727+
logger.info('Withheld the free allowance from a returning deleted-account identity', {
728+
userId,
729+
})
730+
} catch (error) {
731+
logger.error('Failed to withhold the free allowance from a new stats row', { userId, error })
732+
}
733+
}
734+
705735
/**
706736
* Syncs the user's personal billing pool from their exact personal subscription.
707737
* Organization subscriptions have a separate pool and never clear personal limits.

‎apps/sim/lib/billing/organizations/membership.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,10 @@ import {
2929
invalidateSecurityPolicyVersionCache,
3030
} from '@/lib/auth/security-policy'
3131
import { applySessionPolicyToNewMember } from '@/lib/auth/session-policy'
32-
import { syncUsageLimitsFromSubscription } from '@/lib/billing/core/usage'
32+
import {
33+
syncUsageLimitsFromSubscription,
34+
withholdFreeAllowanceFromCreatedStatsRow,
35+
} from '@/lib/billing/core/usage'
3336
import {
3437
assertNoUnresolvedEnterpriseIssuance,
3538
resolveEnterpriseMetadataIntent,
@@ -1935,14 +1938,18 @@ export async function transferOrganizationOwnership(
19351938
.limit(1)
19361939

19371940
if (oldStats) {
1938-
await tx
1941+
const createdStats = await tx
19391942
.insert(userStats)
19401943
.values({
19411944
id: generateId(),
19421945
userId: newOwnerUserId,
19431946
usageLimitUpdatedAt: new Date(),
19441947
})
19451948
.onConflictDoNothing({ target: userStats.userId })
1949+
.returning({ id: userStats.id })
1950+
if (createdStats.length > 0) {
1951+
await withholdFreeAllowanceFromCreatedStatsRow(tx, newOwnerUserId)
1952+
}
19461953

19471954
const overage = oldStats.billedOverageThisPeriod || '0'
19481955
const overageNum = toNumber(toDecimal(overage))

‎apps/sim/lib/billing/storage/tracking.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import { isRecordLike } from '@sim/utils/object'
2424
import { and, asc, eq, inArray, sql } from 'drizzle-orm'
2525
import { maybeNotifyLimit } from '@/lib/billing/core/limit-notifications'
2626
import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
27+
import { withholdFreeAllowanceFromCreatedStatsRow } from '@/lib/billing/core/usage'
2728
import type { BillingEntity } from '@/lib/billing/core/usage-log'
2829
import type {
2930
ProjectStorageBillingContext,
@@ -771,7 +772,8 @@ export async function incrementAdmittedStorageUsageForBillingContextInTx(
771772
* matches 0 rows, and therefore reads as "quota exceeded", if that row
772773
* doesn't exist yet. `ensureUserStatsExists` itself isn't reused here since
773774
* it writes through the standalone `db` client, which would open a second
774-
* pooled connection while this transaction's is held.
775+
* pooled connection while this transaction's is held; a row created here gets
776+
* its deleted-identity check on `tx` instead.
775777
*/
776778
export async function checkAndIncrementStorageUsageInTx(
777779
tx: DbOrTx,
@@ -792,7 +794,7 @@ export async function checkAndIncrementStorageUsageInTx(
792794
const orgScoped = isOrgScopedSubscription(sub, userId) && sub
793795

794796
if (!orgScoped) {
795-
await tx
797+
const created = await tx
796798
.insert(userStats)
797799
.values({
798800
id: generateId(),
@@ -801,6 +803,8 @@ export async function checkAndIncrementStorageUsageInTx(
801803
usageLimitUpdatedAt: new Date(),
802804
})
803805
.onConflictDoNothing({ target: userStats.userId })
806+
.returning({ id: userStats.id })
807+
if (created.length > 0) await withholdFreeAllowanceFromCreatedStatsRow(tx, userId)
804808
}
805809

806810
const [updated] = orgScoped

‎apps/sim/lib/users/deleted-identity.integration.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
handleNewUser,
1818
syncUsageLimitsFromSubscription,
1919
} from '@/lib/billing/core/usage'
20+
import { checkAndIncrementStorageUsageInTx } from '@/lib/billing/storage/tracking'
2021
import { getFreeTierLimit } from '@/lib/billing/subscriptions/utils'
2122
import { deleteUserAccount } from '@/lib/users/account-deletion'
2223
import { purgeExpiredDeletedAccountIdentities } from '@/lib/users/deleted-identity'
@@ -250,6 +251,28 @@ describe('deleted account identities and the free allowance', () => {
250251
expect(await statsOf(stranger)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() })
251252
})
252253

254+
it('withholds the allowance from a stats row a storage write creates', async () => {
255+
const email = `storage.${generateId()}@deleted-identity.test`
256+
await deleteOwnAccount(await signUp(email))
257+
const returning = await insertUser(email)
258+
const stranger = await insertUser(`storage-stranger.${generateId()}@deleted-identity.test`)
259+
const store = (userId: string) =>
260+
db.transaction((tx) => checkAndIncrementStorageUsageInTx(tx, null, userId, 1024))
261+
262+
const first = await store(returning)
263+
const second = await store(returning)
264+
await store(stranger)
265+
266+
expect([first.allowed, second.allowed]).toEqual([true, true])
267+
expect(await statsOf(returning)).toEqual({ freeTierEligible: false, limit: 0 })
268+
expect(await statsOf(stranger)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() })
269+
const [stored] = await db
270+
.select({ bytes: userStats.storageUsedBytes })
271+
.from(userStats)
272+
.where(eq(userStats.userId, returning))
273+
expect(stored.bytes).toBe(2048)
274+
})
275+
253276
it('never withholds from a stats-less user’s fallback row for a sign-in linked in a later request', async () => {
254277
const microsoft = { providerId: 'microsoft', accountId: `ms-${generateId()}` }
255278
await deleteOwnAccount(await signUp(`ms-gone.${generateId()}@deleted-identity.test`, microsoft))

‎apps/sim/lib/users/deleted-identity.ts‎

Lines changed: 40 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -214,27 +214,45 @@ export async function recordDeletedAccountIdentities(tx: DbOrTx, userId: string)
214214
}
215215
}
216216

217+
async function findDeletedIdentityMatch(
218+
executor: DbOrTx,
219+
identity: SignInIdentity
220+
): Promise<boolean> {
221+
const hashes = identityHashes(identity)
222+
if (hashes.length === 0) return false
223+
const [match] = await executor
224+
.select({ identityHash: deletedAccountIdentity.identityHash })
225+
.from(deletedAccountIdentity)
226+
.where(
227+
and(
228+
inArray(deletedAccountIdentity.identityHash, hashes),
229+
gt(deletedAccountIdentity.expiresAt, sql`now()`)
230+
)
231+
)
232+
.limit(1)
233+
return Boolean(match)
234+
}
235+
236+
/**
237+
* Runs a lookup on `executor`. Inside a caller's transaction it runs in a savepoint, so a failed
238+
* lookup rolls back only itself instead of aborting the caller's transaction.
239+
*/
240+
function readIdentity<T>(executor: DbOrTx, read: (reader: DbOrTx) => Promise<T>): Promise<T> {
241+
return executor === db ? read(db) : executor.transaction(read)
242+
}
243+
217244
/**
218245
* Whether any of these identities belonged to an account deleted within the
219246
* retention window. Fails open: a lookup error returns `false` and is logged,
220247
* because refusing a legitimate person's allowance is worse than one missed match.
221248
*/
222-
export async function isDeletedAccountIdentity(identity: SignInIdentity): Promise<boolean> {
249+
export async function isDeletedAccountIdentity(
250+
identity: SignInIdentity,
251+
executor: DbOrTx = db
252+
): Promise<boolean> {
223253
if (!isBillingEnabled) return false
224254
try {
225-
const hashes = identityHashes(identity)
226-
if (hashes.length === 0) return false
227-
const [match] = await db
228-
.select({ identityHash: deletedAccountIdentity.identityHash })
229-
.from(deletedAccountIdentity)
230-
.where(
231-
and(
232-
inArray(deletedAccountIdentity.identityHash, hashes),
233-
gt(deletedAccountIdentity.expiresAt, sql`now()`)
234-
)
235-
)
236-
.limit(1)
237-
return Boolean(match)
255+
return await readIdentity(executor, (reader) => findDeletedIdentityMatch(reader, identity))
238256
} catch (error) {
239257
logger.error('Deleted-identity lookup failed; granting the free allowance', { error })
240258
return false
@@ -246,16 +264,20 @@ export async function isDeletedAccountIdentity(identity: SignInIdentity): Promis
246264
* the retention window: always its email, and its linked social sign-ins only
247265
* when `includeSignIns` is set, which callers do only for a user created in the
248266
* current signup request (a sign-in linked later never withholds anything).
249-
* Fails open like {@link isDeletedAccountIdentity}.
267+
* Pass the caller's transaction as `executor` to read on it. Fails open like
268+
* {@link isDeletedAccountIdentity}.
250269
*/
251270
export async function userHasDeletedAccountIdentity(
252271
userId: string,
253-
{ includeSignIns }: { includeSignIns: boolean }
272+
{ includeSignIns }: { includeSignIns: boolean },
273+
executor: DbOrTx = db
254274
): Promise<boolean> {
255275
if (!isBillingEnabled) return false
256276
try {
257-
const identity = await loadSignInIdentity(db, userId)
258-
return await isDeletedAccountIdentity(includeSignIns ? identity : { email: identity.email })
277+
return await readIdentity(executor, async (reader) => {
278+
const identity = await loadSignInIdentity(reader, userId)
279+
return findDeletedIdentityMatch(reader, includeSignIns ? identity : { email: identity.email })
280+
})
259281
} catch (error) {
260282
logger.error('Deleted-identity lookup failed; granting the free allowance', { userId, error })
261283
return false

0 commit comments

Comments
 (0)