diff --git a/.gitignore b/.gitignore index 9dd3935..f214a38 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ launch.exe # Internal testing /extra-docs pythonlib/test* +!pythonlib/tests/ jsonvv/test* /.vscode /bundle/fonts/extra diff --git a/pythonlib/camoufox/virtdisplay.py b/pythonlib/camoufox/virtdisplay.py index 4321835..6a45b83 100644 --- a/pythonlib/camoufox/virtdisplay.py +++ b/pythonlib/camoufox/virtdisplay.py @@ -1,10 +1,9 @@ import os +import select import subprocess # nosec -from glob import glob -from multiprocessing import Lock -from random import randrange +import time from shutil import which -from typing import List, Optional +from typing import Optional from camoufox.exceptions import ( CannotExecuteXvfb, @@ -13,20 +12,12 @@ from camoufox.exceptions import ( ) from camoufox.pkgman import OS_NAME +# Safe timeout for Xvfb writing display num, prevents infinite hang. +DISPLAYFD_READ_TIMEOUT_S = 10.0 + class VirtualDisplay: - """ - A minimal virtual display implementation for Linux. - """ - - def __init__(self, debug: Optional[bool] = False) -> None: - """ - Constructor for the VirtualDisplay class (singleton object). - """ - self.debug = debug - self.proc: Optional[subprocess.Popen] = None - self._display: Optional[int] = None - self._lock = Lock() + """A minimal virtual display implementation for Linux.""" xvfb_args = ( # fmt: off @@ -39,18 +30,19 @@ class VirtualDisplay: "-extension", "XVideo", "-extension", "XVideo-MotionCompensation", "-extension", "XINERAMA", - "-shmem", "-fp", "built-ins", "-nocursor", "-br", # fmt: on ) + def __init__(self, debug: bool = False) -> None: + self.debug = debug + self.proc: Optional[subprocess.Popen] = None + self._display: Optional[int] = None + @property def xvfb_path(self) -> str: - """ - Get the path to the xvfb executable - """ path = which("Xvfb") if not path: raise CannotFindXvfb("Please install Xvfb to use headless mode.") @@ -58,89 +50,74 @@ class VirtualDisplay: raise CannotExecuteXvfb(f"I do not have permission to execute Xvfb: {path}") return path - @property - def xvfb_cmd(self) -> List[str]: - """ - Get the xvfb command - """ - return [self.xvfb_path, f':{self.display}', *self.xvfb_args] - - def execute_xvfb(self): - """ - Spawn a detatched process - """ - if self.debug: - print('Starting virtual display:', ' '.join(self.xvfb_cmd)) - self.proc = subprocess.Popen( # nosec - self.xvfb_cmd, - stdout=None if self.debug else subprocess.DEVNULL, - stderr=None if self.debug else subprocess.DEVNULL, - ) - def get(self) -> str: - """ - Get the display number - """ - self.assert_linux() + self._assert_linux() - with self._lock: - if self.proc is None: - self.execute_xvfb() - elif self.debug: - print(f'Using virtual display: {self.display}') - return f':{self.display}' + if self.proc is None: + # Launch Xvfb with -displayfd so Xvfb itself picks a free display + # number atomically and reports it back. Avoids userspace races. + # subprocess.Popen's pass_fds keeps an fd at its parent number in + # the child (unlike Node's `stdio: [..., 'pipe']` which renumbers + # to 3), so we tell Xvfb that exact number. + read_fd, write_fd = os.pipe() + cmd = [self.xvfb_path, "-displayfd", str(write_fd), *self.xvfb_args] + if self.debug: + print("Starting virtual display:", " ".join(cmd)) + self.proc = subprocess.Popen( # nosec + cmd, + stdin=subprocess.DEVNULL, + stdout=None if self.debug else subprocess.DEVNULL, + stderr=None if self.debug else subprocess.DEVNULL, + start_new_session=True, + pass_fds=(write_fd,), + env={ + **os.environ, + # Force Mesa software GLX; we don't use the GPU anyway. + "__GLX_VENDOR_LIBRARY_NAME": "mesa", + "LIBGL_ALWAYS_SOFTWARE": "1", + }, + ) + os.close(write_fd) # so the read end EOFs when Xvfb closes its end - def kill(self): - """ - Terminate the xvfb process - """ - with self._lock: - if self.proc and self.proc.poll() is None: - if self.debug: - print('Terminating virtual display:', self.display) - self.proc.terminate() + buf = b"" + deadline = time.monotonic() + DISPLAYFD_READ_TIMEOUT_S + try: + while b"\n" not in buf: + remaining = deadline - time.monotonic() + if remaining <= 0 or not select.select([read_fd], [], [], remaining)[0]: + self.kill() + raise CannotExecuteXvfb( + f"Xvfb did not report a display within " + f"{int(DISPLAYFD_READ_TIMEOUT_S * 1000)}ms" + ) + chunk = os.read(read_fd, 64) + if not chunk: + self.kill() + raise CannotExecuteXvfb( + f"Xvfb did not report a display " + f"(got {buf!r}, exit={self.proc.poll()})" + ) + buf += chunk + finally: + os.close(read_fd) - def __del__(self): - """ - Kill and delete the VirtualDisplay object - """ - self.kill() + try: + self._display = int(buf.strip()) + except ValueError: + self.kill() + raise CannotExecuteXvfb(f"Xvfb wrote non-integer display: {buf!r}") + elif self.debug: + print(f"Using virtual display: {self._display}") + + return f":{self._display}" + + def kill(self) -> None: + if self.proc and self.proc.poll() is None: + if self.debug: + print("Terminating virtual display:", self._display) + self.proc.terminate() @staticmethod - def _get_lock_files() -> List[str]: - """ - Get list of lock files in /tmp - """ - tmpd = os.environ.get('TMPDIR', '/tmp') # nosec - try: - lock_files = glob(os.path.join(tmpd, ".X*-lock")) - except FileNotFoundError: - return [] - return [p for p in lock_files if os.path.isfile(p)] - - @staticmethod - def _free_display() -> int: - """ - Search for free display - """ - ls = list( - map(lambda x: int(x.split("X")[1].split("-")[0]), VirtualDisplay._get_lock_files()) - ) - return max(99, max(ls) + randrange(3, 20)) if ls else 99 # nosec - - @property - def display(self) -> int: - """ - Get the display number - """ - if self._display is None: - self._display = self._free_display() - return self._display - - @staticmethod - def assert_linux(): - """ - Assert that the current OS is Linux - """ - if OS_NAME != 'lin': + def _assert_linux() -> None: + if OS_NAME != "lin": raise VirtualDisplayNotSupported("Virtual display is only supported on Linux.") diff --git a/pythonlib/tests/test_virtdisplay.py b/pythonlib/tests/test_virtdisplay.py new file mode 100644 index 0000000..7cc454f --- /dev/null +++ b/pythonlib/tests/test_virtdisplay.py @@ -0,0 +1,131 @@ +""" +Tests for camoufox.virtdisplay. + +Mirrors camoufox-js/test/virtdisplay.test.ts. + +Run with: + cd pythonlib && python -m pytest tests/test_virtdisplay.py -v + +VIRTDISPLAY_TEST_N controls the concurrent-launch count. Default is kept +low so the test passes on any developer box; set to e.g. 1000 to exercise +real scaling. At high N you will need `ulimit -n` headroom -- each Xvfb +takes one X11 socket plus our -displayfd pipe. +""" + +import os +import re +import sys +import time +from typing import List, Set + +import pytest + +# Make `import camoufox` resolve to the in-tree pythonlib without an install. +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from camoufox.virtdisplay import VirtualDisplay # noqa: E402 + +DISPLAY_RE = re.compile(r"^:\d+$") +N = int(os.environ.get("VIRTDISPLAY_TEST_N", "50")) + +pytestmark = pytest.mark.skipif( + sys.platform != "linux", reason="VirtualDisplay is Linux-only" +) + + +# Track every VirtualDisplay we spawn so the fixture can guarantee +# cleanup even if an assertion fails mid-test. +@pytest.fixture +def tracked() -> List[VirtualDisplay]: + items: List[VirtualDisplay] = [] + yield items + for vd in items: + try: + vd.kill() + except Exception: + pass + + +def _track(items: List[VirtualDisplay], vd: VirtualDisplay) -> VirtualDisplay: + items.append(vd) + return vd + + +def _wait_for_exit(vd: VirtualDisplay, timeout_s: float = 5.0) -> None: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if vd.proc is None or vd.proc.poll() is not None: + return + time.sleep(0.025) + + +def test_single_launch_returns_valid_display_and_kill_terminates_xvfb(tracked): + vd = _track(tracked, VirtualDisplay()) + display = vd.get() + assert DISPLAY_RE.match(display), display + assert vd.proc is not None + assert vd.proc.poll() is None # alive + + vd.kill() + _wait_for_exit(vd) + assert vd.proc.poll() is not None # exited + + +def test_get_is_idempotent_within_one_virtual_display(tracked): + vd = _track(tracked, VirtualDisplay()) + a = vd.get() + b = vd.get() + assert a == b + + +def test_concurrent_reservations_all_get_unique_displays(tracked): + # Every VirtualDisplay spawns its own Xvfb. Each Xvfb scans up from :0 + # and atomically claims the first free X11 socket (kernel-mediated + # bind, no userspace race). -displayfd reports the chosen number back + # to us. A duplicate here would mean we mis-parsed or mis-routed the + # displayfd output, or two Xvfbs somehow bound the same socket. + vds = [_track(tracked, VirtualDisplay()) for _ in range(N)] + + displays = [vd.get() for vd in vds] + + for d in displays: + assert DISPLAY_RE.match(d), d + + unique: Set[str] = set(displays) + assert len(unique) == len(displays), ( + f"duplicate displays: {len(displays) - len(unique)} of {len(displays)}" + ) + + # Every Xvfb is alive. + for vd in vds: + assert vd.proc is not None and vd.proc.poll() is None + + # Tear them all down and confirm every Xvfb actually exited -- no + # leaked processes. + for vd in vds: + vd.kill() + for vd in vds: + _wait_for_exit(vd) + for vd in vds: + assert vd.proc.poll() is not None + + +def test_released_display_numbers_can_be_reused_on_the_next_launch(tracked): + a = _track(tracked, VirtualDisplay()) + a_display = a.get() + + a.kill() + _wait_for_exit(a) + + # Spawning a new Xvfb after release must succeed. The new display + # number may or may not equal a_display -- Xvfb's allocation order is + # its concern -- but we must get *some* display. + b = _track(tracked, VirtualDisplay()) + b_display = b.get() + assert DISPLAY_RE.match(b_display), b_display + + b.kill() + _wait_for_exit(b) + + # Sanity: a_display was a valid form too. + assert DISPLAY_RE.match(a_display), a_display