fix(quota): eliminate TOCTOU race conditions in quota service#1608
fix(quota): eliminate TOCTOU race conditions in quota service#1608atul-upadhyay-7 wants to merge 1 commit into
Conversation
|
@anshika1179 is attempting to deploy a commit to the Nisshchaya's projects Team on Vercel. A member of the Team first needs to authorize it. |
📝 WalkthroughWalkthroughThis PR expands ChangesQuota and Rate-Limit Service Expansion
Sequence DiagramsequenceDiagram
participant Client
participant checkAndReserveQuota
participant validateInstallationId
participant getQuotaMax
participant updateMany
participant fetchQuota
Client->>checkAndReserveQuota: installationId
checkAndReserveQuota->>validateInstallationId: validate
validateInstallationId-->>checkAndReserveQuota: valid/null
checkAndReserveQuota->>getQuotaMax: get maxRequests
getQuotaMax-->>checkAndReserveQuota: maxRequests
checkAndReserveQuota->>updateMany: atomic increment where requestsUsed < max
updateMany-->>checkAndReserveQuota: count affected rows
alt count > 0 (reserved)
checkAndReserveQuota-->>Client: true (allowed)
else count == 0 (update failed)
checkAndReserveQuota->>fetchQuota: fetch current quota
fetchQuota-->>checkAndReserveQuota: quota or null
alt quota null (record missing)
checkAndReserveQuota->>updateMany: upsert and retry
updateMany-->>checkAndReserveQuota: count
checkAndReserveQuota-->>Client: count > 0
else windowStart expired
checkAndReserveQuota->>updateMany: reset window and increment
updateMany-->>checkAndReserveQuota: count
checkAndReserveQuota-->>Client: count > 0
else true exhaustion
checkAndReserveQuota-->>Client: false (denied)
end
else DB error
checkAndReserveQuota-->>Client: false (fail-closed)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
lib/services/quotaService.ts (1)
336-351:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftRace condition remains in window reset path.
When
updateManyreturns 0 rows and the window is expired, multiple concurrent requests can all enter this branch simultaneously. Each executesupdate({ data: { requestsUsed: 1, ... } })and returnstrue. The result is that N concurrent requests all succeed whilerequestsUsedis only set to 1.While the main TOCTOU fix (atomic
updateManywithltcondition) is correct, this secondary race at window expiry allows quota bypass. The impact is limited to the moment of window transition (once per 24h typically), but could allow burst exceeding quota at that instant.Consider using
updateManywith a condition to ensure the window hasn't been reset by another request:Proposed atomic window reset approach
// Check if the window has expired and needs reset if (quota.quotaWindowEnd < now) { - // Reset the window: set requestsUsed to 1 (this request counts), - // clear token consumption, and start a new window. - await prisma.aiQuota.update({ - where: { id: quota.id }, + // Atomically reset only if window still expired (guards against concurrent resets). + const resetResult = await prisma.aiQuota.updateMany({ + where: { + installationId, + quotaWindowEnd: { lt: now }, + }, data: { requestsUsed: 1, tokensConsumed: 0, quotaWindowStart: now, quotaWindowEnd: windowEnd, warningPosted: false, lastAnalysisAt: now, }, }); - return true; + // If reset succeeded, this request reserved. If 0 rows, another request + // already reset the window - retry the main updateMany path. + if (resetResult.count > 0) { + return true; + } + // Retry: window was reset by another request, try atomic increment again. + const retryResult = await prisma.aiQuota.updateMany({ + where: { + installationId, + requestsUsed: { lt: defaultMaxAnalyses }, + }, + data: { + requestsUsed: { increment: 1 }, + lastAnalysisAt: now, + }, + }); + return retryResult.count > 0; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/services/quotaService.ts` around lines 336 - 351, The window-reset branch that calls prisma.aiQuota.update({ where: { id: quota.id }, ... }) is racy because multiple concurrent requests can each set requestsUsed to 1; replace that single-row update with an atomic updateMany using the same lt condition as the earlier attempt: call prisma.aiQuota.updateMany({ where: { id: quota.id, quotaWindowEnd: { lt: now } }, data: { requestsUsed: 1, tokensConsumed: 0, quotaWindowStart: now, quotaWindowEnd: windowEnd, warningPosted: false, lastAnalysisAt: now } }) and only treat the reset as successful if the returned count > 0 (otherwise fall through/retry); keep references to quota.quotaWindowEnd, prisma.aiQuota.updateMany, and quota.id to locate and implement the change.
🧹 Nitpick comments (2)
lib/services/__tests__/quotaService.test.ts (2)
261-277: ⚡ Quick winRestore
process.env.AI_QUOTA_PER_WINDOWinafterEachto prevent cross-test leakage.Tests 5.1–5.3 set and
deletethe env var inline. If any assertion throws before thedeleteline runs, the value leaks into later tests — notably 5.18 (Line 372) which depends on the default max of 250 to expectremainingRequests === 240. Snapshot/restore in the existingafterEachmakes isolation robust regardless of failures.♻️ Suggested isolation
afterEach(() => { jest.useRealTimers(); + delete process.env.AI_QUOTA_PER_WINDOW; });Then the inline
delete process.env.AI_QUOTA_PER_WINDOW;calls in 5.1–5.3 can be dropped.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/services/__tests__/quotaService.test.ts` around lines 261 - 277, The tests mutate process.env.AI_QUOTA_PER_WINDOW inline in the "Scenario 5.1/5.2/5.3" cases which can leak if an assertion throws; update the existing test teardown to snapshot the original process.env.AI_QUOTA_PER_WINDOW at the start of the suite and restore it in afterEach (so the environment is reset even on failure), remove the inline delete process.env.AI_QUOTA_PER_WINDOW calls from the three it blocks, and ensure QuotaService.getQuotaMax() tests still assert expected values while relying on the afterEach restoration for isolation.
256-259: 💤 Low valueTrim the "compliance checklist" to what the suite actually verifies.
The block is titled a "SECURITY COMPLIANCE VERIFICATION CHECKLIST" but many enumerated items have no corresponding test — e.g. concurrency/thread-safety (Lines 405–406), BigInt overflow (Line 396), SQL-injection exclusion (Line 400), cleanup batch-size enforcement (Line 399), and
lastAnalysisAtupdates (Line 432). Presenting these as verified can give a false sense of assurance during an audit. The "GSSoC '26 authorization" references (Lines 258, 395) also don't map to any control here. Consider reducing this to items the tests genuinely cover, or relabel it as an explicit TODO/backlog of gaps.Also applies to: 377-436
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/services/__tests__/quotaService.test.ts` around lines 256 - 259, The test-suite comment block titled "SECURITY COMPLIANCE VERIFICATION CHECKLIST" inside the describe("Scenario 5: Additional Quota Configuration & Exception Validations") is overstating coverage; remove or shrink the checklist to only items that have explicit tests, and convert remaining unsupported items (e.g., concurrency/thread-safety, BigInt overflow, SQL-injection exclusion, cleanup batch-size enforcement, lastAnalysisAt updates, and any "GSSoC '26 authorization" references) into a TODO/backlog note or separate risk list so the suite doesn't claim untested controls; update the comment text accordingly to accurately reflect what the tests in quotaService.test.ts actually verify.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@lib/services/quotaService.ts`:
- Around line 336-351: The window-reset branch that calls
prisma.aiQuota.update({ where: { id: quota.id }, ... }) is racy because multiple
concurrent requests can each set requestsUsed to 1; replace that single-row
update with an atomic updateMany using the same lt condition as the earlier
attempt: call prisma.aiQuota.updateMany({ where: { id: quota.id, quotaWindowEnd:
{ lt: now } }, data: { requestsUsed: 1, tokensConsumed: 0, quotaWindowStart:
now, quotaWindowEnd: windowEnd, warningPosted: false, lastAnalysisAt: now } })
and only treat the reset as successful if the returned count > 0 (otherwise fall
through/retry); keep references to quota.quotaWindowEnd,
prisma.aiQuota.updateMany, and quota.id to locate and implement the change.
---
Nitpick comments:
In `@lib/services/__tests__/quotaService.test.ts`:
- Around line 261-277: The tests mutate process.env.AI_QUOTA_PER_WINDOW inline
in the "Scenario 5.1/5.2/5.3" cases which can leak if an assertion throws;
update the existing test teardown to snapshot the original
process.env.AI_QUOTA_PER_WINDOW at the start of the suite and restore it in
afterEach (so the environment is reset even on failure), remove the inline
delete process.env.AI_QUOTA_PER_WINDOW calls from the three it blocks, and
ensure QuotaService.getQuotaMax() tests still assert expected values while
relying on the afterEach restoration for isolation.
- Around line 256-259: The test-suite comment block titled "SECURITY COMPLIANCE
VERIFICATION CHECKLIST" inside the describe("Scenario 5: Additional Quota
Configuration & Exception Validations") is overstating coverage; remove or
shrink the checklist to only items that have explicit tests, and convert
remaining unsupported items (e.g., concurrency/thread-safety, BigInt overflow,
SQL-injection exclusion, cleanup batch-size enforcement, lastAnalysisAt updates,
and any "GSSoC '26 authorization" references) into a TODO/backlog note or
separate risk list so the suite doesn't claim untested controls; update the
comment text accordingly to accurately reflect what the tests in
quotaService.test.ts actually verify.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5bd9427b-8d6d-470e-8505-70300c10c60a
📒 Files selected for processing (2)
lib/services/__tests__/quotaService.test.tslib/services/quotaService.ts
Closes #1588
Summary by CodeRabbit
New Features
Tests