diff options
Diffstat (limited to 'python/mozperftest/mozperftest/system/proxy.py')
| -rw-r--r-- | python/mozperftest/mozperftest/system/proxy.py | 232 |
1 files changed, 232 insertions, 0 deletions
diff --git a/python/mozperftest/mozperftest/system/proxy.py b/python/mozperftest/mozperftest/system/proxy.py new file mode 100644 index 0000000000..d43a65a0eb --- /dev/null +++ b/python/mozperftest/mozperftest/system/proxy.py @@ -0,0 +1,232 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. +import json +import os +import pathlib +import re +import signal +import tempfile +import threading + +from mozdevice import ADBDevice +from mozlog import get_proxy_logger +from mozprocess import ProcessHandler + +from mozperftest.layers import Layer +from mozperftest.utils import ON_TRY, download_file, get_output_dir, install_package + +LOG = get_proxy_logger(component="proxy") +HERE = os.path.dirname(__file__) + + +class OutputHandler(object): + def __init__(self): + self.proc = None + self.port = None + self.port_event = threading.Event() + + def __call__(self, line): + line = line.strip() + if not line: + return + line = line.decode("utf-8", errors="replace") + try: + data = json.loads(line) + except ValueError: + self.process_output(line) + return + + if isinstance(data, dict) and "action" in data: + # Retrieve the port number for the proxy server from the logs of + # our subprocess. + m = re.match(r"Proxy running on port (\d+)", data.get("message", "")) + if m: + self.port = int(m.group(1)) + self.port_event.set() + LOG.log_raw(data) + else: + self.process_output(json.dumps(data)) + + def finished(self): + self.port_event.set() + + def process_output(self, line): + if self.proc is None: + LOG.process_output(line) + else: + LOG.process_output(self.proc.pid, line) + + def wait_for_port(self): + self.port_event.wait() + return self.port + + +class ProxyRunner(Layer): + """Use a proxy""" + + name = "proxy" + activated = False + + arguments = { + "mode": { + "type": str, + "choices": ["record", "playback"], + "help": "Proxy server mode. Use `playback` to replay from the provided file(s). " + "Use `record` to generate a new recording at the path specified by `--file`. " + "playback - replay from provided file. " + "record - generate a new recording at the specified path.", + }, + "file": { + "type": str, + "nargs": "+", + "help": "The playback files to replay, or the file that a recording will be saved to. " + "For playback, it can be any combination of the following: zip file, manifest file, " + "or a URL to zip/manifest file. " + "For recording, it's a zip fle.", + }, + "perftest-page": { + "type": str, + "default": None, + "help": "This option can be used to specify a single test to record rather than " + "having to continuously modify the pageload_sites.json. This flag should only be " + "used by the perftest team and selects items from " + "`testing/performance/pageload_sites.json` based on the name field. Note that " + "the login fields won't be checked with a request such as this (i.e. it overrides " + "those settings).", + }, + } + + def __init__(self, env, mach_cmd): + super(ProxyRunner, self).__init__(env, mach_cmd) + self.proxy = None + self.tmpdir = None + + def setup(self): + try: + import mozproxy # noqa: F401 + except ImportError: + # Install mozproxy and its vendored deps. + mozbase = pathlib.Path(self.mach_cmd.topsrcdir, "testing", "mozbase") + mozproxy_deps = ["mozinfo", "mozlog", "mozproxy"] + for i in mozproxy_deps: + install_package( + self.mach_cmd.virtualenv_manager, pathlib.Path(mozbase, i) + ) + + # set MOZ_HOST_BIN to find cerutil. Required to set certifcates on android + os.environ["MOZ_HOST_BIN"] = self.mach_cmd.bindir + + def run(self, metadata): + self.metadata = metadata + replay_file = self.get_arg("file") + + # Check if we have a replay file + if replay_file is None: + raise ValueError("Proxy file not provided!!") + + if replay_file is not None and replay_file.startswith("http"): + self.tmpdir = tempfile.TemporaryDirectory() + target = pathlib.Path(self.tmpdir.name, "recording.zip") + self.info("Downloading %s" % replay_file) + download_file(replay_file, target) + replay_file = target + + self.info("Setting up the proxy") + + command = [ + self.mach_cmd.virtualenv_manager.python_path, + "-m", + "mozproxy.driver", + "--topsrcdir=" + self.mach_cmd.topsrcdir, + "--objdir=" + self.mach_cmd.topobjdir, + "--profiledir=" + self.get_arg("profile-directory"), + ] + + if not ON_TRY: + command.extend(["--local"]) + + if metadata.flavor == "mobile-browser": + command.extend(["--tool=%s" % "mitmproxy-android"]) + command.extend(["--binary=android"]) + else: + command.extend(["--tool=%s" % "mitmproxy"]) + # XXX See bug 1712337, we need a single point where we can get the binary used from + # this is required to make it work localy + binary = self.get_arg("browsertime-binary") + if binary is None: + binary = self.mach_cmd.get_binary_path() + command.extend(["--binary=%s" % binary]) + + if self.get_arg("mode") == "record": + output = self.get_arg("output") + if output is None: + output = pathlib.Path(self.mach_cmd.topsrcdir, "artifacts") + results_dir = get_output_dir(output) + + command.extend(["--mode", "record"]) + command.append(str(pathlib.Path(results_dir, replay_file))) + elif self.get_arg("mode") == "playback": + command.extend(["--mode", "playback"]) + command.append(str(replay_file)) + else: + raise ValueError("Proxy mode not provided please provide proxy mode") + + inject_deterministic = self.get_arg("deterministic") + if inject_deterministic: + command.extend(["--deterministic"]) + + print(" ".join(command)) + self.output_handler = OutputHandler() + self.proxy = ProcessHandler( + command, + processOutputLine=self.output_handler, + onFinish=self.output_handler.finished, + ) + self.output_handler.proc = self.proxy + self.proxy.run() + + # Wait until we've retrieved the proxy server's port number so we can + # configure the browser properly. + port = self.output_handler.wait_for_port() + if port is None: + raise ValueError("Unable to retrieve the port number from mozproxy") + self.info("Received port number %s from mozproxy" % port) + + prefs = { + "network.proxy.type": 1, + "network.proxy.http": "127.0.0.1", + "network.proxy.http_port": port, + "network.proxy.ssl": "127.0.0.1", + "network.proxy.ssl_port": port, + "network.proxy.no_proxies_on": "127.0.0.1", + } + browser_prefs = metadata.get_options("browser_prefs") + browser_prefs.update(prefs) + + if metadata.flavor == "mobile-browser": + self.info("Setting reverse port fw for android device") + device = ADBDevice() + device.create_socket_connection("reverse", "tcp:%s" % port, "tcp:%s" % port) + + return metadata + + def teardown(self): + err = None + if self.proxy is not None: + returncode = self.proxy.wait(0) + if returncode is not None: + err = ValueError( + "mozproxy terminated early with return code %d" % returncode + ) + else: + kill_signal = getattr(signal, "CTRL_BREAK_EVENT", signal.SIGINT) + os.kill(self.proxy.pid, kill_signal) + self.proxy.wait() + self.proxy = None + if self.tmpdir is not None: + self.tmpdir.cleanup() + self.tmpdir = None + + if err: + raise err |