add rate limiter to subscriptions

cashu/mint/events/client_manager.py

@@ -4,6 +4,8 @@
 from fastapi import WebSocket
 from loguru import logger
 
+from cashu.mint.limit import assert_limit
+
 from ...core.json_rpc.base import (
     JSONRPCError,
     JSONRPCErrorCode,
@@ -34,6 +36,23 @@ async def start(self):
         while True:
             json_data = await self.websocket.receive_text()
 
+            # Check the rate limit
+            try:
+                assert_limit(
+                    self.websocket.client.host if self.websocket.client else "unknown"
+                )
+            except Exception as e:
+                logger.error(f"Error: {e}")
+                resp = JSONRPCErrorResponse(
+                    error=JSONRPCError(
+                        code=JSONRPCErrorCode.SERVER_ERROR,
+                        message=f"Error: {e}",
+                    ),
+                    id=0,
+                )
+                await self._send_msg(resp)
+                continue
+
             # Parse the JSON data
             try:
                 data = json.loads(json_data)

cashu/mint/limit.py

@@ -1,5 +1,6 @@
 from fastapi import status
 from fastapi.responses import JSONResponse
+from limits import RateLimitItemPerMinute
 from loguru import logger
 from slowapi import Limiter
 from slowapi.util import get_remote_address
@@ -39,3 +40,16 @@ def get_remote_address_excluding_local(request: Request) -> str:
     default_limits=[f"{settings.mint_transaction_rate_limit_per_minute}/minute"],
     enabled=settings.mint_rate_limit,
 )
+
+
+def assert_limit(identifier: str):
+    global limiter
+    success = limiter._limiter.hit(
+        RateLimitItemPerMinute(settings.mint_transaction_rate_limit_per_minute),
+        identifier,
+    )
+    if not success:
+        logger.warning(
+            f"Rate limit {settings.mint_transaction_rate_limit_per_minute}/minute exceeded: {identifier}"
+        )
+        raise Exception("Rate limit exceeded")

cashu/mint/router.py

@@ -1,5 +1,4 @@
 from fastapi import APIRouter, Request, WebSocket
-from limits import RateLimitItemPerMinute
 from loguru import logger
 
 from cashu.mint.events.client_manager import LedgerEventClientManager
@@ -28,7 +27,7 @@
 from ..core.errors import KeysetNotFoundError
 from ..core.settings import settings
 from ..mint.startup import ledger
-from .limit import limiter
+from .limit import assert_limit, limiter
 
 router: APIRouter = APIRouter()
 
@@ -190,13 +189,7 @@ async def get_mint_quote(request: Request, quote: str) -> PostMintQuoteResponse:
 
 @router.websocket("/v1/ws", name="Websocket endpoint for subscriptions")
 async def websocket_endpoint(websocket: WebSocket):
-    success = limiter._limiter.hit(
-        RateLimitItemPerMinute(settings.mint_transaction_rate_limit_per_minute),
-        websocket.client.host if websocket.client else "unknown",
-    )
-    if not success:
-        await websocket.close(code=1008, reason="Rate limit exceeded")
-        return
+    assert_limit(websocket.client.host if websocket.client else "unknown")
     client = LedgerEventClientManager(websocket=websocket)
     success = ledger.events.add_client(client)
     if not success: