fix: close two real money/spend races (Tracecoin double-credit, unbounded AI overspend)
Some checks failed
build-and-release / build (ugc-content-creators) (push) Waiting to run
build-and-release / build (users) (push) Waiting to run
build-and-release / build (video-editors) (push) Waiting to run
build-and-release / build (catering-services) (push) Successful in 1m47s
build-and-release / build (companies) (push) Successful in 2m0s
build-and-release / build (cron) (push) Successful in 2m12s
build-and-release / build (customers) (push) Successful in 2m42s
build-and-release / build (gateway) (push) Successful in 1m0s
build-and-release / build (developers) (push) Successful in 1m27s
build-and-release / build (fitness-trainers) (push) Successful in 1m54s
build-and-release / build (employees) (push) Successful in 1m59s
build-and-release / build (job-seekers) (push) Successful in 1m57s
build-and-release / build (makeup-artists) (push) Successful in 1m39s
build-and-release / build (graphic-designers) (push) Has been cancelled
build-and-release / build (payments) (push) Has been cancelled
build-and-release / build (jobs) (push) Has been cancelled
build-and-release / build (social-media-managers) (push) Has been cancelled
build-and-release / build (tutors) (push) Has been cancelled
build-and-release / build (photographers) (push) Has been cancelled

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.
This commit is contained in:
Ashwin Kumar Sivakumar 2026-07-21 05:51:23 +05:30
parent 92ce2d2a86
commit 3e701f2fe6
4 changed files with 102 additions and 59 deletions

View file

@ -380,23 +380,6 @@ pub async fn run_auto_apply(pool: &PgPool) -> Result<(), Box<dyn std::error::Err
.await .await
.ok(); .ok();
// Increment daily_actions_used on subscription
sqlx::query(
r#"
UPDATE user_ai_subscriptions
SET daily_actions_used = daily_actions_used + 1,
updated_at = NOW()
WHERE user_id = $1
AND status = 'active'
AND NOW() >= current_period_start
AND NOW() < current_period_end
"#,
)
.bind(seeker.user_id)
.execute(pool)
.await
.ok();
tracing::info!( tracing::info!(
"Auto-applied user {} to job '{}' ({})", "Auto-applied user {} to job '{}' ({})",
seeker.user_id, seeker.user_id,

View file

@ -480,29 +480,45 @@ async fn verify_payment(
return Err((StatusCode::BAD_REQUEST, "Payment was not successful".to_string())); 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>( let payment = sqlx::query_as::<_, PaymentRow>(
r#" r#"
SELECT id, reference_number, user_id, package_id, tracecoins_credited, amount_inr, status, payu_mihpayid SELECT id, reference_number, user_id, package_id, tracecoins_credited, amount_inr, status, payu_mihpayid
FROM payments FROM payments
WHERE payu_txnid = $1 AND status = 'PENDING' WHERE payu_txnid = $1 AND status = 'PENDING'
FOR UPDATE
"#, "#,
) )
.bind(&payload.txnid) .bind(&payload.txnid)
.fetch_optional(&state.pool) .fetch_optional(&mut *tx)
.await .await
.map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?;
let payment = match payment { let payment = match payment {
Some(payment) => payment, Some(payment) => payment,
None => { None => {
let _ = tx.rollback().await;
return Err(error_response( return Err(error_response(
StatusCode::NOT_FOUND, StatusCode::NOT_FOUND,
"Payment not found or already processed", "Payment not found or already processed",
)) ));
} }
}; };
if payment.user_id != auth.user_id { if payment.user_id != auth.user_id {
let _ = tx.rollback().await;
return Err(error_response( return Err(error_response(
StatusCode::FORBIDDEN, StatusCode::FORBIDDEN,
"Payment does not belong to user", "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") sqlx::query("UPDATE payments SET status = 'FAILED', payu_mihpayid = $1 WHERE id = $2")
.bind(&payload.mihpayid) .bind(&payload.mihpayid)
.bind(payment.id) .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 .await
.map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; .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(&payload.mihpayid)
.bind(payment.id) .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 .await
.map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?;
sqlx::query( sqlx::query(
r#" r#"
INSERT INTO tracecoin_wallets (user_id, balance, reserved) INSERT INTO tracecoin_ledger (wallet_id, transaction_type, amount, reference_type, reference_id)
VALUES ($1, $2, 0) VALUES ($1, 'CREDIT', $2, 'PAYMENT', $3)
ON CONFLICT (user_id) DO UPDATE SET
balance = tracecoin_wallets.balance + excluded.balance
"#, "#,
) )
.bind(payment.user_id) .bind(wallet_id)
.bind(tracecoins as i64) .bind(tracecoins as i64)
.execute(&state.pool) .bind(payment.id)
.execute(&mut *tx)
.await .await
.map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?;
if let Ok(Some(wallet_id)) = sqlx::query_scalar::<_, Uuid>( tx.commit()
"SELECT id FROM tracecoin_wallets WHERE user_id = $1", .await
) .map_err(|e| error_response(StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?;
.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;
}
generate_purchase_invoice(&state.pool, &payment, &payload).await; generate_purchase_invoice(&state.pool, &payment, &payload).await;

View file

@ -85,22 +85,31 @@ pub async fn charge_feature(
.await? .await?
.ok_or(CreditError::InsufficientCredits)?; .ok_or(CreditError::InsufficientCredits)?;
// Check daily action limit. // Single atomic conditional UPDATE — enforces the credit balance and
if sub.daily_actions_used >= plan.daily_action_limit { // daily action limit and deducts in the same statement, so concurrent
return Err(CreditError::DailyActionLimitReached); // 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. let sub = match charged {
if remaining_credits(&sub) < cost.credit_cost { Some(updated) => updated,
return Err(CreditError::InsufficientCredits); None => {
} if sub.daily_actions_used >= plan.daily_action_limit {
return Err(CreditError::DailyActionLimitReached);
// Charge. }
UserAiSubscriptionRepository::charge_credits(pool, user_id, cost.credit_cost).await?; return Err(CreditError::InsufficientCredits);
UserAiSubscriptionRepository::increment_daily_actions(pool, user_id).await?; }
};
// Record transaction. // Record transaction.
let balance_after = remaining_credits(&sub) - cost.credit_cost; let balance_after = remaining_credits(&sub);
AiCreditTransactionRepository::create( AiCreditTransactionRepository::create(
pool, pool,
user_id, user_id,

View file

@ -183,6 +183,38 @@ impl UserAiSubscriptionRepository {
Ok(()) 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<Option<UserAiSubscription>, 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( pub async fn increment_daily_actions(
pool: &PgPool, pool: &PgPool,
user_id: Uuid, user_id: Uuid,