Skip to content

Commit c48df42

Browse files
committed
improvement(aws-marketplace): close billing gaps and simplify the license write paths
- Refuse linking a Marketplace license to an organization already paying Stripe for a self-serve plan, so it is never billed twice - Count only listings this deployment sells when blocking Stripe checkout - Never add to a usage record after it is frozen for sending, even under a lock taken for another organization - Container license: require a grant the account may use, and stagger replica restarts on a license change - CloudFormation prerequisites: a fixed URL-safe database password instead of RDS's rotating one; placeholder product ID in the example values - One insert path, one claim invariant, and one once-only notice helper in the repository; shared claim lookup, lock recheck, and audit entry in the use cases
1 parent cff97b2 commit c48df42

14 files changed

Lines changed: 421 additions & 231 deletions

‎apps/sim/lib/billing/authorization.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { subscription } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
44
import { APIError } from 'better-auth/api'
55
import { and, eq, isNotNull } from 'drizzle-orm'
6-
import { isAwsMarketplaceEnabled } from '@/lib/billing/aws-marketplace/config'
6+
import { getAwsMarketplaceProducts } from '@/lib/billing/aws-marketplace/config'
77
import { holdsMarketplaceLicense } from '@/lib/billing/aws-marketplace/repository'
88
import { isOrganizationOwnerOrAdmin } from '@/lib/billing/core/organization'
99
import {
@@ -64,8 +64,9 @@ async function assertNoBoundIncompleteSubscription(referenceId: string): Promise
6464
* license lookup fails.
6565
*/
6666
async function assertNotBilledThroughMarketplace(referenceId: string): Promise<void> {
67-
if (!isAwsMarketplaceEnabled()) return
68-
if (!(await holdsMarketplaceLicense(referenceId))) return
67+
const products = getAwsMarketplaceProducts()
68+
if (products.size === 0) return
69+
if (!(await holdsMarketplaceLicense(products, referenceId))) return
6970

7071
logger.warn('Blocking checkout - reference is billed through AWS Marketplace', { referenceId })
7172
throw new APIError('CONFLICT', {

‎apps/sim/lib/billing/aws-marketplace/application/admin.ts‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
1+
import { AuditAction, recordAudit } from '@sim/audit'
22
import { db } from '@sim/db'
33
import type { AdminMutationActor } from '@/lib/admin/dashboard'
44
import {
5-
licenseAuditMetadata,
5+
licenseAuditEntry,
66
requireLinkableTo,
7+
requireNotBilledThroughStripe,
78
} from '@/lib/billing/aws-marketplace/application/use-cases'
89
import {
910
countUnlinkedLicenses,
@@ -57,6 +58,7 @@ export async function linkAwsMarketplaceLicenseAsAdmin(
5758
if (locked.organizationId === input.organizationId) {
5859
return { license: locked, newlyLinked: false }
5960
}
61+
await requireNotBilledThroughStripe(tx, input.organizationId)
6062
return {
6163
license: await linkLicenseForTeam(tx, input.licenseId, input.organizationId),
6264
newlyLinked: true,
@@ -67,12 +69,12 @@ export async function linkAwsMarketplaceLicenseAsAdmin(
6769
actorId: actor.id,
6870
actorName: actor.name,
6971
actorEmail: actor.email,
70-
action: AuditAction.AWS_MARKETPLACE_LICENSE_LINKED,
71-
resourceType: AuditResourceType.AWS_MARKETPLACE_LICENSE,
72-
resourceId: license.id,
73-
resourceName: license.dimension ?? 'AWS Marketplace license',
74-
description: 'Admin linked an AWS Marketplace subscription to an organization',
75-
metadata: { ...licenseAuditMetadata(license), organizationId: input.organizationId },
72+
...licenseAuditEntry(
73+
license,
74+
AuditAction.AWS_MARKETPLACE_LICENSE_LINKED,
75+
'Admin linked an AWS Marketplace subscription to an organization',
76+
{ organizationId: input.organizationId }
77+
),
7678
})
7779
}
7880
return { id: license.id, organizationId: input.organizationId }

‎apps/sim/lib/billing/aws-marketplace/application/use-cases.ts‎

Lines changed: 71 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { AuditAction, AuditResourceType } from '@sim/audit'
1+
import { AuditAction, type AuditActionType, AuditResourceType } from '@sim/audit'
22
import { db } from '@sim/db'
33
import type { AwsMarketplaceClaim } from '@/lib/api/contracts/aws-marketplace'
44
import { awsMarketplaceOperations } from '@/lib/billing/aws-marketplace/application/operations'
@@ -10,13 +10,15 @@ import {
1010
import { verifyLicense } from '@/lib/billing/aws-marketplace/entitlements'
1111
import {
1212
type AwsMarketplaceLicenseRecord,
13+
type BuyerContact,
1314
consumeLicenseClaim,
1415
findLicenseByClaimHash,
1516
isNewerRead,
1617
listAccountLicenses,
1718
listAdministeredOrganizations,
1819
listOrganizationLicenses,
1920
lockLicense,
21+
paysForSelfServePlan,
2022
recordLicenseVerification,
2123
saveLicenseContact,
2224
} from '@/lib/billing/aws-marketplace/repository'
@@ -28,6 +30,7 @@ import { recordProjectedUseCaseAuditEntries } from '@/lib/core/application/autho
2830
import { assertOperationPrincipal, type OperationUseCase } from '@/lib/core/application/operation'
2931
import { requireOrganizationSubjectMembership } from '@/lib/core/application/organization-authorization'
3032
import { OrchestrationError } from '@/lib/core/orchestration/types'
33+
import type { DbTransaction } from '@/lib/db/types'
3134

3235
/** The same answer as an unknown route, so a deployment that does not sell on AWS shows nothing. */
3336
function requireMarketplaceEnabled(): void {
@@ -93,6 +96,22 @@ export function requireLinkableTo(
9396
}
9497
}
9598

99+
/**
100+
* AWS bills a linked organization for the listing, so one already paying Stripe for a plan would
101+
* pay twice. Its subscription is cancelled first, by its owner or by support.
102+
*/
103+
export async function requireNotBilledThroughStripe(
104+
tx: DbTransaction,
105+
organizationId: string
106+
): Promise<void> {
107+
if (await paysForSelfServePlan(tx, organizationId)) {
108+
throw new OrchestrationError(
109+
'conflict',
110+
'This organization already pays for a Sim plan. Cancel that subscription, or contact support, before linking an AWS Marketplace subscription to it.'
111+
)
112+
}
113+
}
114+
96115
/** Enough for a buyer to recognize their account, and no more for someone holding their claim. */
97116
function maskAwsAccount(accountId: string): string {
98117
return `••••${accountId.slice(-4)}`
@@ -149,12 +168,6 @@ interface ClaimInput {
149168
claimToken: string | null
150169
}
151170

152-
interface BuyerContact {
153-
/** Whom to contact about provisioning; AWS asks the registration page to collect it. */
154-
contactEmail: string
155-
companyName: string | null
156-
}
157-
158171
type TakeInput = ClaimInput & BuyerContact
159172

160173
interface LinkLicenseInput extends TakeInput {
@@ -167,8 +180,20 @@ interface LinkedLicense {
167180
newlyLinked: boolean
168181
}
169182

170-
function contactOf(input: BuyerContact) {
171-
return { email: input.contactEmail, companyName: input.companyName }
183+
/** The license this browser's claim names, refused unless the claim is still open. */
184+
async function findOpenClaim(input: ClaimInput) {
185+
const claimHash = input.claimToken ? hashClaimToken(input.claimToken) : null
186+
const claimed = requireOpenClaim(
187+
claimHash ? await findLicenseByClaimHash(claimHash) : null,
188+
input.licenseId,
189+
claimHash
190+
)
191+
return { claimHash, claimed }
192+
}
193+
194+
/** Locks the claimed license and checks the claim again, so two tabs racing it take it once. */
195+
async function lockOpenClaim(tx: DbTransaction, input: ClaimInput, claimHash: string | null) {
196+
return requireOpenClaim(await lockLicense(tx, input.licenseId), input.licenseId, claimHash)
172197
}
173198

174199
/**
@@ -178,12 +203,7 @@ function contactOf(input: BuyerContact) {
178203
*/
179204
async function verifyClaimedLicense(input: ClaimInput, organizationId: string | null) {
180205
requireMarketplaceEnabled()
181-
const claimHash = input.claimToken ? hashClaimToken(input.claimToken) : null
182-
const claimed = requireOpenClaim(
183-
claimHash ? await findLicenseByClaimHash(claimHash) : null,
184-
input.licenseId,
185-
claimHash
186-
)
206+
const { claimHash, claimed } = await findOpenClaim(input)
187207
requireLinkableTo(claimed, organizationId)
188208

189209
const verifiedAt = new Date()
@@ -205,11 +225,25 @@ async function verifyClaimedLicense(input: ClaimInput, organizationId: string |
205225
return { claimHash, verification }
206226
}
207227

208-
export function licenseAuditMetadata(license: AwsMarketplaceLicenseRecord) {
228+
/** The audit entry for a license changing hands; `metadata` adds to its terms. */
229+
export function licenseAuditEntry(
230+
license: AwsMarketplaceLicenseRecord,
231+
action: AuditActionType,
232+
description: string,
233+
metadata: Record<string, unknown> = {}
234+
) {
209235
return {
210-
dimension: license.dimension,
211-
quantity: license.quantity,
212-
expiresAt: license.expiresAt?.toISOString() ?? null,
236+
action,
237+
resourceType: AuditResourceType.AWS_MARKETPLACE_LICENSE,
238+
resourceId: license.id,
239+
resourceName: license.dimension ?? 'AWS Marketplace license',
240+
description,
241+
metadata: {
242+
dimension: license.dimension,
243+
quantity: license.quantity,
244+
expiresAt: license.expiresAt?.toISOString() ?? null,
245+
...metadata,
246+
},
213247
}
214248
}
215249

@@ -226,11 +260,7 @@ export const linkAwsMarketplaceLicense = defineAuthorizedOrganizationUseCase({
226260
}: OrganizationUseCaseContext<LinkLicenseInput>): Promise<LinkedLicense> {
227261
const { claimHash, verification } = await verifyClaimedLicense(input, input.organizationId)
228262
return db.transaction(async (tx) => {
229-
const locked = requireOpenClaim(
230-
await lockLicense(tx, input.licenseId),
231-
input.licenseId,
232-
claimHash
233-
)
263+
const locked = await lockOpenClaim(tx, input, claimHash)
234264
requireLinkableTo(locked, input.organizationId)
235265
requireNotOverruled(locked, verification.verifiedAt)
236266
// Holds the admin's membership until commit, so a concurrent demotion waits for the link.
@@ -243,24 +273,22 @@ export const linkAwsMarketplaceLicense = defineAuthorizedOrganizationUseCase({
243273
{ executor: tx, forShare: true }
244274
)
245275
const newlyLinked = locked.organizationId !== input.organizationId
276+
if (newlyLinked) await requireNotBilledThroughStripe(tx, input.organizationId)
246277
const license = await consumeLicenseClaim(tx, locked, {
247278
link: newlyLinked ? { organizationId: input.organizationId, userId: context.userId } : null,
248-
contact: contactOf(input),
279+
contact: input,
249280
verification,
250281
})
251282
return { license, newlyLinked }
252283
})
253284
},
254285
projectAudit: ({ result }) =>
255286
result.newlyLinked
256-
? {
257-
action: AuditAction.AWS_MARKETPLACE_LICENSE_LINKED,
258-
resourceType: AuditResourceType.AWS_MARKETPLACE_LICENSE,
259-
resourceId: result.license.id,
260-
resourceName: result.license.dimension ?? 'AWS Marketplace license',
261-
description: 'Linked an AWS Marketplace subscription to the organization',
262-
metadata: licenseAuditMetadata(result.license),
263-
}
287+
? licenseAuditEntry(
288+
result.license,
289+
AuditAction.AWS_MARKETPLACE_LICENSE_LINKED,
290+
'Linked an AWS Marketplace subscription to the organization'
291+
)
264292
: [],
265293
})
266294

@@ -280,17 +308,13 @@ export const registerAwsMarketplaceLicenseToAccount: OperationUseCase<
280308
assertOperationPrincipal(principal, awsMarketplaceOperations.registerToAccount)
281309
const { claimHash, verification } = await verifyClaimedLicense(input, null)
282310
const { license, newlyRegistered } = await db.transaction(async (tx) => {
283-
const locked = requireOpenClaim(
284-
await lockLicense(tx, input.licenseId),
285-
input.licenseId,
286-
claimHash
287-
)
311+
const locked = await lockOpenClaim(tx, input, claimHash)
288312
requireLinkableTo(locked, null)
289313
requireNotOverruled(locked, verification.verifiedAt)
290314
const newlyRegistered = locked.linkedByUserId !== principal.userId
291315
const consumed = await consumeLicenseClaim(tx, locked, {
292316
link: newlyRegistered ? { organizationId: null, userId: principal.userId } : null,
293-
contact: contactOf(input),
317+
contact: input,
294318
verification,
295319
})
296320
return { license: consumed, newlyRegistered }
@@ -302,14 +326,11 @@ export const registerAwsMarketplaceLicenseToAccount: OperationUseCase<
302326
principal,
303327
request,
304328
[
305-
{
306-
action: AuditAction.AWS_MARKETPLACE_LICENSE_REGISTERED,
307-
resourceType: AuditResourceType.AWS_MARKETPLACE_LICENSE,
308-
resourceId: license.id,
309-
resourceName: license.dimension ?? 'AWS Marketplace license',
310-
description: 'Registered an AWS Marketplace subscription to an account',
311-
metadata: licenseAuditMetadata(license),
312-
},
329+
licenseAuditEntry(
330+
license,
331+
AuditAction.AWS_MARKETPLACE_LICENSE_REGISTERED,
332+
'Registered an AWS Marketplace subscription to an account'
333+
),
313334
]
314335
)
315336
}
@@ -330,22 +351,13 @@ export const saveAwsMarketplaceClaimContact: OperationUseCase<
330351
async execute({ principal, input }) {
331352
assertOperationPrincipal(principal, awsMarketplaceOperations.saveClaimContact)
332353
requireMarketplaceEnabled()
333-
const claimHash = input.claimToken ? hashClaimToken(input.claimToken) : null
334-
const claimed = requireOpenClaim(
335-
claimHash ? await findLicenseByClaimHash(claimHash) : null,
336-
input.licenseId,
337-
claimHash
338-
)
354+
const { claimHash, claimed } = await findOpenClaim(input)
339355
const verifiedAt = new Date()
340356
const summary = await verifyLicense(requireLicenseProduct(claimed), claimed.licenseArn)
341357
return db.transaction(async (tx) => {
342-
const locked = requireOpenClaim(
343-
await lockLicense(tx, input.licenseId),
344-
input.licenseId,
345-
claimHash
346-
)
358+
const locked = await lockOpenClaim(tx, input, claimHash)
347359
const license = await saveLicenseContact(tx, locked, {
348-
contact: contactOf(input),
360+
contact: input,
349361
verification: { summary, verifiedAt },
350362
})
351363
return { status: license.status }

‎apps/sim/lib/billing/aws-marketplace/checkout-guard.integration.ts‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,20 @@ import { generateId } from '@sim/utils/id'
44
import { inArray } from 'drizzle-orm'
55
import { afterAll, afterEach, describe, expect, it, vi } from 'vitest'
66

7-
const PRODUCT = 'product-placeholder'
87
const ACCOUNT = '000000000000'
98
const TEAM_PLAN = 'team_6000'
9+
const RETIRED_PRODUCT = 'retired-product-placeholder'
1010

11-
const marketplace = vi.hoisted(() => ({ enabled: true }))
11+
const marketplace = vi.hoisted(() => ({ enabled: true, product: 'product-placeholder' }))
12+
const PRODUCT = marketplace.product
1213

13-
vi.mock('@/lib/billing/aws-marketplace/config', () => ({
14-
isAwsMarketplaceEnabled: () => marketplace.enabled,
15-
}))
14+
vi.mock('@/lib/billing/aws-marketplace/config', async () => {
15+
const { parseAwsMarketplaceProducts } = await vi.importActual<
16+
typeof import('@/lib/billing/aws-marketplace/products')
17+
>('@/lib/billing/aws-marketplace/products')
18+
const sold = parseAwsMarketplaceProducts(`${marketplace.product}=contract`)
19+
return { getAwsMarketplaceProducts: () => (marketplace.enabled ? sold : new Map()) }
20+
})
1621

1722
import {
1823
assertPersonalCheckoutAllowed,
@@ -61,12 +66,13 @@ async function insertLicense(link: {
6166
organizationId: string | null
6267
linkedByUserId: string
6368
status: LicenseStatus
69+
productCode?: string
6470
}) {
6571
const now = new Date()
6672
await db.insert(awsMarketplaceLicense).values({
6773
id: generateId(),
6874
licenseArn: `license-${generateId()}`,
69-
productCode: PRODUCT,
75+
productCode: link.productCode ?? PRODUCT,
7076
customerAwsAccountId: ACCOUNT,
7177
dimension: 'enterprise',
7278
quantity: 25,
@@ -135,6 +141,21 @@ describe('Stripe checkout for an AWS Marketplace-billed reference', () => {
135141
).resolves.toBe(true)
136142
})
137143

144+
it('allows an organization whose only license is on a listing this deployment no longer sells', async () => {
145+
const ownerId = await createUser()
146+
const organizationId = await createOwnedOrganization(ownerId)
147+
await insertLicense({
148+
organizationId,
149+
linkedByUserId: ownerId,
150+
status: 'active',
151+
productCode: RETIRED_PRODUCT,
152+
})
153+
154+
await expect(
155+
authorizeSubscriptionReference(ownerId, organizationId, 'upgrade-subscription', TEAM_PLAN)
156+
).resolves.toBe(true)
157+
})
158+
138159
it('allows an organization with no license', async () => {
139160
const { ownerId, organizationId } = await organizationWithLicense(null)
140161

‎apps/sim/lib/billing/aws-marketplace/client.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,10 @@ let queueClient: SQSClient | null = null
2929
/**
3030
* Credentials for the Marketplace and event-queue clients only. On ECS the app's env also carries
3131
* a static IAM user's keys, which the SDK's default chain prefers and which Sim's general AWS
32-
* clients (`getAwsCredentials()`) sign with. Infra grants ResolveCustomer, GetEntitlements, and the
33-
* queue permissions to the ECS task role alone, and AWS asks for temporary role credentials here,
34-
* so whenever the container credential endpoint is present these clients sign as the task role.
35-
* Elsewhere (local development, tests) they keep the default chain.
32+
* clients (`getAwsCredentials()`) sign with. Infra grants ResolveCustomer, GetEntitlements,
33+
* BatchMeterUsage, and the queue permissions to the ECS task role alone, and AWS asks for
34+
* temporary role credentials here, so whenever the container credential endpoint is present these
35+
* clients sign as the task role. Elsewhere (local development, tests) they keep the default chain.
3636
*/
3737
function marketplaceClientConfig() {
3838
const relativeUri = env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI

0 commit comments

Comments
 (0)