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..2b8c4c9 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.""" @@ -748,15 +753,22 @@ def __init__(self): self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.sock.setblocking(1) - def _sock_recv(self, conn) -> str: + def _sock_recv_bytes(self, conn, max_length: int | None = None) -> bytes: buf = b"" while True: b = conn.recv(1) + if not b: + raise ConnectionError("StatusDB connection closed unexpectedly") if b == self.delim: break buf += b + if max_length is not None and len(buf) > max_length: + break - return buf.decode() + return buf + + def _sock_recv(self, conn) -> str: + return self._sock_recv_bytes(conn).decode() def _sock_send(self, conn, msg: str): conn.send(msg.encode() + self.delim) @@ -765,6 +777,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 +797,16 @@ def run_server(self): t.start() def run_connection(self, conn): - with suppress(ConnectionError): + with conn, suppress(ConnectionError): + expected_token = self.token.encode("ascii") + authenticated = secrets.compare_digest( + self._sock_recv_bytes(conn, max_length=len(expected_token)), + expected_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 +873,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..d45ba34 100644 --- a/tests/test_pytest_rerunfailures.py +++ b/tests/test_pytest_rerunfailures.py @@ -8,6 +8,8 @@ from pytest_rerunfailures import ( HAS_PYTEST_HANDLECRASHITEM, + ServerStatusDB, + SocketDB, StatusDB, SubtestReport, XDistHooks, @@ -367,6 +369,97 @@ def mark_test_pending(_): assert db.get_suite_reruns() == 0 +def test_sock_recv_raises_connection_error_on_eof(): + db = SocketDB.__new__(SocketDB) + StatusDB.__init__(db) + connection = mock.MagicMock() + connection.recv.side_effect = [ + b"", + AssertionError("recv called again after EOF"), + ] + + with pytest.raises( + ConnectionError, match="StatusDB connection closed unexpectedly" + ): + db._sock_recv(connection) + + connection.recv.assert_called_once_with(1) + + +@pytest.mark.parametrize( + "authentication", + [ + pytest.param(b"invalid-token", id="incorrect-token"), + pytest.param(b"\xff", id="invalid-utf8"), + pytest.param("é".encode(), id="non-ascii"), + ], +) +def test_statusdb_rejects_unauthenticated_commands(authentication): + 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 = authentication + b"\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_statusdb_accepts_64_byte_authentication_token(): + server = ServerStatusDB.__new__(ServerStatusDB) + StatusDB.__init__(server) + server.rerunfailures_db = {} + server.token = "a" * 64 + + connection = mock.MagicMock() + wire_data = server.token.encode() + b"\nset|test|r|1\n" + connection.recv.side_effect = [bytes((byte,)) for byte in wire_data] + [b""] + + server.run_connection(connection) + + connection.send.assert_called_once_with(b"1\n") + assert server._get("test", "r") == 1 + + +def test_statusdb_rejects_oversized_authentication(): + server = ServerStatusDB.__new__(ServerStatusDB) + StatusDB.__init__(server) + server.rerunfailures_db = {} + server.token = "a" * 64 + server._set("test", "r", 1) + + connection = mock.MagicMock() + oversized_authentication = server.token.encode() + b"x" + wire_data = oversized_authentication + b"\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 connection.recv.call_count == 65 + 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"""