diff --git a/src/modules/budgets/budget.module.ts b/src/modules/budgets/budget.module.ts index 2a51b3b..bebe0bc 100644 --- a/src/modules/budgets/budget.module.ts +++ b/src/modules/budgets/budget.module.ts @@ -6,6 +6,7 @@ import { BudgetReservationService } from './services/budget-reservation.service' import { PolicyEvaluatorService } from './services/policy-evaluator.service'; import { RedisLock } from '../../common/locks/redis-lock.util'; import { RollingWindowBudgetService } from './services/rolling-window-budget.service'; +import { SpendingLimitGuardService } from './services/spending-limit-guard.service'; /** * Budget module. Exports the service so the transactions pipeline can enforce @@ -20,14 +21,7 @@ import { RollingWindowBudgetService } from './services/rolling-window-budget.ser */ @Module({ controllers: [BudgetController], - providers: [ - BudgetService, - BudgetRepository, - BudgetReservationService, - RedisLock, - PolicyEvaluatorService, - RollingWindowBudgetService, - ], - exports: [BudgetService, PolicyEvaluatorService, RollingWindowBudgetService], + providers: [BudgetService, BudgetRepository, PolicyEvaluatorService, RollingWindowBudgetService, SpendingLimitGuardService], + exports: [BudgetService, PolicyEvaluatorService, RollingWindowBudgetService, SpendingLimitGuardService], }) export class BudgetModule {} diff --git a/src/modules/budgets/services/spending-limit-guard.service.spec.ts b/src/modules/budgets/services/spending-limit-guard.service.spec.ts new file mode 100644 index 0000000..1ace7d3 --- /dev/null +++ b/src/modules/budgets/services/spending-limit-guard.service.spec.ts @@ -0,0 +1,76 @@ +import { describe, expect, it, vi, beforeEach } from 'vitest'; +import { SpendingLimitGuardService } from './spending-limit-guard.service'; +import { RedisLock } from '../../../common/locks/redis-lock.util'; +import { BudgetService } from '../budget.service'; + +describe('SpendingLimitGuardService', () => { + let service: SpendingLimitGuardService; + let redisLock: RedisLock; + let budgets: BudgetService; + + beforeEach(() => { + redisLock = { + withLock: vi.fn(), + } as unknown as RedisLock; + + budgets = { + assertWithinBudget: vi.fn(), + consume: vi.fn(), + } as unknown as BudgetService; + + service = new SpendingLimitGuardService(redisLock, budgets); + }); + + it('should acquire lock and call assertWithinBudget + consume', async () => { + const fakeBudget = { id: 'b1', spent: 0, limitAmount: { toFixed: () => '100' } }; + vi.mocked(redisLock.withLock).mockImplementation(async (_key, fn) => fn()); + vi.mocked(budgets.assertWithinBudget).mockResolvedValue(undefined as never); + vi.mocked(budgets.consume).mockResolvedValue(fakeBudget as never); + + const result = await service.guardAndConsume('org1', 'b1', 50); + + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:guard:b1', + expect.any(Function), + 5000, + ); + expect(budgets.assertWithinBudget).toHaveBeenCalledWith('org1', 'b1', 50); + expect(budgets.consume).toHaveBeenCalledWith('org1', 'b1', 50); + expect(result).toBe(fakeBudget); + }); + + it('should not call consume if assertWithinBudget throws', async () => { + vi.mocked(redisLock.withLock).mockImplementation(async (_key, fn) => fn()); + vi.mocked(budgets.assertWithinBudget).mockRejectedValueOnce( + new Error('Budget exceeded'), + ); + + await expect(service.guardAndConsume('org1', 'b1', 200)).rejects.toThrow('Budget exceeded'); + expect(budgets.consume).not.toHaveBeenCalled(); + }); + + it('should use custom TTL when provided', async () => { + vi.mocked(redisLock.withLock).mockImplementation(async (_key, fn) => fn()); + vi.mocked(budgets.assertWithinBudget).mockResolvedValue(undefined as never); + vi.mocked(budgets.consume).mockResolvedValue({} as never); + + await service.guardAndConsume('org1', 'b1', 10, 10_000); + + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:guard:b1', + expect.any(Function), + 10_000, + ); + }); + + it('should propagate lock acquisition failure', async () => { + vi.mocked(redisLock.withLock).mockRejectedValue( + new Error('Lock not acquired'), + ); + + await expect(service.guardAndConsume('org1', 'b1', 50)).rejects.toThrow( + 'Lock not acquired', + ); + expect(budgets.assertWithinBudget).not.toHaveBeenCalled(); + }); +}); diff --git a/src/modules/budgets/services/spending-limit-guard.service.ts b/src/modules/budgets/services/spending-limit-guard.service.ts new file mode 100644 index 0000000..07b4a9a --- /dev/null +++ b/src/modules/budgets/services/spending-limit-guard.service.ts @@ -0,0 +1,53 @@ +import { Injectable } from '@nestjs/common'; +import { RedisLock } from '../../../common/locks/redis-lock.util'; +import { BudgetService } from '../budget.service'; + +/** + * Wraps the budget check + consume cycle in a single Redis distributed lock + * to eliminate the TOCTOU race between assertWithinBudget (step 5) and + * consume (step 7) in the transaction pipeline. + * + * Without this guard, concurrent transactions can both pass the headroom + * check before either consumes, overshooting the budget limit. + * + * The lock key is `lock:budget:guard:{budgetId}` — distinct from the + * individual check/consume keys in BudgetService. + */ +@Injectable() +export class SpendingLimitGuardService { + constructor( + private readonly redisLock: RedisLock, + private readonly budgets: BudgetService, + ) {} + + /** + * Atomically validates headroom and debits the budget under a single + * distributed lock. Throws BudgetExceededException if the spend would + * breach the limit (no state is mutated in that case). + * + * @param organizationId - Owning organization. + * @param budgetId - Budget to check and consume from. + * @param amount - Amount to spend. + * @param ttlMs - Lock TTL in milliseconds (default 5 000). + * @returns The budget after consumption. + */ + async guardAndConsume( + organizationId: string, + budgetId: string, + amount: number, + ttlMs = 5_000, + ) { + const lockKey = `budget:guard:${budgetId}`; + return this.redisLock.withLock( + lockKey, + async () => { + // Step 1: Validate headroom (same logic as assertWithinBudget). + await this.budgets.assertWithinBudget(organizationId, budgetId, amount); + + // Step 2: Debit the budget (same logic as consume). + return this.budgets.consume(organizationId, budgetId, amount); + }, + ttlMs, + ); + } +} diff --git a/src/modules/transactions/transaction.service.ts b/src/modules/transactions/transaction.service.ts index 10fb939..ff20d76 100644 --- a/src/modules/transactions/transaction.service.ts +++ b/src/modules/transactions/transaction.service.ts @@ -16,6 +16,7 @@ import { WalletService, toNetworkName } from '../wallets/wallet.service'; import { PolicyService } from '../policies/policy.service'; import { RiskService } from '../risk/risk.service'; import { BudgetService } from '../budgets/budget.service'; +import { SpendingLimitGuardService } from '../budgets/services/spending-limit-guard.service'; import { StellarService } from '../stellar/stellar.service'; import { AgentService } from '../agents/agent.service'; import { TransactionIntent } from '../policies/policy.types'; @@ -65,6 +66,7 @@ export class TransactionService { private readonly policies: PolicyService, private readonly risk: RiskService, private readonly budgets: BudgetService, + private readonly spendingGuard: SpendingLimitGuardService, private readonly stellar: StellarService, private readonly eventBus: EventBusService, private readonly prisma: PrismaService, @@ -98,46 +100,41 @@ export class TransactionService { { actorId }, ); - // 5. Reserve the budget headroom (authoritative — persisted atomically under - // a lock, so concurrent agent requests cannot overspend a single budget). - if (input.budgetId) { - await this.budgets.reserve(organizationId, input.budgetId, amount); - } - const requiresApproval = policyResult.requiresApproval || !assessment.canAutoExecute; - // 6. Persist the transaction row. If persistence fails, return the reserved - // headroom so the budget is not silently eaten. - let transaction: Transaction; - try { - transaction = await this.repository.create({ - organization: { connect: { id: organizationId } }, - wallet: { connect: { id: wallet.id } }, - ...(input.agentId ? { agent: { connect: { id: input.agentId } } } : {}), - ...(policyResult.matchedPolicyId - ? { policy: { connect: { id: policyResult.matchedPolicyId } } } - : {}), - ...(input.budgetId ? { budget: { connect: { id: input.budgetId } } } : {}), - asset: input.asset, - amount: new Decimal(input.amount), - senderAddress: wallet.stellarAddress, - recipientAddress: input.recipientAddress, - memo: memoValue, - purpose: input.purpose, - status: TransactionStatus.DRAFT, - riskScore: assessment.score, - riskBand: assessment.band, - requiresApproval, - metadata: input.metadata as Prisma.InputJsonValue, - }); - } catch (error) { - if (input.budgetId) { - await this.budgets - .release(organizationId, input.budgetId, amount) - .catch(() => undefined); + // 5. Budget check — atomically guard+consume for auto-executable txns, + // or check-only for approval-required txns (consume deferred to execute). + let budgetConsumed = false; + if (input.budgetId) { + if (requiresApproval) { + await this.budgets.assertWithinBudget(organizationId, input.budgetId, amount); + } else { + await this.spendingGuard.guardAndConsume(organizationId, input.budgetId, amount); + budgetConsumed = true; } - throw error; } + + // 6. Persist the transaction row. + const transaction = await this.repository.create({ + organization: { connect: { id: organizationId } }, + wallet: { connect: { id: wallet.id } }, + ...(input.agentId ? { agent: { connect: { id: input.agentId } } } : {}), + ...(policyResult.matchedPolicyId + ? { policy: { connect: { id: policyResult.matchedPolicyId } } } + : {}), + ...(input.budgetId ? { budget: { connect: { id: input.budgetId } } } : {}), + asset: input.asset, + amount: new Decimal(input.amount), + senderAddress: wallet.stellarAddress, + recipientAddress: input.recipientAddress, + memo: memoValue, + purpose: input.purpose, + status: TransactionStatus.DRAFT, + riskScore: assessment.score, + riskBand: assessment.band, + requiresApproval, + metadata: input.metadata as Prisma.InputJsonValue, + }); await this.eventBus.emit( DomainEventName.TransactionCreated, { transactionId: transaction.id, amount: input.amount, riskBand: assessment.band }, @@ -156,7 +153,7 @@ export class TransactionService { }; } - const executed = await this.execute(organizationId, transaction.id, actorId); + const executed = await this.execute(organizationId, transaction.id, actorId, budgetConsumed); return { transaction: executed, requiresApproval: false, risk: assessment }; } @@ -165,7 +162,7 @@ export class TransactionService { * auto-executable transactions and by the approvals module once a proposal has * gathered the required approvals. */ - async execute(organizationId: string, transactionId: string, actorId?: string): Promise { + async execute(organizationId: string, transactionId: string, actorId?: string, budgetConsumed = false): Promise { const tx = await this.getOrThrow(organizationId, transactionId); if ( tx.status === TransactionStatus.COMPLETED || @@ -202,9 +199,7 @@ export class TransactionService { confirmationCount: result.ledger ? 1 : 0, }); - if (result.successful && tx.budgetId) { - // Settles the reservation booked at creation — spend was already - // reserved, so this only records the event and warns near the cap. + if (result.successful && tx.budgetId && !budgetConsumed) { await this.budgets.consume(organizationId, tx.budgetId, Number(tx.amount)); }