Reshape the stream claim into an explicit state machine

maxisbey · maxisbey · commit 41b9faaf9b96 · 2026-07-21T21:14:05.000Z
The claim registry maps each standard descriptor to at most one owning
transport; the wire duplicate is allocated where it cannot land in the
standard range (F_DUPFD_CLOEXEC above fd 2 on POSIX) and recorded on
the claim before the descriptor is moved; release restores with a
single dup2 and deregisters only on success. Every failure, modeled or
not, lands on the safe side: the claim is retained and later
transports are refused rather than handed a diverted descriptor.

Deleted by the same design: the roll-forward path (its premise, a
reliably reportable dup2 outcome, does not exist on Windows), and the
restore-time flush (user flush() side effects can destroy the wire;
unflushed stray output now drains after the session instead). The
design was validated through three adversarial verification passes;
the write-up with invariants and accepted residues is on the PR.
diff --git a/src/mcp/server/stdio.py b/src/mcp/server/stdio.py
@@ -16,6 +16,7 @@ async def run_server():
 import threading
 from collections.abc import Callable
 from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass
 from io import TextIOWrapper
 from typing import BinaryIO, Literal, TextIO
 
@@ -27,9 +28,25 @@ async def run_server():
 from mcp.shared._context_streams import create_context_streams
 from mcp.shared.message import SessionMessage
 
-# fds whose real stream a serving transport owns, whether diverted or served in place.
-_claimed: set[int] = set()
-_claimed_lock = threading.Lock()
+if sys.platform != "win32":  # pragma: no branch
+    import fcntl
+
+# Stream-claim contract (design and attack log in PR #3117):
+# - _claims is the single authority for who owns fd 0/1; mutated only under the
+#   lock, only by acquire's insert and release's deregister.
+# - private_fd is recorded the instant the wire duplicate exists, before fd is
+#   ever moved, and is never closed while the claim is registered.
+# - Release deregisters only after dup2(private_fd, fd) restores the wire; a
+#   failed release keeps the claim, so successors are refused, never fed a
+#   diverted descriptor. Every failure lands on that safe side.
+_claims: dict[int, "_StreamClaim"] = {}
+_claims_lock = threading.Lock()
+
+
+@dataclass
+class _StreamClaim:
+    fd: int
+    private_fd: int | None = None
 
 
 def _is_backed_by_fd(stream: TextIO, fd: int) -> bool:
@@ -39,14 +56,15 @@ def _is_backed_by_fd(stream: TextIO, fd: int) -> bool:
         return False
 
 
-def _std_descriptors_open() -> bool:
-    """Whether fds 0-2 are all open, so a dup cannot land in the standard range."""
-    try:
-        for fd in (0, 1, 2):
-            os.fstat(fd)
-    except OSError:
-        return False
-    return True
+def _dup_above_std(fd: int) -> int:
+    """Duplicate fd onto a descriptor that cannot land in the standard range."""
+    if sys.platform == "win32":  # pragma: no cover
+        duplicate = os.dup(fd)
+        if duplicate <= 2:
+            os.close(duplicate)
+            raise OSError(f"duplicate of fd {fd} landed in the standard range")
+        return duplicate
+    return fcntl.fcntl(fd, fcntl.F_DUPFD_CLOEXEC, 3)
 
 
 def _open_stdin_diversion() -> int:
@@ -60,68 +78,69 @@ def _open_stdout_diversion() -> int:
         return os.open(os.devnull, os.O_WRONLY)
 
 
+def _restore_fd(fd: int, private_fd: int) -> bool:
+    """Point fd back at the wire; the Windows handle rebind never affects the outcome."""
+    try:
+        os.dup2(private_fd, fd)
+    except OSError:
+        return False
+    if sys.platform == "win32":  # pragma: no cover
+        with suppress(OSError):
+            rebind_std_handle_to_fd(fd)
+    return True
+
+
 def _claim_fd(
     fd: int, stream: TextIO, mode: Literal["rb", "wb"], open_diversion: Callable[[], int]
 ) -> tuple[BinaryIO, Callable[[], None] | None]:
-    """Move the protocol pipe to a private descriptor and divert fd while serving.
+    """Claim a standard stream: divert fd and serve the wire from a private duplicate.
+
+    Best-effort: when descriptors cannot be duplicated or diverted, serves the
+    sys stream's buffer in place, exactly as v1 did, with the claim held.
 
     Raises:
         RuntimeError: fd is already claimed by another transport in this process.
     """
     if not _is_backed_by_fd(stream, fd):
         return stream.buffer, None
-    with _claimed_lock:
-        if fd in _claimed:
+    claim = _StreamClaim(fd)
+    with _claims_lock:
+        if fd in _claims:
             raise RuntimeError(f"another stdio_server() in this process has already claimed fd {fd}")
-        _claimed.add(fd)
+        _claims[fd] = claim
 
-    def unclaim() -> None:
-        with _claimed_lock:
-            _claimed.discard(fd)
+    def release() -> None:
+        if claim.private_fd is None or _restore_fd(fd, claim.private_fd):
+            with _claims_lock:
+                del _claims[fd]
 
-    if not _std_descriptors_open():
-        return stream.buffer, unclaim
-    private_fd = None
     try:
-        private_fd = os.dup(fd)
-        diversion_fd = open_diversion()
-        try:
-            os.dup2(diversion_fd, fd)
-        finally:
-            os.close(diversion_fd)
-        if sys.platform == "win32":  # pragma: no cover
-            rebind_std_handle_to_fd(fd)
+        private_fd = _dup_above_std(fd)
     except OSError:
-        # Roll back and serve in place; if the rollback itself fails, fall
-        # through and roll forward instead: fd is stuck on the diversion, but
-        # the private duplicate still holds the wire, so serve from it with
-        # the claim held (the restore at exit gets another chance).
-        if private_fd is None or _restore_fd(fd, private_fd):
-            if private_fd is not None:
-                os.close(private_fd)
-            return stream.buffer, unclaim
-
-    def restore() -> None:
-        # Drain text buffered during the claim (a stray print) to the diversion.
-        with suppress(OSError, ValueError):
-            stream.flush()
-        # A failed restore leaves fd diverted; keep it claimed so later transports are refused.
-        if _restore_fd(fd, private_fd):
-            unclaim()
-
-    # closefd=False: a blocked worker thread may still read this descriptor after exit.
-    return os.fdopen(private_fd, mode, closefd=False), restore
+        return stream.buffer, release
+    claim.private_fd = private_fd
 
-
-def _restore_fd(fd: int, private_fd: int) -> bool:
-    """Point fd back at the protocol pipe; a failure never masks the transport's exit."""
     try:
-        os.dup2(private_fd, fd)
-        if sys.platform == "win32":  # pragma: no cover
-            rebind_std_handle_to_fd(fd)
+        diversion_fd = open_diversion()
     except OSError:
-        return False
-    return True
+        return stream.buffer, release
+    try:
+        os.dup2(diversion_fd, fd)
+    except OSError:
+        # The divert did not land; fd still carries the wire, so serve it in
+        # place through the shared buffer (two writers on one pipe tear frames).
+        with suppress(OSError):
+            os.close(diversion_fd)
+        return stream.buffer, release
+    with suppress(OSError):
+        os.close(diversion_fd)
+    if sys.platform == "win32":  # pragma: no cover
+        with suppress(OSError):
+            rebind_std_handle_to_fd(fd)
+
+    # closefd=False: a worker thread can still block on this descriptor after
+    # the transport exits, so it must never be closed and recycled under it.
+    return os.fdopen(private_fd, mode, closefd=False), release
 
 
 @asynccontextmanager
diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py
@@ -216,30 +216,26 @@ async def test_stdio_server_takes_stdin_off_the_descriptor_table_while_serving(
 async def test_stdio_server_reads_stdin_in_place_when_descriptor_isolation_fails(
     failing_call: str, monkeypatch: pytest.MonkeyPatch
 ) -> None:
-    """A descriptor-table failure while claiming stdin degrades to reading sys.stdin in place.
+    """A descriptor failure while claiming stdin degrades to reading sys.stdin in place.
 
-    SDK-defined behavior: isolation is best-effort; the dup2 variant fails after the private duplicate exists.
+    SDK-defined behavior: isolation is best-effort; when duplicating fd 0 or diverting
+    it fails, the transport serves over the original stdin exactly as v1 did.
     """
     request = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
     with _pipe_planted_on_fd0(monkeypatch) as (in_r, in_w):
         os.write(in_w, _frame(request))
         os.close(in_w)
         monkeypatch.setattr(sys, "stdout", TextIOWrapper(io.BytesIO(), encoding="utf-8"))
 
-        # Injectors fire once, then pass through: pytest capture also calls os.dup/os.dup2,
-        # and a still-armed injector detonating there corrupts every later test's capture.
         if failing_call == "dup":
-            real_dup = os.dup
-            armed = [True]
 
-            def failing_dup(fd: int) -> int:
-                if fd == 0 and armed[0]:
-                    armed[0] = False
-                    raise OSError("injected descriptor failure")
-                return real_dup(fd)
+            def failing_dup_above_std(fd: int) -> int:
+                raise OSError("injected descriptor failure")
 
-            monkeypatch.setattr(os, "dup", failing_dup)
+            monkeypatch.setattr("mcp.server.stdio._dup_above_std", failing_dup_above_std)
         else:
+            # Fires once at the divert, then passes through: pytest's capture
+            # machinery also calls os.dup2 at phase transitions.
             real_dup2 = os.dup2
             armed = [True]
 
@@ -256,12 +252,6 @@ def failing_dup2(fd: int, fd2: int, inheritable: bool = True) -> int:
                 async with read_stream:  # pragma: no branch
                     # Isolation was skipped: fd 0 is still the protocol pipe.
                     assert os.path.sameopenfile(0, in_r)
-                    # In-place transports still own the stream: a second one is refused.
-                    with pytest.raises(RuntimeError, match="already claimed fd 0"):
-                        async with stdio_server():
-                            pytest.fail("unreachable")  # pragma: no cover
-                    # The spent injector passes calls through untouched.
-                    os.close(os.dup(0))
                     received = await read_stream.receive()
                     assert isinstance(received, SessionMessage)
                     assert received.message == request
@@ -278,8 +268,7 @@ async def test_stdio_server_exits_cleanly_when_the_stdin_restore_fails(
     still-diverted fd must refuse later transports rather than serve them the diversion.
     """
     request = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
-    fresh_claims: set[int] = set()
-    monkeypatch.setattr("mcp.server.stdio._claimed", fresh_claims)  # this test leaves fd 0 claimed
+    monkeypatch.setattr("mcp.server.stdio._claims", {})  # this test leaves fd 0 claimed
     with _pipe_planted_on_fd0(monkeypatch) as (_, in_w):
         os.write(in_w, _frame(request))
         os.close(in_w)
@@ -374,6 +363,8 @@ def failing_dup(fd: int) -> int:
         with anyio.fail_after(5):
             async with stdio_server(stdin=anyio.AsyncFile(io.StringIO())) as (read_stream, write_stream):
                 read_stream.close()
+                # The spent injector passes later duplications through untouched.
+                os.close(os.dup(0))
                 devnull_probe = os.open(os.devnull, os.O_WRONLY)
                 try:
                     assert os.path.sameopenfile(1, devnull_probe)
@@ -443,12 +434,14 @@ async def test_a_refused_claim_releases_the_stream_it_already_took(
 
 
 @pytest.mark.anyio
-async def test_stdio_server_serves_in_place_when_the_standard_descriptors_are_incomplete(
+async def test_the_claim_engages_even_when_stderr_is_closed(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """A process missing a standard descriptor is served in place, without surgery.
+    """A process missing fd 2 still gets full isolation.
 
-    With fd 2 closed, a duplicate could land in the standard range: the transport must not touch the table.
+    SDK-defined behavior: the wire duplicate is allocated above the standard range
+    atomically, so a hole in the descriptor table cannot capture it; the stdout
+    diversion falls back to the null device.
     """
     request = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
     response = JSONRPCResponse(jsonrpc="2.0", id=1, result={})
@@ -459,8 +452,12 @@ async def test_stdio_server_serves_in_place_when_the_standard_descriptors_are_in
             with anyio.fail_after(5):
                 async with stdio_server() as (read_stream, write_stream):
                     async with read_stream:  # pragma: no branch
-                        assert os.path.sameopenfile(0, in_r)
-                        assert os.path.sameopenfile(1, out_w)
+                        # Claimed: fd 0 reads the null device, not the pipe.
+                        devnull_probe = os.open(os.devnull, os.O_RDONLY)
+                        try:
+                            assert os.path.sameopenfile(0, devnull_probe)
+                        finally:
+                            os.close(devnull_probe)
 
                         os.write(in_w, _frame(request))
                         received = await read_stream.receive()
@@ -471,69 +468,39 @@ async def test_stdio_server_serves_in_place_when_the_standard_descriptors_are_in
                         assert jsonrpc_message_adapter.validate_json(line.decode().strip()) == response
                         os.close(in_w)
                         await write_stream.aclose()
+
+            assert os.path.sameopenfile(0, in_r)
+            assert os.path.sameopenfile(1, out_w)
         finally:
             os.dup2(saved2, 2)
             os.close(saved2)
 
 
 @pytest.mark.anyio
-async def test_a_claim_that_cannot_roll_back_serves_the_wire_from_the_private_duplicate(
+async def test_stdio_server_serves_in_place_when_the_diversion_cannot_be_opened(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """A mid-claim failure whose rollback also fails rolls forward instead of unclaiming.
-
-    SDK-defined behavior: with fd 0 stuck on the diversion, the transport keeps the
-    claim and serves the protocol from the private duplicate; the restore at exit
-    still runs and puts fd 0 back on the pipe.
-    """
+    """A diversion that cannot be opened leaves fd 0 untouched and serves in place."""
     request = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
     with _pipe_planted_on_fd0(monkeypatch) as (in_r, in_w):
+        os.write(in_w, _frame(request))
+        os.close(in_w)
         monkeypatch.setattr(sys, "stdout", TextIOWrapper(io.BytesIO(), encoding="utf-8"))
 
-        # The first close (the diversion fd) and the second dup2 (the rollback)
-        # fail: the claim can neither complete nor be undone. One-shot and
-        # call-counted injectors, as elsewhere in this file.
-        real_close = os.close
-        close_armed = [True]
+        def failing_diversion() -> int:
+            raise OSError("injected diversion failure")
 
-        def failing_first_close(fd: int) -> None:
-            if close_armed[0]:
-                close_armed[0] = False
-                raise OSError("injected close failure")
-            real_close(fd)
-
-        real_dup2 = os.dup2
-        dup2_calls: list[int] = []
-
-        def flaky_dup2(fd: int, fd2: int, inheritable: bool = True) -> int:
-            dup2_calls.append(fd)
-            if len(dup2_calls) == 2:
-                raise OSError("injected rollback failure")
-            return real_dup2(fd, fd2, inheritable)
-
-        monkeypatch.setattr(os, "close", failing_first_close)
-        monkeypatch.setattr(os, "dup2", flaky_dup2)
+        monkeypatch.setattr("mcp.server.stdio._open_stdin_diversion", failing_diversion)
 
         with anyio.fail_after(5):
-            async with stdio_server() as (read_stream, write_stream):
+            async with stdio_server() as (read_stream, write_stream):  # pragma: no branch
                 async with read_stream:  # pragma: no branch
-                    # fd 0 is stuck on the null device, yet the wire still flows.
-                    devnull_probe = os.open(os.devnull, os.O_RDONLY)
-                    try:
-                        assert os.path.sameopenfile(0, devnull_probe)
-                    finally:
-                        os.close(devnull_probe)
-
-                    os.write(in_w, _frame(request))
+                    assert os.path.sameopenfile(0, in_r)
                     received = await read_stream.receive()
                     assert isinstance(received, SessionMessage)
                     assert received.message == request
-                    os.close(in_w)
                     await write_stream.aclose()
 
-        # The exit restore (third dup2) succeeded: fd 0 is back on the pipe.
-        assert os.path.sameopenfile(0, in_r)
-
 
 class _GatedStdin(io.RawIOBase):
     """Raw stdin double: serves its frames, then blocks until released before EOF.
