From 95a0f9c836b1a82d96fd4cbf5b25fcb4aa22e921 Mon Sep 17 00:00:00 2001 From: teknium1 <127238744+teknium1@users.noreply.github.com> Date: Thu, 16 Jul 2026 05:02:35 -0700 Subject: [PATCH] fix(mcp): close select-to-bind TOCTOU on the OAuth callback port MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _find_free_port() closed its probe socket before HTTPServer re-bound the port minutes later, leaving a window where another process could steal it (#22161 by @amathxbt). _reserve_callback_port() now keeps the selected socket bound (bounded FIFO pool) until _wait_for_callback adopts it via bind_and_activate=False. Also sets allow_reuse_address BEFORE binding — the cherry-picked #44872 set it after the constructor had already bound, where it is a no-op. Also updates the three #57836 non-interactive-guard tests to the closure-factory API from #44872. --- tests/tools/test_mcp_oauth.py | 103 +++++++++++++++++++++++++++++++++- tools/mcp_oauth.py | 70 ++++++++++++++++++++--- 2 files changed, 163 insertions(+), 10 deletions(-) diff --git a/tests/tools/test_mcp_oauth.py b/tests/tools/test_mcp_oauth.py index 547e31afde40..7ff4232f80aa 100644 --- a/tests/tools/test_mcp_oauth.py +++ b/tests/tools/test_mcp_oauth.py @@ -423,6 +423,103 @@ class TestOAuthPortSharing: assert 1024 <= mod._oauth_port <= 65535 +# --------------------------------------------------------------------------- +# TOCTOU port reservation (#22161) +# --------------------------------------------------------------------------- + +class TestCallbackPortReservation: + """The socket picked at selection time stays bound until callback bind. + + _find_free_port() closed its probe socket before HTTPServer re-bound the + port, leaving a race window where another process could steal it + (#22161). _reserve_callback_port() keeps the bound socket parked in + _reserved_sockets until _wait_for_callback adopts it. + """ + + def test_reserved_port_cannot_be_stolen(self): + import socket as sock + import tools.mcp_oauth as mod + + port = mod._reserve_callback_port() + try: + # The reservation holds the bind — a competing bind must fail. + thief = sock.socket(sock.AF_INET, sock.SOCK_STREAM) + with pytest.raises(OSError): + thief.bind(("127.0.0.1", port)) + thief.close() + finally: + reserved = mod._reserved_sockets.pop(port, None) + if reserved is not None: + reserved.close() + + def test_configure_callback_port_reserves_ephemeral(self): + import tools.mcp_oauth as mod + + cfg: dict = {} + port = mod._configure_callback_port(cfg) + try: + assert cfg["_resolved_port"] == port + assert port in mod._reserved_sockets + finally: + reserved = mod._reserved_sockets.pop(port, None) + if reserved is not None: + reserved.close() + + def test_pinned_port_is_not_reserved(self): + import tools.mcp_oauth as mod + + cfg: dict = {"redirect_port": 49399} + port = mod._configure_callback_port(cfg) + assert port == 49399 + assert 49399 not in mod._reserved_sockets + + def test_reservation_pool_is_bounded(self): + import tools.mcp_oauth as mod + + ports = [mod._reserve_callback_port() for _ in range(mod._MAX_RESERVED_SOCKETS + 3)] + try: + assert len(mod._reserved_sockets) <= mod._MAX_RESERVED_SOCKETS + # newest reservations survive + assert ports[-1] in mod._reserved_sockets + finally: + for p in list(mod._reserved_sockets): + mod._reserved_sockets.pop(p).close() + + def test_wait_for_callback_adopts_reserved_socket(self, monkeypatch): + """E2E: reserve → _wait_for_callback binds the SAME socket and the + callback round-trips through it.""" + import asyncio + import threading + import urllib.request + import tools.mcp_oauth as mod + + cfg: dict = {} + port = mod._configure_callback_port(cfg) + monkeypatch.setattr(mod, "_is_interactive", lambda: False) + # Bypass the non-interactive guard — this test drives the flow directly. + monkeypatch.setattr(mod, "_raise_if_non_interactive", lambda lead: None) + + async def drive(): + task = asyncio.create_task(mod._wait_for_callback()) + await asyncio.sleep(0.2) # let the server adopt the socket + + def hit(): + urllib.request.urlopen( + f"http://127.0.0.1:{port}/callback?code=abc123&state=xyz", + timeout=5, + ) + + t = threading.Thread(target=hit, daemon=True) + t.start() + return await asyncio.wait_for(task, timeout=10) + + code, state = asyncio.run(drive()) + assert code == "abc123" + assert state == "xyz" + # Reservation was consumed by adoption. + assert port not in mod._reserved_sockets + + # --------------------------------------------------------------------------- # remove_oauth_tokens # --------------------------------------------------------------------------- @@ -661,7 +758,7 @@ class TestNonInteractiveFailFastAtCallbackBoundary: ) with pytest.raises(OAuthNonInteractiveError, match="browser authorization"): - asyncio.run(mod._redirect_handler("https://idp.example.com/authorize?x=1")) + asyncio.run(mod._make_redirect_handler(49300)("https://idp.example.com/authorize?x=1")) err = capsys.readouterr().err assert "https://idp.example.com/authorize" not in err @@ -673,7 +770,7 @@ class TestNonInteractiveFailFastAtCallbackBoundary: monkeypatch.setattr(mod, "_is_interactive", lambda: False) with pytest.raises(OAuthNonInteractiveError, match="hermes mcp login"): - asyncio.run(mod._redirect_handler("https://idp.example.com/authorize")) + asyncio.run(mod._make_redirect_handler(49301)("https://idp.example.com/authorize")) mod._oauth_port = _find_free_port() with pytest.raises(OAuthNonInteractiveError, match="hermes mcp login"): @@ -698,7 +795,7 @@ class TestNonInteractiveFailFastAtCallbackBoundary: monkeypatch.delenv("SSH_TTY", raising=False) monkeypatch.setattr(mod, "_can_open_browser", lambda: False) - asyncio.run(mod._redirect_handler("https://idp.example.com/authorize?x=9")) + asyncio.run(mod._make_redirect_handler(49302)("https://idp.example.com/authorize?x=9")) err = capsys.readouterr().err assert "https://idp.example.com/authorize?x=9" in err diff --git a/tools/mcp_oauth.py b/tools/mcp_oauth.py index a5946544dec2..9003f980fcd7 100644 --- a/tools/mcp_oauth.py +++ b/tools/mcp_oauth.py @@ -155,6 +155,43 @@ def _find_free_port() -> int: return s.getsockname()[1] +# Bound-but-not-listening sockets reserved for pending OAuth callback flows, +# keyed by port. Holding the socket from port-selection time until +# _wait_for_callback adopts it closes the TOCTOU window where another process +# could grab the port between _find_free_port() closing its probe socket and +# HTTPServer binding minutes later (#22161). Bounded FIFO so repeated +# build_oauth_auth calls (reconnect loops) cannot leak fds. +_reserved_sockets: "dict[int, socket.socket]" = {} +_MAX_RESERVED_SOCKETS = 8 + + +def _reserve_callback_port() -> int: + """Pick an ephemeral callback port and keep its socket bound. + + Returns the port. The bound (not yet listening) socket is parked in + ``_reserved_sockets`` so no other process can bind the port before + ``_wait_for_callback`` adopts it. Adoption (or ``server_close``) owns + the socket's lifetime from there. + """ + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + s.bind(("127.0.0.1", 0)) + except OSError: + s.close() + raise + port = s.getsockname()[1] + # Evict oldest reservations past the cap (dict preserves insertion order). + while len(_reserved_sockets) >= _MAX_RESERVED_SOCKETS: + _, stale = next(iter(_reserved_sockets.items())) + _reserved_sockets.pop(next(iter(_reserved_sockets)), None) + try: + stale.close() + except OSError: + pass + _reserved_sockets[port] = s + return port + + def _is_interactive() -> bool: """Return True if we can reasonably expect to interact with a user.""" if not _oauth_interactive_enabled.get(): @@ -659,13 +696,29 @@ async def _wait_for_callback() -> tuple[str, str | None]: # We just need to poll for the result. handler_cls, result = _make_callback_handler() - # Start a temporary server on the known port. - # allow_reuse_address prevents TIME_WAIT on the socket after the flow - # completes, so the next OAuth flow on the same port can bind immediately - # (fixes #44590). + # Start a temporary server on the known port, adopting the socket + # reserved at port-selection time when one exists. Holding the bound + # socket from _reserve_callback_port() until here closes the TOCTOU + # window where another process could steal the port between selection + # and bind (#22161). allow_reuse_address is set BEFORE binding (setting + # it after the constructor has already bound is a no-op) so a lingering + # TIME_WAIT socket from a previous flow cannot block the next one + # (#44590). try: - server = HTTPServer(("127.0.0.1", _oauth_port), handler_cls) - server.allow_reuse_address = True + server = HTTPServer( + ("127.0.0.1", _oauth_port), handler_cls, bind_and_activate=False + ) + reserved = _reserved_sockets.pop(_oauth_port, None) + if reserved is not None: + # Adopt the reserved (already bound) socket and start listening. + server.socket.close() + server.socket = reserved + server.server_address = reserved.getsockname() + server.server_activate() + else: + server.allow_reuse_address = True + server.server_bind() + server.server_activate() except OSError: # Port already in use — the server from build_oauth_auth is running. # Fall back to polling the server started by build_oauth_auth. @@ -840,7 +893,10 @@ def _configure_callback_port(cfg: dict) -> int: """ global _oauth_port requested = int(cfg.get("redirect_port", 0)) - port = _find_free_port() if requested == 0 else requested + # Ephemeral selection reserves the bound socket until _wait_for_callback + # adopts it, closing the select→bind TOCTOU race (#22161). An explicit + # user-pinned port is used as-is. + port = _reserve_callback_port() if requested == 0 else requested cfg["_resolved_port"] = port _oauth_port = port # legacy consumer: _wait_for_callback reads this return port