Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/358.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Authenticate xdist StatusDB connections with a per-session token.
23 changes: 19 additions & 4 deletions src/pytest_rerunfailures.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import os
import platform
import re
import secrets
import socket
import sys
import threading
Expand Down Expand Up @@ -595,7 +596,10 @@ def pytest_configure(config):
if is_master(config):
config.failures_db = ServerStatusDB()
else:
config.failures_db = ClientStatusDB(config.workerinput["sock_port"])
config.failures_db = ClientStatusDB(
config.workerinput["sock_port"],
config.workerinput["statusdb_token"],
)
else:
config.failures_db = StatusDB() # no-op db

Expand Down Expand Up @@ -624,8 +628,9 @@ def pytest_runtest_logreport(self, report):
)

def pytest_configure_node(self, node):
"""Configure xdist hook for node sock_port."""
"""Configure xdist hook with StatusDB connection details."""
node.workerinput["sock_port"] = node.config.failures_db.sock_port
node.workerinput["statusdb_token"] = node.config.failures_db.token

def pytest_handlecrashitem(self, crashitem, report, sched):
"""Return the crashitem from pending and collection."""
Expand Down Expand Up @@ -762,6 +767,7 @@ def _sock_send(self, conn, msg: str):
class ServerStatusDB(SocketDB):
def __init__(self) -> None:
super().__init__()
self.token = secrets.token_hex(32)
self.sock.bind(("127.0.0.1", 0))
self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

Expand All @@ -781,7 +787,12 @@ def run_server(self):
t.start()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: the handshake buffers unboundedly before authentication, and connections are unlimited.

_sock_recv appends to buf with no length cap, and every accepted connection gets its own daemon thread with no connection limit — both before any authentication has happened. Any local process can connect and stream bytes without a newline until the controller runs out of memory.

Only a local DoS, but capping the pre-auth read at the known token length (64 hex chars) is nearly free and fits the threat model this PR is addressing.

Comment created by Claude


def run_connection(self, conn):
with suppress(ConnectionError):
with conn, suppress(ConnectionError):
authenticated = secrets.compare_digest(self._sock_recv(conn), self.token)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-ASCII or invalid-UTF-8 handshake input kills this thread instead of rejecting the client.

secrets.compare_digest raises TypeError when either str argument contains non-ASCII characters, and _sock_recv's buf.decode() raises UnicodeDecodeError on invalid UTF-8. Neither is a ConnectionError, so suppress(ConnectionError) does not catch them.

Sending "\u00e9\n" to the port on this branch dumps

Exception in thread Thread-3 (run_connection):
  ...
  authenticated = secrets.compare_digest(self._sock_recv(conn), self.token)
TypeError: comparing strings with non-ASCII characters is not supported

to stderr in the middle of the user's test run, and the client receives b'' rather than the intended "0" rejection.

Since this handshake is exactly the code now expected to face arbitrary local input, it should compare raw bytes (recv the token undecoded and compare_digest(raw, self.token.encode())), or wrap the comparison so any malformed input yields a clean "0".

Comment created by Claude

self._sock_send(conn, "1" if authenticated else "0")
if not authenticated:
return

while True:
op, i, k, v = self._sock_recv(conn).split("|")
if op == "set":
Expand Down Expand Up @@ -848,9 +859,13 @@ def get_suite_reruns(self) -> int:


class ClientStatusDB(SocketDB):
def __init__(self, sock_port):
def __init__(self, sock_port, token):
super().__init__()
self.sock.connect(("127.0.0.1", sock_port))
self._sock_send(self.sock, token)
if self._sock_recv(self.sock) != "1":

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A worker can hang forever here (busy-spinning) if the server never answers the handshake.

_sock_recv loops on conn.recv(1) and only breaks on the b"\n" delimiter. At EOF recv returns b"", which is neither the delimiter nor an error, so the loop never terminates.

That is now directly reachable: whenever the server-side run_connection thread dies before replying (see the compare_digest comment above, or any other non-ConnectionError exception), the new with conn closes the socket, and this call sits in _sock_recv burning 100% of a core — the ConnectionError("StatusDB authentication failed") below is never raised. The xdist controller then waits on a worker that never finishes pytest_configure, hanging the whole session.

The same EOF gap leaves a live spinning thread on the server side: connecting and immediately closing takes threading.active_count() from 2 to 3 permanently.

_sock_recv should treat b"" as end-of-stream and raise ConnectionError; a settimeout on the handshake would be worth considering too.

Comment created by Claude

self.sock.close()
raise ConnectionError("StatusDB authentication failed")

def _set(self, i: str, k: str, v: int):
self._sock_send(self.sock, "|".join(("set", i, k, str(v))))
Expand Down
32 changes: 32 additions & 0 deletions tests/test_pytest_rerunfailures.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

from pytest_rerunfailures import (
HAS_PYTEST_HANDLECRASHITEM,
ServerStatusDB,
StatusDB,
SubtestReport,
XDistHooks,
Expand Down Expand Up @@ -367,6 +368,37 @@ def mark_test_pending(_):
assert db.get_suite_reruns() == 0


def test_statusdb_rejects_unauthenticated_commands():
server = ServerStatusDB.__new__(ServerStatusDB)
StatusDB.__init__(server)
server.rerunfailures_db = {}
server.token = str(mock.sentinel.statusdb_token)
server._set("test", "r", 1)

connection = mock.MagicMock()
wire_data = b"invalid-token\nset|test|r|2\n"
connection.recv.side_effect = [bytes((byte,)) for byte in wire_data]

server.run_connection(connection)

connection.send.assert_called_once_with(b"0\n")
assert server._get("test", "r") == 1


def test_xdist_configure_node_passes_statusdb_connection_details():
failures_db = SimpleNamespace(sock_port=12345, token=mock.sentinel.statusdb_token)
node = SimpleNamespace(
config=SimpleNamespace(failures_db=failures_db), workerinput={}
)

XDistHooks().pytest_configure_node(node)

assert node.workerinput == {
"sock_port": 12345,
"statusdb_token": mock.sentinel.statusdb_token,
}


def test_rerun_passes_after_temporary_test_failure_with_flaky_mark(testdir):
testdir.makepyfile(
f"""
Expand Down
Loading