fix: key API rate budget only on a validated bearer, never fall back to client-chosen session id (#6215)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

src/praisonai-bot/praisonai_bot/gateway/api_endpoints.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 import asyncio
 import contextvars
 import hashlib
+import hmac
 import json
 import time
 import uuid
@@ -377,7 +378,7 @@ class GatewayApiEndpoints:
                 # Issue #6215: charge a stable caller identity, only when a
                 # budget is set (keeps duck-typed gates' signature intact).
                 rk = (
-                    {"rate_key": self._rate_key_var.get()}
+                    {"rate_key": self._rate_key_var.get() or "ip:unknown"}
                     if getattr(gate, "_rate_limit_enabled", False) else {}
                 )
                 async with gate.admit(session_id=session.session_id, **rk):
@@ -568,15 +569,33 @@ class GatewayApiEndpoints:
     )
 
     @staticmethod
-    def _rate_key(request: Any) -> Optional[str]:
-        """Stable caller identity that a new session header cannot reset."""
+    def _rate_key(request: Any, auth_token: Optional[str] = None) -> str:
+        """Stable caller identity that a client-chosen value cannot reset.
+
+        A bearer token is used only when it *is* the gateway's configured
+        ``auth_token`` (so it was validated); any other bearer value is
+        attacker-chosen and would mint a fresh budget per request, so such
+        callers are charged to their peer address. Never returns ``None``, so the
+        budget never falls back to the client-chosen session id.
+        """
         auth = request.headers.get("authorization", "")
-        if auth.startswith("Bearer ") and auth[7:]:
-            digest = hashlib.sha256(auth[7:].encode("utf-8")).hexdigest()[:16]
-            return "token:" + digest
+        bearer = auth[7:] if auth.startswith("Bearer ") else ""
+        if bearer and auth_token and hmac.compare_digest(
+            bearer.encode("utf-8"), str(auth_token).encode("utf-8")
+        ):
+            return "token:" + hashlib.sha256(bearer.encode("utf-8")).hexdigest()[:16]
         client = getattr(request, "client", None)
         host = getattr(client, "host", None)
-        return f"ip:{host}" if host else None
+        return f"ip:{host}" if host else "ip:unknown"
+
+    def _bind_rate_key(self, request: Any) -> None:
+        """Bind this request's per-user budget key for its turn (#6215).
+
+        Called by the auth-guarded route wrapper after auth has passed.
+        """
+        cfg = getattr(self._gw, "config", None)
+        token = getattr(cfg, "auth_token", None)
+        self._rate_key_var.set(self._rate_key(request, token))
 
     @staticmethod
     def _caller_key(request: Any) -> str:
@@ -592,10 +611,6 @@ class GatewayApiEndpoints:
         (Issue #2715 review, Greptile P2). Such callers are stateless by design;
         they should pass a session header to pin a conversation.
         """
-        # Issue #6215: the session header is client-chosen, so charge the
-        # per-user request budget to a stable caller identity (bearer token
-        # digest, else peer address) bound for this request's turn.
-        GatewayApiEndpoints._rate_key_var.set(GatewayApiEndpoints._rate_key(request))
         for header in ("openai-session", "x-session-id"):
             val = request.headers.get(header)
             if val:

src/praisonai-bot/praisonai_bot/gateway/server.py

@@ -2631,6 +2631,8 @@ class WebSocketGateway:
                     auth_err = _check_auth(request)
                     if auth_err:
                         return auth_err
+                    # Issue #6215: bind the per-user budget key post-auth.
+                    self._api_endpoints._bind_rate_key(request)
                     return await handler(request)
                 return _wrapped
 
@@ -5901,16 +5903,17 @@ class WebSocketGateway:
     
     # ── Client Management (Public API) ────────────────────────────────
     
-    def _client_rate_key(self, client_id: Optional[str]) -> Optional[str]:
+    def _client_rate_key(self, client_id: Optional[str]) -> str:
         """Stable per-user budget key for a WebSocket turn (#6215).
 
-        Session ids are client-chosen, so charge the peer address instead;
-        ``None`` falls back to the session id (e.g. no live connection).
+        Session ids are client-chosen, so charge the peer address instead.
+        With no known peer, callers share one ``ip:unknown`` budget rather
+        than falling back to the client-chosen session id.
         """
         ws = self._clients.get(client_id) if client_id else None
         peer = getattr(ws, "client", None)
         host = getattr(peer, "host", None)
-        return f"ip:{host}" if host else None
+        return f"ip:{host}" if host else "ip:unknown"
 
     def add_client(self, client_id: str, websocket: Any) -> None:
         """Register a client WebSocket connection.

src/praisonai-bot/tests/unit/test_admission_rate_limit.py

@@ -177,5 +177,64 @@ def test_api_rate_key_ignores_session_header():
     a = GatewayApiEndpoints._rate_key(_Req({"x-session-id": "one"}))
     b = GatewayApiEndpoints._rate_key(_Req({"x-session-id": "two"}))
     assert a == b == "ip:10.0.0.1"
-    tok = GatewayApiEndpoints._rate_key(_Req({"authorization": "Bearer secret-token"}))
+    tok = GatewayApiEndpoints._rate_key(
+        _Req({"authorization": "Bearer secret-token"}), "secret-token"
+    )
     assert tok.startswith("token:") and "secret" not in tok
+
+
+def test_api_rate_key_rotating_unvalidated_bearer_shares_ip_budget():
+    """An unvalidated bearer (auth open, or wrong token) cannot mint budgets."""
+    try:
+        from praisonai_bot.gateway.api_endpoints import GatewayApiEndpoints
+    except ModuleNotFoundError:  # pragma: no cover - optional deps
+        pytest.skip("gateway optional dependency not installed")
+
+    class _Req:
+        def __init__(self, headers, host="10.0.0.1"):
+            self.headers = headers
+            self.client = type("C", (), {"host": host})()
+
+    for configured in (None, "", "real-token"):
+        keys = {
+            GatewayApiEndpoints._rate_key(
+                _Req({"authorization": f"Bearer rotate-{i}"}), configured
+            )
+            for i in range(5)
+        }
+        assert keys == {"ip:10.0.0.1"}
+
+
+def test_api_rate_key_missing_peer_never_falls_back_to_session():
+    try:
+        from praisonai_bot.gateway.api_endpoints import GatewayApiEndpoints
+    except ModuleNotFoundError:  # pragma: no cover - optional deps
+        pytest.skip("gateway optional dependency not installed")
+
+    class _Req:
+        headers = {"x-session-id": "chosen"}
+        client = None
+
+    assert GatewayApiEndpoints._rate_key(_Req()) == "ip:unknown"
+
+    gate = build_admission_gate(per_user_per_minute=1)
+
+    async def run():
+        async with gate.admit(session_id="conv-1", rate_key="ip:unknown"):
+            pass
+        with pytest.raises(AdmissionRejected):
+            async with gate.admit(session_id="conv-2", rate_key="ip:unknown"):
+                pass
+
+    asyncio.run(run())
+
+
+def test_ws_rate_key_missing_peer_is_shared():
+    try:
+        from praisonai_bot.gateway.server import WebSocketGateway
+    except ModuleNotFoundError:  # pragma: no cover - optional deps
+        pytest.skip("gateway optional dependency not installed")
+    gw = WebSocketGateway.__new__(WebSocketGateway)
+    gw._clients = {}
+    assert gw._client_rate_key(None) == "ip:unknown"
+    assert gw._client_rate_key("gone") == "ip:unknown"