Merge pull request #30 from Routstr/codex/investigate-logical-error-in-payment-code

Fix balance race conditions

mypy.ini

@@ -0,0 +1,2 @@
+[mypy]
+ignore_missing_imports = True

router/admin.py

@@ -108,8 +108,9 @@ async def dashboard(request: Request) -> str:
             f"<tr><td>{key.hashed_key}</td><td>{key.balance}</td><td>{key.total_spent}</td><td>{key.total_requests}</td><td>{key.refund_address}</td><td>{'{} ({} UTC)'.format(key.key_expiry_time, expiry_time_human_readable) if key.key_expiry_time else key.key_expiry_time}</td></tr>"
         )
 
-    # Calculate the total balance of all API keys
-    total_user_balance = int(sum(key.balance / 1000 for key in api_keys))
+    # Calculate the total balance of all API keys using integer arithmetic to
+    # avoid rounding issues.
+    total_user_balance = sum(key.balance for key in api_keys) // 1000
     # Fetch balance from cashu
     current_balance = (await WALLET.fetch_wallet_state()).balance
     owner_balance = current_balance - total_user_balance

router/auth.py

@@ -6,6 +6,7 @@ from typing import Optional
 
 
 from fastapi import HTTPException, Request
+from sqlmodel import update, col
 
 from .cashu import credit_balance, pay_out
 from .db import ApiKey, AsyncSession
@@ -149,12 +150,31 @@ async def pay_for_request(
             },
         )
 
-    # Charge the base cost for the request
-    key.balance -= COST_PER_REQUEST
-    key.total_spent += COST_PER_REQUEST
-    key.total_requests += 1
-    session.add(key)
+    # Charge the base cost for the request atomically to avoid race conditions
+    stmt = (
+        update(ApiKey)
+        .where(col(ApiKey.hashed_key) == key.hashed_key)
+        .where(col(ApiKey.balance) >= COST_PER_REQUEST)
+        .values(
+            balance=col(ApiKey.balance) - COST_PER_REQUEST,
+            total_spent=col(ApiKey.total_spent) + COST_PER_REQUEST,
+            total_requests=col(ApiKey.total_requests) + 1,
+        )
+    )
+    result = await session.exec(stmt)  # type: ignore[call-overload]
     await session.commit()
+    if result.rowcount == 0:
+        # Another concurrent request spent the balance first
+        raise HTTPException(
+            status_code=402,
+            detail={
+                "error": {
+                    "message": f"Insufficient balance: {COST_PER_REQUEST} mSats required. {key.balance} available.",
+                    "type": "insufficient_quota",
+                    "code": "insufficient_balance",
+                }
+            },
+        )
     await session.refresh(key)
 
 
@@ -232,6 +252,7 @@ async def adjust_payment_for_tokens(
     cost_difference = token_based_cost - COST_PER_REQUEST
 
     if cost_difference == 0:
+        await session.commit()
         return cost_data  # No adjustment needed
 
     if cost_difference > 0:
@@ -240,23 +261,39 @@ async def adjust_payment_for_tokens(
             print(
                 f"Warning: Insufficient balance for token-based pricing adjustment: {key.hashed_key[:10]}..."
             )
-            # Still proceed but log the issue - we already provided the service
-            # Add information about insufficient balance to cost data
             cost_data["warning"] = "Insufficient balance for full token-based pricing"
             cost_data["balance_shortage_msats"] = cost_difference - key.balance
+            await session.commit()
         else:
-            key.balance -= cost_difference
-            key.total_spent += cost_difference
-            cost_data["total_msats"] = COST_PER_REQUEST + cost_difference
+            charge_stmt = (
+                update(ApiKey)
+                .where(col(ApiKey.hashed_key) == key.hashed_key)
+                .where(col(ApiKey.balance) >= cost_difference)
+                .values(
+                    balance=col(ApiKey.balance) - cost_difference,
+                    total_spent=col(ApiKey.total_spent) + cost_difference,
+                )
+            )
+            result = await session.exec(charge_stmt)  # type: ignore[call-overload]
+            await session.commit()
+            if result.rowcount:
+                cost_data["total_msats"] = COST_PER_REQUEST + cost_difference
+                await session.refresh(key)
     else:
         # Refund some of the base cost
         refund = abs(cost_difference)
-        key.balance += refund
-        key.total_spent -= refund
+        refund_stmt = (
+            update(ApiKey)
+            .where(col(ApiKey.hashed_key) == key.hashed_key)
+            .values(
+                balance=col(ApiKey.balance) + refund,
+                total_spent=col(ApiKey.total_spent) - refund,
+            )
+        )
+        await session.exec(refund_stmt)  # type: ignore[call-overload]
+        await session.commit()
         cost_data["total_msats"] = COST_PER_REQUEST - refund
-
-    session.add(key)
-    await session.commit()
+        await session.refresh(key)
 
     asyncio.create_task(pay_out())
 

router/cashu.py

@@ -3,7 +3,7 @@ import asyncio
 import time
 
 from sixty_nuts import Wallet
-from sqlmodel import select, func, col
+from sqlmodel import select, func, col, update
 from .db import ApiKey, AsyncSession, get_session
 
 
@@ -84,11 +84,24 @@ async def credit_balance(cashu_token: str, key: ApiKey, session: AsyncSession) -
 
     amount_msats = amount_sats * 1000
     key.balance += amount_msats
+
     session.add(key)
+    await session.flush()
+
+    # Apply the balance change atomically to avoid race conditions when topping
+    # up the same key concurrently.
+    stmt = (
+        update(ApiKey)
+        .where(col(ApiKey.hashed_key) == key.hashed_key)
+        .values(balance=col(ApiKey.balance) + amount)
+    )
+    await session.exec(stmt)  # type: ignore[call-overload]
     await session.commit()
+
     return amount_msats
 
 
+
 async def check_for_refunds() -> None:
     """
     Periodically checks for API keys that are eligible for refunds and processes them.
@@ -129,8 +142,6 @@ async def check_for_refunds() -> None:
 
 
 async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession) -> int:
-    if key.balance < amount_msats:
-        raise ValueError("Insufficient balance.")
     if amount_msats <= 0:
         amount_msats = key.balance
 
@@ -139,9 +150,19 @@ async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession)
     if amount_sats == 0:
         raise ValueError("Amount too small to refund (less than 1 sat)")
 
-    key.balance -= amount_msats
-    session.add(key)
+    # Atomically deduct the balance to avoid race conditions when multiple
+    # refunds are triggered concurrently.
+    stmt = (
+        update(ApiKey)
+        .where(col(ApiKey.hashed_key) == key.hashed_key)
+        .where(col(ApiKey.balance) >= amount_msats)
+        .values(balance=col(ApiKey.balance) - amount_msats)
+    )
+    result = await session.exec(stmt)  # type: ignore[call-overload]
     await session.commit()
+    if result.rowcount == 0:
+        raise ValueError("Insufficient balance.")
+    await session.refresh(key)
 
     if key.refund_address is None:
         raise ValueError("Refund address not set.")