From 67f7bc5cb09e64d568050d444abf7519ecb5fb0f Mon Sep 17 00:00:00 2001 From: RerankerGuo Date: Tue, 13 Oct 2026 16:54:31 +0800 Subject: [PATCH] fix(api): strip Bearer scheme so keyed requests reach per-key rate-limit bucket (#2428) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _get_client_key in src/memos/api/middleware/rate_limit.py tested whether request.headers['Authorization'] started with 'krlk_', but real clients send 'Authorization: Bearer krlk_' per RFC 6750. The prefix check never matched and every request — including ones carrying a valid API key — fell through to IP-based limiting. Behind a shared egress IP all keyed clients shared one bucket, so the per-key rate-limit was dead. The fix strips the 'Bearer ' scheme (case-insensitive, per RFC 7235) before the prefix check, so 'Bearer krlk_' correctly maps to 'ratelimit:key:'. Backwards compatibility is preserved: a bare 'krlk_' header still matches. - New tests/api/test_rate_limit_bearer.py with 7 cases covering Bearer prefix, lowercase bearer, non-key Bearer tokens, missing header, X-Forwarded-For preference, and the bare-krlk_ backwards-compat path. - No new imports, no new dependencies, no public API change. Fixes #2428. --- src/memos/api/middleware/rate_limit.py | 7 ++- tests/api/test_rate_limit_bearer.py | 78 ++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 tests/api/test_rate_limit_bearer.py diff --git a/src/memos/api/middleware/rate_limit.py b/src/memos/api/middleware/rate_limit.py index c547378ca..72808a13d 100644 --- a/src/memos/api/middleware/rate_limit.py +++ b/src/memos/api/middleware/rate_limit.py @@ -57,8 +57,13 @@ def _get_client_key(request: Request) -> str: Uses API key if available, otherwise falls back to IP. """ - # Try to get API key from header + # Try to get API key from header. Real clients send a Bearer scheme + # per RFC 6750 (``Authorization: Bearer krlk_``); we strip the + # scheme before the prefix check so keyed requests actually land on + # a per-key bucket (issue #2428). auth_header = request.headers.get("Authorization", "") + if auth_header.lower().startswith("bearer "): + auth_header = auth_header[len("bearer ") :].lstrip() if auth_header.startswith("krlk_"): # Use first 20 chars of key as identifier return f"ratelimit:key:{auth_header[:20]}" diff --git a/tests/api/test_rate_limit_bearer.py b/tests/api/test_rate_limit_bearer.py new file mode 100644 index 000000000..710346cf0 --- /dev/null +++ b/tests/api/test_rate_limit_bearer.py @@ -0,0 +1,78 @@ +"""Regression tests for issue #2428. + +``_get_client_key`` in ``src/memos/api/middleware/rate_limit.py`` checked +whether ``request.headers['Authorization']`` started with ``krlk_`` — but +real clients send ``Authorization: Bearer krlk_`` per RFC 6750, so +the prefix check never matched and every request fell through to IP-based +limiting. Behind a shared egress IP all keyed clients shared one bucket +and the per-key rate-limit was dead code. + +The fix strips the ``Bearer `` scheme before testing the key prefix, so +``Bearer krlk_`` correctly maps to ``ratelimit:key:``. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from memos.api.middleware.rate_limit import _get_client_key + + +def _fake_request(authorization: str | None, *, client_host: str = "10.0.0.1"): + headers: dict[str, str] = {} + if authorization is not None: + headers["Authorization"] = authorization + return SimpleNamespace( + headers=headers, + client=SimpleNamespace(host=client_host), + ) + + +class TestGetClientKey: + def test_bearer_krlk_key_uses_per_key_bucket(self): + body = "krlk_" + "a" * 64 + request = _fake_request("Bearer " + body) + key = _get_client_key(request) + assert key == "ratelimit:key:" + body[:20] + assert "10.0.0.1" not in key + + def test_bearer_krlk_key_first_twenty_chars_of_actual_key(self): + body = "krlk_" + "f" * 64 + request = _fake_request("Bearer " + body) + key = _get_client_key(request) + assert key == "ratelimit:key:" + body[:20] + + def test_lowercase_bearer_scheme(self): + # RFC 7235: the scheme name is case-insensitive. We normalise to + # lowercase before the prefix check. + body = "krlk_" + "b" * 64 + request = _fake_request("bearer " + body) + key = _get_client_key(request) + assert key == "ratelimit:key:" + body[:20] + + def test_bearer_with_non_key_token_falls_back_to_ip(self): + request = _fake_request("Bearer some-jwt-token") + key = _get_client_key(request) + assert key.startswith("ratelimit:ip:") + assert "10.0.0.1" in key + + def test_missing_authorization_header_falls_back_to_ip(self): + request = _fake_request(None) + key = _get_client_key(request) + assert key.startswith("ratelimit:ip:") + assert "10.0.0.1" in key + + def test_x_forwarded_for_header_preferred_over_ip(self): + request = _fake_request(None, client_host="203.0.113.1") + request.headers["X-Forwarded-For"] = "198.51.100.1, 10.0.0.1" + key = _get_client_key(request) + assert key == "ratelimit:ip:198.51.100.1" + + def test_krlk_prefix_without_bearer_scheme_still_matches(self): + # Backwards compatibility: clients that sent a bare krlk_ header + # (before the Bearer scheme was added) must still match. + body = "krlk_" + "c" * 64 + request = _fake_request(body) + key = _get_client_key(request) + assert key == "ratelimit:key:" + body[:20]