fix: bound _refresh_locks with LRU eviction to prevent memory leak (#3968)

This commit is contained in:
Jeremiah Lowin 2026-04-18 11:07:38 -04:00 committed by GitHub
commit 801385df44
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 49 additions and 4 deletions

View file

@ -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

View file

@ -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"}