python/mozperftest/mozperftest/system/proxy.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4 import json
   5 import os
   6 import pathlib
   7 import re
   8 import signal
   9 import tempfile
  10 import threading
  11
  12 from mozdevice import ADBDevice
  13 from mozlog import get_proxy_logger
  14 from mozprocess import ProcessHandler
  15
  16 from mozperftest.layers import Layer
  17 from mozperftest.utils import (
  18     ON_TRY,
  19     download_file,
  20     get_output_dir,
  21     get_pretty_app_name,
  22     install_package,
  23 )
  24
  25 LOG = get_proxy_logger(component="proxy")
  26 HERE = os.path.dirname(__file__)
  27
  28
  29 class OutputHandler(object):
  30     def __init__(self):
  31         self.proc = None
  32         self.port = None
  33         self.port_event = threading.Event()
  34
  35     def __call__(self, line):
  36         line = line.strip()
  37         if not line:
  38             return
  39         line = line.decode("utf-8", errors="replace")
  40         try:
  41             data = json.loads(line)
  42         except ValueError:
  43             self.process_output(line)
  44             return
  45
  46         if isinstance(data, dict) and "action" in data:
  47             # Retrieve the port number for the proxy server from the logs of
  48             # our subprocess.
  49             m = re.match(r"Proxy running on port (\d+)", data.get("message", ""))
  50             if m:
  51                 self.port = int(m.group(1))
  52                 self.port_event.set()
  53             LOG.log_raw(data)
  54         else:
  55             self.process_output(json.dumps(data))
  56
  57     def finished(self):
  58         self.port_event.set()
  59
  60     def process_output(self, line):
  61         if self.proc is None:
  62             LOG.process_output(line)
  63         else:
  64             LOG.process_output(self.proc.pid, line)
  65
  66     def wait_for_port(self):
  67         self.port_event.wait()
  68         return self.port
  69
  70
  71 class ProxyRunner(Layer):
  72     """Use a proxy"""
  73
  74     name = "proxy"
  75     activated = False
  76
  77     arguments = {
  78         "mode": {
  79             "type": str,
  80             "choices": ["record", "playback"],
  81             "help": "Proxy server mode. Use `playback` to replay from the provided file(s). "
  82             "Use `record` to generate a new recording at the path specified by `--file`. "
  83             "playback - replay from provided file. "
  84             "record - generate a new recording at the specified path.",
  85         },
  86         "file": {
  87             "type": str,
  88             "nargs": "+",
  89             "help": "The playback files to replay, or the file that a recording will be saved to. "
  90             "For playback, it can be any combination of the following: zip file, manifest file, "
  91             "or a URL to zip/manifest file. "
  92             "For recording, it's a zip fle.",
  93         },
  94         "perftest-page": {
  95             "type": str,
  96             "default": None,
  97             "help": "This option can be used to specify a single test to record rather than "
  98             "having to continuously modify the pageload_sites.json. This flag should only be "
  99             "used by the perftest team and selects items from "
 100             "`testing/performance/pageload_sites.json` based on the name field. Note that "
 101             "the login fields won't be checked with a request such as this (i.e. it overrides "
 102             "those settings).",
 103         },
 104         "deterministic": {
 105             "action": "store_true",
 106             "default": False,
 107             "help": "If set, the deterministic JS script will be injected into the pages.",
 108         },
 109     }
 110
 111     def __init__(self, env, mach_cmd):
 112         super(ProxyRunner, self).__init__(env, mach_cmd)
 113         self.proxy = None
 114         self.tmpdir = None
 115
 116     def setup(self):
 117         try:
 118             import mozproxy  # noqa: F401
 119         except ImportError:
 120             # Install mozproxy and its vendored deps.
 121             mozbase = pathlib.Path(self.mach_cmd.topsrcdir, "testing", "mozbase")
 122             mozproxy_deps = ["mozinfo", "mozlog", "mozproxy"]
 123             for i in mozproxy_deps:
 124                 install_package(
 125                     self.mach_cmd.virtualenv_manager, pathlib.Path(mozbase, i)
 126                 )
 127
 128         # set MOZ_HOST_BIN to find cerutil. Required to set certifcates on android
 129         os.environ["MOZ_HOST_BIN"] = self.mach_cmd.bindir
 130
 131     def run(self, metadata):
 132         self.metadata = metadata
 133         replay_file = self.get_arg("file")
 134
 135         # Check if we have a replay file
 136         if replay_file is None:
 137             raise ValueError("Proxy file not provided!!")
 138
 139         if replay_file is not None and replay_file.startswith("http"):
 140             self.tmpdir = tempfile.TemporaryDirectory()
 141             target = pathlib.Path(self.tmpdir.name, "recording.zip")
 142             self.info("Downloading %s" % replay_file)
 143             download_file(replay_file, target)
 144             replay_file = target
 145
 146         self.info("Setting up the proxy")
 147
 148         command = [
 149             self.mach_cmd.virtualenv_manager.python_path,
 150             "-m",
 151             "mozproxy.driver",
 152             "--topsrcdir=" + self.mach_cmd.topsrcdir,
 153             "--objdir=" + self.mach_cmd.topobjdir,
 154             "--profiledir=" + self.get_arg("profile-directory"),
 155         ]
 156
 157         if not ON_TRY:
 158             command.extend(["--local"])
 159
 160         if metadata.flavor == "mobile-browser":
 161             command.extend(["--tool=%s" % "mitmproxy-android"])
 162             command.extend(["--binary=android"])
 163             command.extend(
 164                 [f"--app={get_pretty_app_name(self.get_arg('android-app-name'))}"]
 165             )
 166         else:
 167             command.extend(["--tool=%s" % "mitmproxy"])
 168             # XXX See bug 1712337, we need a single point where we can get the binary used from
 169             # this is required to make it work localy
 170             binary = self.get_arg("browsertime-binary")
 171             if binary is None:
 172                 binary = self.mach_cmd.get_binary_path()
 173             command.extend(["--binary=%s" % binary])
 174
 175         if self.get_arg("mode") == "record":
 176             output = self.get_arg("output")
 177             if output is None:
 178                 output = pathlib.Path(self.mach_cmd.topsrcdir, "artifacts")
 179             results_dir = get_output_dir(output)
 180
 181             command.extend(["--mode", "record"])
 182             command.append(str(pathlib.Path(results_dir, replay_file)))
 183         elif self.get_arg("mode") == "playback":
 184             command.extend(["--mode", "playback"])
 185             command.append(str(replay_file))
 186         else:
 187             raise ValueError("Proxy mode not provided please provide proxy mode")
 188
 189         inject_deterministic = self.get_arg("deterministic")
 190         if inject_deterministic:
 191             command.extend(["--deterministic"])
 192
 193         print(" ".join(command))
 194         self.output_handler = OutputHandler()
 195         self.proxy = ProcessHandler(
 196             command,
 197             processOutputLine=self.output_handler,
 198             onFinish=self.output_handler.finished,
 199         )
 200         self.output_handler.proc = self.proxy
 201         self.proxy.run()
 202
 203         # Wait until we've retrieved the proxy server's port number so we can
 204         # configure the browser properly.
 205         port = self.output_handler.wait_for_port()
 206         if port is None:
 207             raise ValueError("Unable to retrieve the port number from mozproxy")
 208         self.info("Received port number %s from mozproxy" % port)
 209
 210         prefs = {
 211             "network.proxy.type": 1,
 212             "network.proxy.http": "127.0.0.1",
 213             "network.proxy.http_port": port,
 214             "network.proxy.ssl": "127.0.0.1",
 215             "network.proxy.ssl_port": port,
 216             "network.proxy.no_proxies_on": "127.0.0.1",
 217         }
 218         browser_prefs = metadata.get_options("browser_prefs")
 219         browser_prefs.update(prefs)
 220
 221         if metadata.flavor == "mobile-browser":
 222             self.info("Setting reverse port fw for android device")
 223             device = ADBDevice()
 224             device.create_socket_connection("reverse", "tcp:%s" % port, "tcp:%s" % port)
 225
 226         return metadata
 227
 228     def teardown(self):
 229         err = None
 230         if self.proxy is not None:
 231             returncode = self.proxy.wait(0)
 232             if returncode is not None:
 233                 err = ValueError(
 234                     "mozproxy terminated early with return code %d" % returncode
 235                 )
 236             else:
 237                 kill_signal = getattr(signal, "CTRL_BREAK_EVENT", signal.SIGINT)
 238                 os.kill(self.proxy.pid, kill_signal)
 239                 self.proxy.wait()
 240             self.proxy = None
 241         if self.tmpdir is not None:
 242             self.tmpdir.cleanup()
 243             self.tmpdir = None
 244
 245         if err:
 246             raise err