feat: add support for marionette

src/ffpuppet/bootstrapper.py

@@ -8,14 +8,16 @@
 from select import select
 from socket import SO_REUSEADDR, SOL_SOCKET, socket
 from time import perf_counter, sleep
-from typing import Callable
-
-from .exceptions import BrowserTerminatedError, BrowserTimeoutError, LaunchError
+from typing import TYPE_CHECKING, Callable
 
 # as of python 3.10 socket.timeout was made an alias of TimeoutError
 # pylint: disable=ungrouped-imports,wrong-import-order
 from socket import timeout as socket_timeout  # isort: skip
 
+from .exceptions import BrowserTerminatedError, BrowserTimeoutError, LaunchError
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
 
 LOG = getLogger(__name__)
 
@@ -63,6 +65,18 @@ def __enter__(self) -> Bootstrapper:
     def __exit__(self, *exc: object) -> None:
         self.close()
 
+    @classmethod
+    def check_port(cls, value: int) -> bool:
+        """Verify port value is in valid range.
+
+        Args:
+            None
+
+        Returns:
+            bool
+        """
+        return value == 0 or 1024 <= value <= 65535
+
     def close(self) -> None:
         """Close listening socket.
 
@@ -85,9 +99,35 @@ def create(cls, attempts: int = 50, port: int = 0) -> Bootstrapper:
         Returns:
             Bootstrapper.
         """
+        sock = cls.create_socket(attempts=attempts, port=port)
+        if sock is None:
+            raise LaunchError("Could not find available port")
+        return cls(sock)
+
+    @classmethod
+    def create_socket(
+        cls,
+        attempts: int = 50,
+        blocked: Iterable[int] | None = BLOCKED_PORTS,
+        port: int = 0,
+    ) -> socket | None:
+        """Create a listening socket.
+
+        Args:
+            attempts: Number of times to attempt to bind.
+            blocked: Ports that cannot be used.
+            port: Port to use. Use 0 for system select.
+
+        Returns:
+            A listening socket.
+        """
         assert attempts > 0
-        assert port == 0 or 1024 <= port <= 65535
-        assert port not in cls.BLOCKED_PORTS
+        if not cls.check_port(port):
+            LOG.debug("requested invalid port: %d", port)
+            return None
+        if blocked and port in blocked:
+            LOG.debug("requested blocked port: %d", port)
+            return None
         for _ in range(attempts):
             sock = socket()
             sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
@@ -100,14 +140,14 @@ def create(cls, attempts: int = 50, port: int = 0) -> Bootstrapper:
                 sleep(0.1)
                 continue
             # avoid blocked ports
-            if sock.getsockname()[1] in cls.BLOCKED_PORTS:
+            if blocked and sock.getsockname()[1] in blocked:
                 LOG.debug("bound to blocked port, retrying...")
                 sock.close()
                 continue
             break
         else:
-            raise LaunchError("Could not find available port")
-        return cls(sock)
+            return None
+        return sock
 
     @property
     def location(self) -> str:

src/ffpuppet/core.py

@@ -109,8 +109,9 @@ class FFPuppet:
         "_logs",
         "_proc_tree",
         "_profile_template",
-        "_xvfb",
         "_working_path",
+        "_xvfb",
+        "marionette",
         "profile",
         "reason",
     )
@@ -136,6 +137,7 @@ def __init__(
         self._profile_template = use_profile
         self._xvfb: Xvfb | None = None
         self._working_path = working_path
+        self.marionette: int | None = None
         self.profile: Profile | None = None
         self.reason: Reason | None = Reason.CLOSED
 
@@ -583,6 +585,7 @@ def close(self, force_close: bool = False) -> None:
 
         # reset remaining to closed state
         try:
+            self.marionette = None
             self._proc_tree = None
             self._logs.close()
             self._checks = []
@@ -677,6 +680,7 @@ def launch(
         launch_timeout: int = 300,
         location: str | None = None,
         log_limit: int = 0,
+        marionette: int | None = None,
         memory_limit: int = 0,
         prefs_js: Path | None = None,
         extension: list[Path] | None = None,
@@ -764,14 +768,30 @@ def launch(
                 "network.proxy.failover_direct": "false",
                 "privacy.partition.network_state": "false",
             }
-            self.profile.add_prefs(prefs)
 
             launch_args = [bootstrapper.location]
             is_windows = system() == "Windows"
             if is_windows:
                 # disable launcher process
                 launch_args.append("-no-deelevate")
                 launch_args.append("-wait-for-browser")
+
+            if marionette is not None:
+                # find/validate port to use
+                free_sock = Bootstrapper.create_socket(port=marionette)
+                if free_sock is None:
+                    if marionette == 0:
+                        LOG.error("Cannot find available port for marionette")
+                    else:
+                        LOG.error("Marionette cannot use port: %d", marionette)
+                    raise LaunchError("Debugging server port unavailable")
+                self.marionette = free_sock.getsockname()[1]
+                free_sock.close()
+                launch_args.append("-marionette")
+                prefs["marionette.port"] = str(self.marionette)
+
+            self.profile.add_prefs(prefs)
+
             cmd = self.build_launch_cmd(str(bin_path), additional_args=launch_args)
 
             # open logs
@@ -830,6 +850,7 @@ def launch(
             if self._proc_tree is None:
                 # only clean up here if a launch was not attempted or Popen failed
                 LOG.debug("process not launched")
+                self.marionette = None
                 self.profile.remove()
                 self.profile = None
                 self.reason = Reason.CLOSED

src/ffpuppet/main.py

@@ -13,6 +13,7 @@
 from tempfile import mkdtemp
 from time import sleep, strftime
 
+from .bootstrapper import Bootstrapper
 from .core import Debugger, FFPuppet, Reason
 from .exceptions import BrowserExecutionError
 from .helpers import certutil_available, certutil_find
@@ -134,6 +135,15 @@ def parse_args(argv: list[str] | None = None) -> Namespace:
         nargs="?",
         help="Headless mode. 'default' uses browser's built-in headless mode.",
     )
+    cfg_group.add_argument(
+        "--marionette",
+        const=0,
+        default=None,
+        nargs="?",
+        type=int,
+        help="Enable marionette. If a port is provided it is used otherwise "
+        "a random port is selected. (default: disabled)",
+    )
     cfg_group.add_argument(
         "-p",
         "--prefs",
@@ -267,6 +277,8 @@ def parse_args(argv: list[str] | None = None) -> Namespace:
         value = int(Path(settings).read_bytes())
         if value > 1:
             parser.error(f"rr needs {settings} <= 1, but it is {value}")
+    if args.marionette is not None and not Bootstrapper.check_port(args.marionette):
+        parser.error("--marionette must be 0 or > 1024 and < 65536")
     if not args.logs.is_dir():
         parser.error(f"Log output directory is invalid '{args.logs}'")
     args.log_level = log_level_map[args.log_level]
@@ -317,6 +329,7 @@ def main(argv: list[str] | None = None) -> None:
             location=args.url,
             launch_timeout=args.launch_timeout,
             log_limit=args.log_limit,
+            marionette=args.marionette,
             memory_limit=args.memory,
             prefs_js=args.prefs,
             extension=args.extension,
@@ -326,6 +339,8 @@ def main(argv: list[str] | None = None) -> None:
             assert ffp.profile is not None
             assert ffp.profile.path is not None
             Profile.check_prefs(ffp.profile.path / "prefs.js", args.prefs)
+        if ffp.marionette is not None:
+            LOG.info("Marionette listening on port: %d", ffp.marionette)
         LOG.info("Running Firefox (pid: %d)...", ffp.get_pid())
         while ffp.is_healthy():
             sleep(args.poll_interval)

src/ffpuppet/resources/testff.py

@@ -29,6 +29,7 @@ def main() -> int:
     parser = ArgumentParser(prog="testff", description="Fake Firefox for testing")
     parser.add_argument("url")
     parser.add_argument("-headless", action="store_true", help="ignored")
+    parser.add_argument("-marionette", nargs="?", type=int, help="ignored")
     parser.add_argument("-new-instance", action="store_true", help="ignored")
     parser.add_argument("-no-deelevate", action="store_true", help="ignored")
     parser.add_argument("-wait-for-browser", action="store_true", help="ignored")

src/ffpuppet/test_bootstrapper.py

@@ -184,23 +184,22 @@ def _fake_browser(port, payload_size=5120):
 
 
 @mark.parametrize(
-    "bind, attempts, raised",
+    "bind, attempts",
     [
         # failed to bind (OSError)
-        ((OSError(0, "foo1"),), 1, LaunchError),
+        ((OSError(0, "foo1"),), 1),
         # failed to bind (PermissionError) - multiple attempts
-        (repeat(PermissionError(10013, "foo2"), 4), 4, LaunchError),
+        (repeat(PermissionError(10013, "foo2"), 4), 4),
     ],
 )
-def test_bootstrapper_08(mocker, bind, attempts, raised):
-    """test Bootstrapper.create() - failures"""
+def test_bootstrapper_08(mocker, bind, attempts):
+    """test Bootstrapper.create_socket() - failures"""
     mocker.patch("ffpuppet.bootstrapper.sleep", autospec=True)
     fake_sock = mocker.MagicMock(spec_set=socket)
     fake_sock.bind.side_effect = bind
     mocker.patch("ffpuppet.bootstrapper.select", return_value=([fake_sock], None, None))
     mocker.patch("ffpuppet.bootstrapper.socket", return_value=fake_sock)
-    with raises(raised), Bootstrapper.create(attempts=attempts):
-        pass
+    assert Bootstrapper.create_socket(attempts=attempts) is None
     assert fake_sock.bind.call_count == attempts
     assert fake_sock.close.call_count == attempts
 
@@ -216,3 +215,33 @@ def test_bootstrapper_09(mocker):
     with Bootstrapper.create(attempts=2):
         pass
     assert fake_sock.close.call_count == 2
+
+
+def test_bootstrapper_10(mocker):
+    """test Bootstrapper.create() - failure"""
+    mocker.patch("ffpuppet.bootstrapper.Bootstrapper.create_socket", return_value=None)
+    with raises(LaunchError), Bootstrapper.create():
+        pass
+
+
+@mark.parametrize("value", [123, 5555])
+def test_bootstrapper_11(value):
+    """test Bootstrapper.create_socket() - unusable ports"""
+    assert Bootstrapper.create_socket(blocked=[5555], port=value) is None
+
+
+@mark.parametrize(
+    "value, result",
+    [
+        (0, True),
+        (1337, True),
+        (32768, True),
+        (-1, False),
+        (1, False),
+        (1023, False),
+        (65536, False),
+    ],
+)
+def test_bootstrapper_12(value, result):
+    """test Bootstrapper.check_port()"""
+    assert Bootstrapper.check_port(value) == result

src/ffpuppet/test_ffpuppet.py

@@ -963,3 +963,25 @@ def test_ffpuppet_32(mocker):
             with raises(NotImplementedError):
                 ffp.dump_coverage()
         ffp._proc_tree = None
+
+
+def test_ffpuppet_33():
+    """test FFPuppet.launch() with marionette"""
+    with FFPuppet() as ffp, HTTPTestServer() as srv:
+        assert ffp.marionette is None
+        ffp.launch(TESTFF_BIN, marionette=0, location=srv.get_addr())
+        assert ffp.marionette is not None
+        ffp.close()
+        assert ffp.marionette is None
+
+
+@mark.parametrize("port", [0, 123])
+def test_ffpuppet_34(mocker, port):
+    """test FFPuppet.launch() with marionette failure"""
+    fake_bts = mocker.patch("ffpuppet.core.Bootstrapper", autospec=True)
+    fake_bts.create.return_value.location = ""
+    fake_bts.create_socket.return_value = None
+    with FFPuppet() as ffp, HTTPTestServer() as srv:
+        with raises(LaunchError):
+            ffp.launch(TESTFF_BIN, marionette=port, location=srv.get_addr())
+        assert ffp.marionette is None

src/ffpuppet/test_main.py

@@ -72,6 +72,12 @@ def test_parse_args_01(capsys, mocker, tmp_path):
     with raises(SystemExit):
         parse_args([str(fake_bin), "--log-limit", "-1"])
     assert "error: --log-limit must be >= 0" in capsys.readouterr()[-1]
+    # invalid marionette port
+    with raises(SystemExit):
+        parse_args([str(fake_bin), "--marionette", "123"])
+    assert (
+        "error: --marionette must be 0 or > 1024 and < 65536" in capsys.readouterr()[-1]
+    )
     # invalid memory limit
     with raises(SystemExit):
         parse_args([str(fake_bin), "--memory", "-1"])