From 51001bb31a53ff94a742e599d54726bec4925055 Mon Sep 17 00:00:00 2001 From: Shroominic Date: Sat, 2 Aug 2025 22:08:04 -0300 Subject: [PATCH] simplify deserialize_token_from_string --- router/payment/helpers.py | 188 +++----------------------------------- router/proxy.py | 13 +-- 2 files changed, 16 insertions(+), 185 deletions(-) diff --git a/router/payment/helpers.py b/router/payment/helpers.py index 5c265ea5..d7d0a5bb 100644 --- a/router/payment/helpers.py +++ b/router/payment/helpers.py @@ -1,12 +1,10 @@ -import base64 import json import os -from typing import Literal -import cbor2 from fastapi import HTTPException, Response from ..core import get_logger +from ..wallet import deserialize_token_from_string from .cost_caculation import COST_PER_REQUEST, MODEL_BASED_PRICING from .models import MODELS @@ -49,17 +47,7 @@ def get_cost_per_request(model: str | None = None) -> int: return COST_PER_REQUEST -def check_token_balance(headers: dict, body: dict) -> Literal["sat", "msat"]: - """Check if the provided token has sufficient balance.""" - logger.debug( - "Checking token balance", - extra={ - "has_x_cashu": "x-cashu" in headers, - "has_authorization": "authorization" in headers, - "model": body.get("model", "unknown"), - }, - ) - +def check_token_balance(headers: dict, body: dict) -> None: if x_cashu := headers.get("x-cashu", None): cashu_token = x_cashu logger.debug( @@ -100,173 +88,27 @@ def check_token_balance(headers: dict, body: dict) -> Literal["sat", "msat"]: # Handle regular API keys (sk-*) if cashu_token.startswith("sk-"): - logger.debug( - "Regular API key detected", extra={"key_preview": cashu_token[:10] + "..."} - ) - return "sat" + return cost = get_cost_per_request(model=body.get("model", None)) - if cashu_token.startswith("cashuA"): - logger.debug("Processing CashuA token", extra={"required_cost_msats": cost}) + token_obj = deserialize_token_from_string(cashu_token) - try: - _token = base64_token_json(cashu_token) - amount = sum(p["amount"] for t in _token["token"] for p in t["proofs"]) - unit: Literal["sat", "msat"] = _token.get("unit", "sat") + amount_msat = ( + token_obj.amount if token_obj.unit == "msat" else token_obj.amount * 1000 + ) - if unit == "sat": - amount *= 1000 - - logger.info( - "CashuA token parsed successfully", - extra={ - "amount": amount, - "unit": unit, - "amount_msats": amount, - "required_cost_msats": cost, - "sufficient_balance": amount >= cost, - }, - ) - - if amount < cost: - logger.warning( - "Insufficient token balance", - extra={ - "amount_msats": amount, - "required_msats": cost, - "shortfall_msats": cost - amount, - "unit": unit, - }, - ) - raise HTTPException(status_code=413, detail="Insufficient balance") - - except Exception as e: - logger.error( - "Failed to parse CashuA token", - extra={ - "error": str(e), - "error_type": type(e).__name__, - "token_preview": cashu_token[:20] + "...", - }, - ) - raise HTTPException(status_code=401, detail="Invalid token format") - - elif cashu_token.startswith("cashuB"): - logger.debug("Processing CashuB token", extra={"required_cost_msats": cost}) - - try: - _token = base64_token_cbor(cashu_token) - amount = sum(p["a"] for t in _token["t"] for p in t["p"]) - unit = _token["u"] - - if unit == "sat": - amount *= 1000 - - logger.info( - "CashuB token parsed successfully", - extra={ - "amount": amount, - "unit": unit, - "amount_msats": amount, - "required_cost_msats": cost, - "sufficient_balance": amount >= cost, - }, - ) - - if amount < cost: - logger.warning( - "Insufficient token balance", - extra={ - "amount_msats": amount, - "required_msats": cost, - "shortfall_msats": cost - amount, - "unit": unit, - }, - ) - raise HTTPException(status_code=413, detail="Insufficient balance") - - except Exception as e: - logger.error( - "Failed to parse CashuB token", - extra={ - "error": str(e), - "error_type": type(e).__name__, - "token_preview": cashu_token[:20] + "...", - }, - ) - raise HTTPException(status_code=401, detail="Invalid token format") - - else: - logger.error( - "Unknown token format", - extra={"token_prefix": cashu_token[:10] if cashu_token else "empty"}, - ) - raise HTTPException(status_code=401, detail="Unauthorized") - - return unit - - -def base64_token_json(cashu_token: str) -> dict: - """Decode a CashuA (JSON) token.""" - logger.debug("Decoding CashuA token", extra={"token_length": len(cashu_token)}) - - try: - # Version 3 - JSON format - encoded = cashu_token[6:] # Remove "cashuA" - # Add correct padding – (-len) % 4 equals 0,1,2,3 - encoded += "=" * ((-len(encoded)) % 4) - - decoded = base64.urlsafe_b64decode(encoded).decode() - token_data = json.loads(decoded) - - logger.debug( - "CashuA token decoded successfully", - extra={ - "token_proofs_count": sum( - len(t.get("proofs", [])) for t in token_data.get("token", []) - ), - "unit": token_data.get("unit", "unknown"), + if cost > amount_msat: + raise HTTPException( + status_code=413, + detail={ + "reason": "Insufficient balance", + "amount_required_msat": cost, + "model": body.get("model", "unknown"), + "type": "minimum_balance_required", }, ) - return token_data - except Exception as e: - logger.error( - "Failed to decode CashuA token", - extra={"error": str(e), "error_type": type(e).__name__}, - ) - raise - - -def base64_token_cbor(cashu_token: str) -> dict: - """Decode a CashuB (CBOR) token.""" - logger.debug("Decoding CashuB token", extra={"token_length": len(cashu_token)}) - - try: - encoded = cashu_token[6:] # Remove "cashuB" - encoded += "=" * ((-len(encoded)) % 4) - decoded_bytes = base64.urlsafe_b64decode(encoded) - token_data = cbor2.loads(decoded_bytes) - - logger.debug( - "CashuB token decoded successfully", - extra={ - "token_proofs_count": sum( - len(t.get("p", [])) for t in token_data.get("t", []) - ), - "unit": token_data.get("u", "unknown"), - }, - ) - - return token_data - except Exception as e: - logger.error( - "Failed to decode CashuB token", - extra={"error": str(e), "error_type": type(e).__name__}, - ) - raise - def get_max_cost_for_model(model: str) -> int: """Get the maximum cost for a specific model.""" diff --git a/router/proxy.py b/router/proxy.py index d89afd8f..e588c4f5 100644 --- a/router/proxy.py +++ b/router/proxy.py @@ -479,18 +479,7 @@ async def proxy( media_type="application/json", ) - # Check token balance for all requests to get currency unit - try: - unit = check_token_balance(headers, request_body_dict) - logger.debug( - "Token balance check completed", extra={"path": path, "unit": unit} - ) - except HTTPException as e: - logger.warning( - "Token balance check failed", - extra={"path": path, "status_code": e.status_code, "detail": str(e.detail)}, - ) - raise + check_token_balance(headers, request_body_dict) # Handle authentication if x_cashu := headers.get("x-cashu", None):