mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 13:34:17 +02:00
fix: bound _refresh_locks with LRU eviction to prevent memory leak (#3968)
This commit is contained in:
parent
64fbc52e1d
commit
801385df44
2 changed files with 49 additions and 4 deletions
|
|
@ -22,6 +22,7 @@ import hashlib
|
|||
import secrets
|
||||
import time
|
||||
from base64 import urlsafe_b64encode
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import urlencode, urlparse, urlunparse
|
||||
|
||||
|
|
@ -92,6 +93,8 @@ from fastmcp.utilities.logging import get_logger
|
|||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_REFRESH_LOCK_CACHE_SIZE = 10_000
|
||||
|
||||
|
||||
def _normalize_resource_url(url: str) -> str:
|
||||
"""Normalize a resource URL by removing query parameters and trailing slashes.
|
||||
|
|
@ -576,7 +579,7 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
|
|||
# refresh the same token within a single process. Does not protect
|
||||
# against cross-process races in distributed deployments — those are
|
||||
# handled by re-reading from storage after refresh failure.
|
||||
self._refresh_locks: dict[str, anyio.Lock] = {}
|
||||
self._refresh_locks: OrderedDict[str, anyio.Lock] = OrderedDict()
|
||||
|
||||
logger.debug(
|
||||
"Initialized OAuth proxy provider with upstream server %s",
|
||||
|
|
@ -647,6 +650,18 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
|
|||
timeout=HTTP_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
def _get_refresh_lock(self, token_id: str) -> anyio.Lock:
|
||||
"""Get or create a per-token refresh lock, evicting LRU entries when at capacity."""
|
||||
lock = self._refresh_locks.get(token_id)
|
||||
if lock is None:
|
||||
lock = anyio.Lock()
|
||||
self._refresh_locks[token_id] = lock
|
||||
if len(self._refresh_locks) > _REFRESH_LOCK_CACHE_SIZE:
|
||||
self._refresh_locks.popitem(last=False)
|
||||
else:
|
||||
self._refresh_locks.move_to_end(token_id)
|
||||
return lock
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# PKCE Helper Methods
|
||||
# -------------------------------------------------------------------------
|
||||
|
|
@ -1656,9 +1671,7 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
|
|||
|
||||
# Advisory lock prevents concurrent requests from racing
|
||||
# to refresh the same upstream token.
|
||||
if token_id not in self._refresh_locks:
|
||||
self._refresh_locks[token_id] = anyio.Lock()
|
||||
lock = self._refresh_locks[token_id]
|
||||
lock = self._get_refresh_lock(token_id)
|
||||
|
||||
async with lock:
|
||||
# Re-read from storage — another task may have
|
||||
|
|
|
|||
|
|
@ -907,6 +907,38 @@ class TestTransparentUpstreamRefresh:
|
|||
assert result is not None
|
||||
assert result.token == "refreshed-upstream-access"
|
||||
|
||||
def test_refresh_lock_cache_bounded(self, proxy, monkeypatch):
|
||||
"""_get_refresh_lock never grows beyond _REFRESH_LOCK_CACHE_SIZE."""
|
||||
monkeypatch.setattr(
|
||||
"fastmcp.server.auth.oauth_proxy.proxy._REFRESH_LOCK_CACHE_SIZE", 3
|
||||
)
|
||||
for i in range(10):
|
||||
proxy._get_refresh_lock(f"token-{i}")
|
||||
assert len(proxy._refresh_locks) == 3
|
||||
|
||||
def test_refresh_lock_lru_evicts_least_recently_used(self, proxy, monkeypatch):
|
||||
"""Touching an entry promotes it; eviction removes the oldest untouched entry."""
|
||||
monkeypatch.setattr(
|
||||
"fastmcp.server.auth.oauth_proxy.proxy._REFRESH_LOCK_CACHE_SIZE", 3
|
||||
)
|
||||
proxy._get_refresh_lock("a")
|
||||
proxy._get_refresh_lock("b")
|
||||
proxy._get_refresh_lock("c")
|
||||
# Touch "a" to move it to MRU position
|
||||
proxy._get_refresh_lock("a")
|
||||
# Adding "d" should evict "b" (oldest untouched)
|
||||
proxy._get_refresh_lock("d")
|
||||
assert "b" not in proxy._refresh_locks
|
||||
assert "a" in proxy._refresh_locks
|
||||
assert "c" in proxy._refresh_locks
|
||||
assert "d" in proxy._refresh_locks
|
||||
|
||||
def test_refresh_lock_same_token_returns_same_lock(self, proxy):
|
||||
"""Requesting the same token ID twice returns the same lock object."""
|
||||
lock1 = proxy._get_refresh_lock("tok")
|
||||
lock2 = proxy._get_refresh_lock("tok")
|
||||
assert lock1 is lock2
|
||||
|
||||
async def test_upstream_claims_propagated(self, proxy):
|
||||
jwt = await self._setup_session_with_claims(
|
||||
proxy, upstream_claims={"sub": "user-123"}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue