diff --git a/changes/358.bugfix.rst b/changes/358.bugfix.rst new file mode 100644 index 0000000..1d92745 --- /dev/null +++ b/changes/358.bugfix.rst @@ -0,0 +1 @@ +Authenticate xdist StatusDB connections with a per-session token. diff --git a/src/pytest_rerunfailures.py b/src/pytest_rerunfailures.py index 22aef22..fba2e83 100644 --- a/src/pytest_rerunfailures.py +++ b/src/pytest_rerunfailures.py @@ -3,6 +3,7 @@ import os import platform import re +import secrets import socket import sys import threading @@ -598,7 +599,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 @@ -627,8 +631,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.""" @@ -765,6 +770,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) @@ -784,7 +790,12 @@ def run_server(self): t.start() def run_connection(self, conn): - with suppress(ConnectionError): + with conn, suppress(ConnectionError): + authenticated = secrets.compare_digest(self._sock_recv(conn), self.token) + 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": @@ -851,9 +862,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": + 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)))) diff --git a/tests/test_pytest_rerunfailures.py b/tests/test_pytest_rerunfailures.py index c257bbf..8141e4b 100644 --- a/tests/test_pytest_rerunfailures.py +++ b/tests/test_pytest_rerunfailures.py @@ -8,6 +8,7 @@ from pytest_rerunfailures import ( HAS_PYTEST_HANDLECRASHITEM, + ServerStatusDB, StatusDB, SubtestReport, XDistHooks, @@ -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"""