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
124 changes: 124 additions & 0 deletions apps/sim/lib/billing/aws-marketplace/container-license.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'

const { mockSend } = vi.hoisted(() => ({ mockSend: vi.fn() }))

vi.mock('@aws-sdk/client-license-manager', () => ({
LicenseManagerClient: class {
send = mockSend
},
CheckoutLicenseCommand: class {
constructor(public input: { ProductSKU: string }) {}
},
}))
vi.mock('@/lib/core/network/proxy-agent.server', () => ({ awsRequestHandler: () => undefined }))
vi.mock('@sim/utils/random', () => ({ randomInt: (min: number) => min }))

import { applyAwsMarketplaceContainerLicense } from '@/lib/billing/aws-marketplace/container-license'

const MINUTE = 60 * 1000

type Answer =
| 'entitled'
| 'NoEntitlementsAllowedException'
| 'ResourceNotFoundException'
| 'AccessDeniedException'
| 'ServerInternalException'

/** License Manager's checkout answers in order; the last one repeats once the list runs out. */
let answers: Answer[] = []
const checkoutRequests: unknown[] = []

function awsError(name: string): Error {
return Object.assign(new Error(name), { name })
}

function enterpriseEnabled(): boolean {
return process.env.ENTERPRISE_ENABLED === 'true'
}

describe('applyAwsMarketplaceContainerLicense', () => {
let kill: ReturnType<typeof vi.spyOn>

beforeEach(() => {
vi.useFakeTimers({ toFake: ['setTimeout'] })
answers = []
checkoutRequests.length = 0
vi.stubEnv('AWS_MARKETPLACE_LICENSE_PRODUCT_SKU', 'product-placeholder')
kill = vi.spyOn(process, 'kill').mockImplementation(() => true)
mockSend.mockReset()
mockSend.mockImplementation(async (command: { input: unknown }) => {
checkoutRequests.push(command.input)
const answer = answers.length > 1 ? answers.shift() : answers[0]
if (answer !== 'entitled') throw awsError(answer ?? 'ServerInternalException')
return { EntitlementsAllowed: [{ Name: 'enterprise_license', Unit: 'None' }] }
})
})

afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
kill.mockRestore()
Reflect.deleteProperty(process.env, 'ENTERPRISE_ENABLED')
Reflect.deleteProperty(process.env, 'NEXT_PUBLIC_ENTERPRISE_ENABLED')
})

it('does nothing outside a Marketplace install', async () => {
vi.stubEnv('AWS_MARKETPLACE_LICENSE_PRODUCT_SKU', undefined)
await applyAwsMarketplaceContainerLicense()
expect(mockSend).not.toHaveBeenCalled()
})

it('turns Enterprise on with an uncounted checkout of the entitlement', async () => {
answers = ['entitled']
await applyAwsMarketplaceContainerLicense()
expect(enterpriseEnabled()).toBe(true)
expect(checkoutRequests).toEqual([
expect.objectContaining({
ProductSKU: 'product-placeholder',
CheckoutType: 'PROVISIONAL',
KeyFingerprint: 'aws:294406891311:AWS/Marketplace:issuer-fingerprint',
Entitlements: [{ Name: 'enterprise_license', Unit: 'None' }],
}),
])
})

it('runs without Enterprise when no license grants the entitlement', async () => {
answers = ['NoEntitlementsAllowedException']
await applyAwsMarketplaceContainerLicense()
expect(enterpriseEnabled()).toBe(false)
})

it('runs without Enterprise when its IAM role cannot check out', async () => {
answers = ['AccessDeniedException']
await applyAwsMarketplaceContainerLicense()
expect(enterpriseEnabled()).toBe(false)
})

it.each<Answer>(['NoEntitlementsAllowedException', 'ResourceNotFoundException'])(
'restarts without Enterprise once the license ends (%s)',
async (answer) => {
answers = ['entitled', answer]
await applyAwsMarketplaceContainerLicense()
await vi.advanceTimersByTimeAsync(16 * MINUTE)
expect(kill).toHaveBeenCalledWith(process.pid, 'SIGTERM')
}
)

it('restarts into Enterprise once the license becomes active', async () => {
answers = ['NoEntitlementsAllowedException', 'entitled']
await applyAwsMarketplaceContainerLicense()
expect(enterpriseEnabled()).toBe(false)
await vi.advanceTimersByTimeAsync(16 * MINUTE)
expect(kill).toHaveBeenCalledWith(process.pid, 'SIGTERM')
})

it.each<Answer>(['ServerInternalException', 'AccessDeniedException'])(
'never restarts a running Enterprise deployment over a failed check (%s)',
async (answer) => {
answers = ['entitled', answer]
await applyAwsMarketplaceContainerLicense()
await vi.advanceTimersByTimeAsync(6 * 60 * MINUTE)
expect(kill).not.toHaveBeenCalled()
}
)
})
151 changes: 77 additions & 74 deletions apps/sim/lib/billing/aws-marketplace/container-license.ts
Original file line number Diff line number Diff line change
@@ -1,26 +1,21 @@
import {
type GrantedLicense,
LicenseManagerClient,
ListReceivedLicensesCommand,
} from '@aws-sdk/client-license-manager'
import { CheckoutLicenseCommand, LicenseManagerClient } from '@aws-sdk/client-license-manager'
import { createLogger } from '@sim/logger'
import { getErrorMessage } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { randomInt } from '@sim/utils/random'
import { awsRequestHandler } from '@/lib/core/network/proxy-agent.server'

/**
* Enterprise entitlement for Sim installed from the AWS Marketplace container listing.
*
* The listing sells an annual contract whose one dimension is `enterprise_license`. AWS Marketplace
* grants the buyer a license in AWS License Manager, and the running product reads it. While the
* license is active, the Enterprise suite is on (`ENTERPRISE_ENABLED`). Without one, Sim runs as the
* open-source core, so a lapsed renewal or an unreachable License Manager never locks a buyer out
* of their deployment.
*
* The check reads the license rather than checking a unit out of it. A contract usually covers the
* whole deployment with one unit, so checking it out from every replica would make them compete
* for it, and a unit left checked out (a failed check-in holds it for an hour) would read exactly
* like an expired license.
* grants the buyer a license in AWS License Manager, and the running product checks it the way AWS
* documents for a contract that unlocks a tier of features rather than metering a quantity: a
* provisional `CheckoutLicense` of the entitlement with unit `None`. Such a checkout is not counted
* against the contract, so every replica can check at once, and it needs only the permission AWS's
* managed `AWSLicenseManagerConsumptionPolicy` grants. While the license is active, the Enterprise
* suite is on (`ENTERPRISE_ENABLED`). Without one, Sim runs as the open-source core, so a lapsed
* renewal or an unreachable License Manager never locks a buyer out of their deployment.
*
* Opt-in: nothing here runs unless {@link PRODUCT_SKU_ENV} is set, which only the Marketplace
* Helm profile does. Enterprise flags are read once at startup by both the server and the
Expand All @@ -34,8 +29,8 @@ const logger = createLogger('AwsMarketplaceContainerLicense')
const PRODUCT_SKU_ENV = 'AWS_MARKETPLACE_LICENSE_PRODUCT_SKU'

/**
* Optional override for where the license is read. License Manager replicates a granted license to
* other Regions, so the default is the pod's own Region, the same chain AWS's EKS walkthrough uses.
* Optional override for where the license is checked. License Manager replicates a granted license
* to other Regions, so the default is the pod's own Region, the same chain AWS's EKS walkthrough uses.
*/
const LICENSE_REGION_ENV = 'AWS_MARKETPLACE_LICENSE_REGION'

Expand All @@ -48,21 +43,26 @@ const ENTERPRISE_ENTITLEMENT = 'enterprise_license'
/** AWS recommends re-checking every 15 minutes to pick up renewals, upgrades, and cancellations. */
const RECHECK_INTERVAL_MS = 15 * 60 * 1000

/**
* Replicas start together, so their re-checks line up; each waits a random part of this before
* restarting, so a renewal or lapse never takes every replica down at once.
*/
const RESTART_SPREAD_MS = 10 * 60 * 1000
/** Replicas start together; a random part of this added to each re-check keeps them apart. */
const RECHECK_JITTER_MS = 2 * 60 * 1000

/**
* How the buyer's account holds the license. A delegate grant only shows licenses distributed to
* other accounts, and a grant still awaiting acceptance, or rejected, disabled, or deleted by the
* license administrator, entitles nothing. No status at all is a license the account bought itself.
* Replicas see a change together; each waits a random part of this before restarting, so a renewal
* or lapse never takes every replica down at once.
*/
const USABLE_RECEIVED_STATUSES = new Set(['ACTIVE', 'WORKFLOW_COMPLETED'])
const RESTART_SPREAD_MS = 10 * 60 * 1000

const ENTERPRISE_FLAGS = ['ENTERPRISE_ENABLED', 'NEXT_PUBLIC_ENTERPRISE_ENABLED'] as const

/** License Manager's answers when the account holds no active license with the entitlement. */
const NO_ENTITLEMENT_ERRORS = new Set([
'NoEntitlementsAllowedException',
'ResourceNotFoundException',
Comment thread
waleedlatif1 marked this conversation as resolved.
Resolved
])

/** The ServiceAccount's IAM role does not allow the check: a setup problem, never a verdict. */
const PERMISSION_ERRORS = new Set(['AccessDeniedException', 'AuthorizationException'])

/** `unknown` is a failed check (network, permissions, throttling), never a verdict on the license. */
type LicenseState = 'entitled' | 'not-entitled' | 'unknown'

Expand All @@ -78,58 +78,62 @@ function createClient(): LicenseManagerClient {
})
}

function isUsableGrant({ ReceivedMetadata: received }: GrantedLicense): boolean {
if (received?.GrantType === 'DelegateGrant') return false
return !received?.ReceivedStatus || USABLE_RECEIVED_STATUSES.has(received.ReceivedStatus)
}

/**
* An AWS Marketplace license this product's buyer holds is current when License Manager reports it
* `AVAILABLE`, its issuer is AWS Marketplace, it carries the Enterprise entitlement, and the account
* holds it through a grant it may use. License Manager moves an ended, cancelled, or suspended
* license out of `AVAILABLE` itself.
*/
function isActiveEnterpriseLicense(license: GrantedLicense): boolean {
return (
license.Status === 'AVAILABLE' &&
isUsableGrant(license) &&
license.Issuer?.KeyFingerprint === MARKETPLACE_KEY_FINGERPRINT &&
(license.Entitlements ?? []).some((entitlement) => entitlement.Name === ENTERPRISE_ENTITLEMENT)
)
function classifyFailure(error: unknown): LicenseState {
const name = error instanceof Error ? error.name : undefined
if (name && NO_ENTITLEMENT_ERRORS.has(name)) {
logger.info('AWS License Manager reports no active Enterprise entitlement', { error: name })
return 'not-entitled'
}
if (name && PERMISSION_ERRORS.has(name)) {
logger.error(
'The IAM role cannot check the AWS Marketplace license; allow license-manager:CheckoutLicense',
{ error: name }
)
return 'unknown'
}
logger.warn('Could not check the AWS Marketplace Enterprise license', {
error: name,
message: getErrorMessage(error, 'Unknown error'),
})
return 'unknown'
}

async function checkLicense(
client: LicenseManagerClient,
productSku: string
): Promise<LicenseState> {
try {
let nextToken: string | undefined
do {
const page = await client.send(
new ListReceivedLicensesCommand({
Filters: [
{ Name: 'ProductSKU', Values: [productSku] },
{ Name: 'Fingerprint', Values: [MARKETPLACE_KEY_FINGERPRINT] },
],
NextToken: nextToken,
})
)
if ((page.Licenses ?? []).some(isActiveEnterpriseLicense)) return 'entitled'
nextToken = page.NextToken
} while (nextToken)
return 'not-entitled'
const checkout = await client.send(
new CheckoutLicenseCommand({
Comment thread
waleedlatif1 marked this conversation as resolved.
Resolved
ProductSKU: productSku,
CheckoutType: 'PROVISIONAL',
KeyFingerprint: MARKETPLACE_KEY_FINGERPRINT,
Entitlements: [{ Name: ENTERPRISE_ENTITLEMENT, Unit: 'None' }],
ClientToken: generateId(),
})
)
const allowed = (checkout.EntitlementsAllowed ?? []).some(
(entitlement) => entitlement.Name === ENTERPRISE_ENTITLEMENT
)
return allowed ? 'entitled' : 'not-entitled'
} catch (error) {
logger.warn('Could not check the AWS Marketplace Enterprise license', {
message: getErrorMessage(error, 'Unknown error'),
})
return 'unknown'
return classifyFailure(error)
}
}

function setEnterpriseFlags(enabled: boolean): void {
for (const flag of ENTERPRISE_FLAGS) process.env[flag] = enabled ? 'true' : 'false'
}

function restartToApply(state: LicenseState): void {
const delayMs = randomInt(0, RESTART_SPREAD_MS)
logger.info('AWS Marketplace Enterprise license changed; restarting to apply it', {
state,
delayMs,
})
setTimeout(() => process.kill(process.pid, 'SIGTERM'), delayMs)
}

/**
* Resolves the Enterprise license before the server loads, so every Enterprise gate reads the
* answer through `ENTERPRISE_ENABLED`, then keeps re-checking. Returns without doing anything on a
Expand All @@ -150,17 +154,16 @@ export async function applyAwsMarketplaceContainerLicense(): Promise<void> {
{ state: bootState }
)

const recheck = setInterval(async () => {
const scheduleRecheck = () => {
setTimeout(recheck, RECHECK_INTERVAL_MS + randomInt(0, RECHECK_JITTER_MS)).unref()
}
const recheck = async () => {
const state = await checkLicense(client, productSku)
const entitledNow = state === 'entitled'
if (state === 'unknown' || entitledNow === enterpriseAtBoot) return
clearInterval(recheck)
const delayMs = randomInt(0, RESTART_SPREAD_MS)
logger.info('AWS Marketplace Enterprise license changed; restarting to apply it', {
state,
delayMs,
})
setTimeout(() => process.kill(process.pid, 'SIGTERM'), delayMs)
}, RECHECK_INTERVAL_MS)
recheck.unref()
// A failed check keeps the current mode; only a definite answer that differs restarts.
const changed =
(state === 'entitled' && !enterpriseAtBoot) || (state === 'not-entitled' && enterpriseAtBoot)
if (changed) return restartToApply(state)
scheduleRecheck()
}
scheduleRecheck()
}
2 changes: 1 addition & 1 deletion helm/sim/Chart.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ apiVersion: v2
name: sim
description: A Helm chart for Sim - the open-source AI workspace where teams build, deploy, and manage AI agents
type: application
version: 1.11.16
version: 1.11.17
appVersion: "v0.8.26"
kubeVersion: ">=1.25.0-0"
home: https://sim.ai
Expand Down
7 changes: 4 additions & 3 deletions helm/sim/examples/aws-marketplace-prerequisites.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -159,9 +159,10 @@ Resources:
Version: '2012-10-17'
Statement:
- Effect: Allow
# Sim only reads the license. ListReceivedLicenses has no resource-level
# permissions, so it cannot be scoped below '*'.
Action: license-manager:ListReceivedLicenses
# Sim checks the license with an uncounted checkout, which AWS's managed
# AWSLicenseManagerConsumptionPolicy also grants. CheckoutLicense has no
# resource-level permissions, so it cannot be scoped below '*'.
Action: license-manager:CheckoutLicense
Resource: '*'

Outputs:
Expand Down
5 changes: 3 additions & 2 deletions helm/sim/examples/values-aws-marketplace.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,9 @@
# provider (for IRSA).
# 2. RDS for PostgreSQL 17 and ElastiCache reachable from the cluster. The listing's
# CloudFormation template creates both in your VPC.
# 3. An IAM role for the ServiceAccount below that allows license-manager:ListReceivedLicenses
# (the CloudFormation template creates it).
# 3. An IAM role for the ServiceAccount below that allows license-manager:CheckoutLicense (the
# CloudFormation template creates it; AWS's managed AWSLicenseManagerConsumptionPolicy
# grants it too).
# 4. An ACM certificate for the two hostnames below.
# 5. Two Secrets in the install namespace (no secret ever goes in this file). The template's
# DatabasePasswordArn and CacheAuthTokenArn outputs are the ARNs of the Secrets Manager
Expand Down
Loading