From 951eb32c2b60df657fb6c4608b0eb2fb7d170e10 Mon Sep 17 00:00:00 2001
From: Shroominic <db2002dominic@gmail.com>
Date: Mon, 25 Aug 2025 15:03:43 -0300
Subject: [PATCH] cached /refund

---
 routstr/balance.py | 51 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 2 deletions(-)

diff --git a/routstr/balance.py b/routstr/balance.py
index 450fce43..97039905 100644
--- a/routstr/balance.py
+++ b/routstr/balance.py
@@ -1,3 +1,7 @@
+import asyncio
+import hashlib
+import os
+from time import monotonic
 from typing import Annotated, NoReturn
 
 from fastapi import APIRouter, Depends, Header, HTTPException
@@ -98,11 +102,52 @@ async def topup_wallet_endpoint(
     return {"msats": amount_msats}
 
 
+_REFUND_CACHE_TTL_SECONDS: int = int(os.environ.get("REFUND_CACHE_TTL_SECONDS", "3600"))
+_refund_cache_lock: asyncio.Lock = asyncio.Lock()
+_refund_cache: dict[str, tuple[float, dict[str, str]]] = {}
+
+
+def _cache_key_for_authorization(authorization: str) -> str:
+    return hashlib.sha256(authorization.strip().encode()).hexdigest()
+
+
+async def _refund_cache_get(authorization: str) -> dict[str, str] | None:
+    key = _cache_key_for_authorization(authorization)
+    async with _refund_cache_lock:
+        item = _refund_cache.get(key)
+        if item is None:
+            return None
+        expires_at, value = item
+        if expires_at <= monotonic():
+            del _refund_cache[key]
+            return None
+        return value
+
+
+async def _refund_cache_set(authorization: str, value: dict[str, str]) -> None:
+    key = _cache_key_for_authorization(authorization)
+    expiry = monotonic() + _REFUND_CACHE_TTL_SECONDS
+    async with _refund_cache_lock:
+        _refund_cache[key] = (expiry, value)
+
+
 @router.post("/refund")
 async def refund_wallet_endpoint(
-    key: ApiKey = Depends(get_key_from_header),
+    authorization: Annotated[str, Header(...)],
     session: AsyncSession = Depends(get_session),
-) -> dict:
+) -> dict[str, str]:
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid authorization. Use 'Bearer <cashu-token>' or 'Bearer <api-key>'",
+        )
+
+    bearer_value: str = authorization[7:]
+
+    if cached := await _refund_cache_get(bearer_value):
+        return cached
+
+    key: ApiKey = await validate_bearer_key(bearer_value, session)
     remaining_balance_msats: int = key.balance
 
     if remaining_balance_msats <= 0:
@@ -153,6 +198,8 @@ async def refund_wallet_endpoint(
         else:
             raise HTTPException(status_code=500, detail="Refund failed")
 
+    await _refund_cache_set(bearer_value, result)
+
     await session.delete(key)
     await session.commit()