From 3e701f2fe6d7001f4609ab0a3242ec296076d941 Mon Sep 17 00:00:00 2001 From: Ashwin Kumar Sivakumar Date: Tue, 21 Jul 2026 05:51:23 +0530 Subject: [PATCH] fix: close two real money/spend races (Tracecoin double-credit, unbounded AI overspend) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Asked to review Tracecoin and AI implementation safety. Found and fixed two exploitable TOCTOU races, plus a data-integrity bug: 1. apps/payments/src/main.rs::verify_payment — the PayU success callback is called directly by the client (not a server-to-server webhook), so a user fully controls how many times they replay a valid success payload. The payment "is it still PENDING" check and the "mark SUCCESS + credit wallet" write were separate, non-transactional queries — concurrent replays could both pass the check before either commits, double- (or N-times-) crediting the wallet for one real payment. Now wrapped in a single transaction with `SELECT ... FOR UPDATE` on the payments row, so a second concurrent call blocks until the first commits, then correctly sees the row is no longer PENDING (Postgres re-evaluates the WHERE clause via EvalPlanQual after the lock is granted). 2. crates/db/src/models/ai/repository.rs — UserAiSubscriptionRepository had the exact same shape of bug: apps/users/src/ai/credits.rs:: charge_feature read the subscription, checked daily-limit and credit balance, THEN issued two separate unconditional `UPDATE ... SET x = x + $1` statements with no WHERE guard on the balance. N concurrent requests from one user all pass the check before any deduction lands, running up unlimited LLM API spend (this endpoint is called before/around real LiteLLM calls, so the cost is real). Added UserAiSubscriptionRepository::try_charge — a single conditional UPDATE that checks the daily limit and credit balance and deducts atomically, returning None (mapped to the existing error types) if either check fails. 3. apps/cron/src/tasks/auto_apply.rs — daily_actions_used was being incremented twice per auto-applied job (once in the credit-deduct UPDATE, once more in a second, redundant UPDATE right after) — silently halving job seekers' effective daily auto-apply limit. Removed the redundant second UPDATE. Also added non-negative CHECK constraints directly to the live database (tracecoin_wallets.balance/reserved, user_ai_subscriptions.daily_actions_used/monthly_credits_used/ purchased_credits_used) as defense in depth — belt-and-suspenders in case a future code path reintroduces a similar bug. --- apps/cron/src/tasks/auto_apply.rs | 17 ------ apps/payments/src/main.rs | 77 +++++++++++++++++---------- apps/users/src/ai/credits.rs | 35 +++++++----- crates/db/src/models/ai/repository.rs | 32 +++++++++++ 4 files changed, 102 insertions(+), 59 deletions(-) diff --git a/apps/cron/src/tasks/auto_apply.rs b/apps/cron/src/tasks/auto_apply.rs index b3d605e..59fc286 100644 --- a/apps/cron/src/tasks/auto_apply.rs +++ b/apps/cron/src/tasks/auto_apply.rs @@ -380,23 +380,6 @@ pub async fn run_auto_apply(pool: &PgPool) -> Result<(), Box= current_period_start - AND NOW() < current_period_end - "#, - ) - .bind(seeker.user_id) - .execute(pool) - .await - .ok(); - tracing::info!( "Auto-applied user {} to job '{}' ({})", seeker.user_id, diff --git a/apps/payments/src/main.rs b/apps/payments/src/main.rs index a8549ac..870a494 100644 --- a/apps/payments/src/main.rs +++ b/apps/payments/src/main.rs @@ -480,29 +480,45 @@ async fn verify_payment( return Err((StatusCode::BAD_REQUEST, "Payment was not successful".to_string())); } + // The whole claim-and-credit sequence runs in one transaction with the + // payments row locked via SELECT ... FOR UPDATE. Without this, two + // concurrent calls with the same (replayed) valid PayU success payload + // could both observe status = 'PENDING' before either commits its + // UPDATE, and both would credit the wallet — a double-credit exploit a + // user fully controls, since this endpoint is called directly by the + // client after PayU redirects back, not by a server-to-server webhook. + let mut tx = state + .pool + .begin() + .await + .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + let payment = sqlx::query_as::<_, PaymentRow>( r#" SELECT id, reference_number, user_id, package_id, tracecoins_credited, amount_inr, status, payu_mihpayid FROM payments WHERE payu_txnid = $1 AND status = 'PENDING' + FOR UPDATE "#, ) .bind(&payload.txnid) - .fetch_optional(&state.pool) + .fetch_optional(&mut *tx) .await .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; let payment = match payment { Some(payment) => payment, None => { + let _ = tx.rollback().await; return Err(error_response( StatusCode::NOT_FOUND, "Payment not found or already processed", - )) + )); } }; if payment.user_id != auth.user_id { + let _ = tx.rollback().await; return Err(error_response( StatusCode::FORBIDDEN, "Payment does not belong to user", @@ -513,7 +529,11 @@ async fn verify_payment( sqlx::query("UPDATE payments SET status = 'FAILED', payu_mihpayid = $1 WHERE id = $2") .bind(&payload.mihpayid) .bind(payment.id) - .execute(&state.pool) + .execute(&mut *tx) + .await + .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + tx.commit() .await .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; @@ -539,43 +559,42 @@ async fn verify_payment( ) .bind(&payload.mihpayid) .bind(payment.id) - .execute(&state.pool) + .execute(&mut *tx) + .await + .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + let wallet_id: Uuid = sqlx::query_scalar( + r#" + INSERT INTO tracecoin_wallets (user_id, balance, reserved) + VALUES ($1, $2, 0) + ON CONFLICT (user_id) DO UPDATE SET + balance = tracecoin_wallets.balance + excluded.balance, + updated_at = NOW() + RETURNING id + "#, + ) + .bind(payment.user_id) + .bind(tracecoins as i64) + .fetch_one(&mut *tx) .await .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; sqlx::query( r#" - INSERT INTO tracecoin_wallets (user_id, balance, reserved) - VALUES ($1, $2, 0) - ON CONFLICT (user_id) DO UPDATE SET - balance = tracecoin_wallets.balance + excluded.balance + INSERT INTO tracecoin_ledger (wallet_id, transaction_type, amount, reference_type, reference_id) + VALUES ($1, 'CREDIT', $2, 'PAYMENT', $3) "#, ) - .bind(payment.user_id) + .bind(wallet_id) .bind(tracecoins as i64) - .execute(&state.pool) + .bind(payment.id) + .execute(&mut *tx) .await .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; - if let Ok(Some(wallet_id)) = sqlx::query_scalar::<_, Uuid>( - "SELECT id FROM tracecoin_wallets WHERE user_id = $1", - ) - .bind(payment.user_id) - .fetch_optional(&state.pool) - .await - { - let _ = sqlx::query( - r#" - INSERT INTO tracecoin_ledger (wallet_id, transaction_type, amount, reference_type, reference_id) - VALUES ($1, 'CREDIT', $2, 'PAYMENT', $3) - "#, - ) - .bind(wallet_id) - .bind(tracecoins as i64) - .bind(payment.id) - .execute(&state.pool) - .await; - } + tx.commit() + .await + .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; generate_purchase_invoice(&state.pool, &payment, &payload).await; diff --git a/apps/users/src/ai/credits.rs b/apps/users/src/ai/credits.rs index 773554c..6999277 100644 --- a/apps/users/src/ai/credits.rs +++ b/apps/users/src/ai/credits.rs @@ -85,22 +85,31 @@ pub async fn charge_feature( .await? .ok_or(CreditError::InsufficientCredits)?; - // Check daily action limit. - if sub.daily_actions_used >= plan.daily_action_limit { - return Err(CreditError::DailyActionLimitReached); - } + // Single atomic conditional UPDATE — enforces the credit balance and + // daily action limit and deducts in the same statement, so concurrent + // requests from the same user can't all pass a separate check before + // any of them commits a deduction (which would let a user run up + // unbounded LLM spend by firing requests in parallel). + let charged = UserAiSubscriptionRepository::try_charge( + pool, + user_id, + cost.credit_cost, + plan.daily_action_limit, + ) + .await?; - // Check credits. - if remaining_credits(&sub) < cost.credit_cost { - return Err(CreditError::InsufficientCredits); - } - - // Charge. - UserAiSubscriptionRepository::charge_credits(pool, user_id, cost.credit_cost).await?; - UserAiSubscriptionRepository::increment_daily_actions(pool, user_id).await?; + let sub = match charged { + Some(updated) => updated, + None => { + if sub.daily_actions_used >= plan.daily_action_limit { + return Err(CreditError::DailyActionLimitReached); + } + return Err(CreditError::InsufficientCredits); + } + }; // Record transaction. - let balance_after = remaining_credits(&sub) - cost.credit_cost; + let balance_after = remaining_credits(&sub); AiCreditTransactionRepository::create( pool, user_id, diff --git a/crates/db/src/models/ai/repository.rs b/crates/db/src/models/ai/repository.rs index 955afdd..b9e7df2 100644 --- a/crates/db/src/models/ai/repository.rs +++ b/crates/db/src/models/ai/repository.rs @@ -183,6 +183,38 @@ impl UserAiSubscriptionRepository { Ok(()) } + /// Atomically checks and deducts credits + daily action count in a single + /// conditional UPDATE, so concurrent calls for the same user can't both + /// pass a separate check-then-act step and overspend past the limit. + /// Returns the updated row if the charge succeeded, or None if the user + /// didn't have enough credits or hit their daily action limit (caller + /// can re-read the subscription to report which one). + pub async fn try_charge( + pool: &PgPool, + user_id: Uuid, + credits: i32, + daily_action_limit: i32, + ) -> Result, sqlx::Error> { + sqlx::query_as::<_, UserAiSubscription>( + r#" + UPDATE user_ai_subscriptions + SET monthly_credits_used = monthly_credits_used + $1, + daily_actions_used = daily_actions_used + 1, + updated_at = NOW() + WHERE user_id = $2 + AND daily_actions_used < $3 + AND (monthly_credits_total + purchased_credits_total + - monthly_credits_used - purchased_credits_used) >= $1 + RETURNING * + "#, + ) + .bind(credits) + .bind(user_id) + .bind(daily_action_limit) + .fetch_optional(pool) + .await + } + pub async fn increment_daily_actions( pool: &PgPool, user_id: Uuid,