From 801385df44ed4f60c95132fc1b2e2cd4132d97aa Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 18 Apr 2026 11:07:38 -0400 Subject: [PATCH] fix: bound _refresh_locks with LRU eviction to prevent memory leak (#3968) --- src/fastmcp/server/auth/oauth_proxy/proxy.py | 21 ++++++++++--- tests/server/auth/oauth_proxy/test_tokens.py | 32 ++++++++++++++++++++ 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/src/fastmcp/server/auth/oauth_proxy/proxy.py b/src/fastmcp/server/auth/oauth_proxy/proxy.py index fb3079dc0..9471ccff7 100644 --- a/src/fastmcp/server/auth/oauth_proxy/proxy.py +++ b/src/fastmcp/server/auth/oauth_proxy/proxy.py @@ -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 diff --git a/tests/server/auth/oauth_proxy/test_tokens.py b/tests/server/auth/oauth_proxy/test_tokens.py index 9e820589e..e802b2ade 100644 --- a/tests/server/auth/oauth_proxy/test_tokens.py +++ b/tests/server/auth/oauth_proxy/test_tokens.py @@ -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"}