From 7d8ed5b8113769796236c27c6b757340d96d31a4 Mon Sep 17 00:00:00 2001 From: lewistemple22 Date: Tue, 1 Sep 2026 01:11:02 +0100 Subject: [PATCH] feat(budgets): add Redis-backed distributed mutex guard for spending limits Closes #86 --- src/modules/budgets/budget.module.ts | 5 +- .../spending-limit-guard.service.spec.ts | 76 +++++++++++++++++++ .../services/spending-limit-guard.service.ts | 53 +++++++++++++ .../transactions/transaction.service.ts | 23 ++++-- 4 files changed, 148 insertions(+), 9 deletions(-) create mode 100644 src/modules/budgets/services/spending-limit-guard.service.spec.ts create mode 100644 src/modules/budgets/services/spending-limit-guard.service.ts diff --git a/src/modules/budgets/budget.module.ts b/src/modules/budgets/budget.module.ts index 87f99b2..8ef2ccd 100644 --- a/src/modules/budgets/budget.module.ts +++ b/src/modules/budgets/budget.module.ts @@ -4,6 +4,7 @@ import { BudgetService } from './budget.service'; import { BudgetRepository } from './budget.repository'; import { PolicyEvaluatorService } from './services/policy-evaluator.service'; 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 @@ -16,7 +17,7 @@ import { RollingWindowBudgetService } from './services/rolling-window-budget.ser */ @Module({ controllers: [BudgetController], - providers: [BudgetService, BudgetRepository, 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 107f0b5..81b29e1 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,13 +100,20 @@ export class TransactionService { { actorId }, ); - // 5. Budget headroom (no mutation yet). + const requiresApproval = policyResult.requiresApproval || !assessment.canAutoExecute; + + // 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) { - await this.budgets.assertWithinBudget(organizationId, input.budgetId, amount); + if (requiresApproval) { + await this.budgets.assertWithinBudget(organizationId, input.budgetId, amount); + } else { + await this.spendingGuard.guardAndConsume(organizationId, input.budgetId, amount); + budgetConsumed = true; + } } - const requiresApproval = policyResult.requiresApproval || !assessment.canAutoExecute; - // 6. Persist the transaction row. const transaction = await this.repository.create({ organization: { connect: { id: organizationId } }, @@ -144,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 }; } @@ -153,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 || @@ -188,7 +197,7 @@ export class TransactionService { confirmationCount: result.ledger ? 1 : 0, }); - if (result.successful && tx.budgetId) { + if (result.successful && tx.budgetId && !budgetConsumed) { await this.budgets.consume(organizationId, tx.budgetId, Number(tx.amount)); }