Authenticate xdist StatusDB connections

changes/358.bugfix.rst

@@ -0,0 +1 @@
+Authenticate xdist StatusDB connections with a per-session token.

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))))

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"""