Use tornado for a local web server if possible

pmeenan · pmeenan · commit 68733aea282a · 2017-11-15T17:28:31.000-05:00
diff --git a/Dockerfile b/Dockerfile
@@ -55,6 +55,7 @@ RUN pip install \
     psutil \
     requests \
     ujson \
+    tornado \
     xvfbwrapper \
     marionette_driver
 
diff --git a/docs/install.md b/docs/install.md
@@ -10,6 +10,7 @@ wptagent currently supports Windows, Linux and OSX for desktop browsers as well
     * pypiwin32 (Windows only)
     * requests
     * ujson
+    * tornado
     * xvfbwrapper (Linux only)
     * bind9utils (Linux only, for rndc)
     * marionette_driver (Firefox)
@@ -21,7 +22,7 @@ wptagent currently supports Windows, Linux and OSX for desktop browsers as well
 * cgroup-tools (Linux only if mobile CPU emulation is desired)
 * traceroute (Mac and Linux)
 * Debian:
-    * ```sudo apt-get install -y python2.7 python-pip imagemagick ffmpeg xvfb dbus-x11 cgroup-tools traceroute && sudo pip install dnspython monotonic pillow psutil requests ujson xvfbwrapper marionette_driver```
+    * ```sudo apt-get install -y python2.7 python-pip imagemagick ffmpeg xvfb dbus-x11 cgroup-tools traceroute && sudo pip install dnspython monotonic pillow psutil requests ujson tornado xvfbwrapper marionette_driver```
 * Chrome Browser
     * Linux stable, beta and unstable channels on Ubuntu/Debian:
         * ```wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | sudo apt-key add -```
diff --git a/internal/tornado_responder.py b/internal/tornado_responder.py
@@ -0,0 +1,74 @@
+# Copyright 2017 Google Inc. All rights reserved.
+# Use of this source code is governed by the Apache 2.0 license that can be
+# found in the LICENSE file.
+import logging
+import tornado.ioloop
+import tornado.web
+
+MESSAGE_SERVER = None
+
+class TornadoRequestHandler(tornado.web.RequestHandler):
+    """Request handler for when we are using tornado"""
+    def get(self):
+        """Handle GET requests"""
+        import ujson as json
+        logging.debug(self.request.uri)
+        response = None
+        content_type = 'text/plain'
+        if self.request.uri == '/ping':
+            response = 'pong'
+        elif self.request.uri == '/config':
+            # JSON config data
+            content_type = 'application/json'
+            response = '{}'
+            if MESSAGE_SERVER.config is not None:
+                response = json.dumps(MESSAGE_SERVER.config)
+        elif self.request.uri == '/config.html':
+            # HTML page that can be queried from the extension for config data
+            content_type = 'text/html'
+            response = "<html><head>\n"
+            if MESSAGE_SERVER.config is not None:
+                import cgi
+                response += '<div id="wptagentConfig" style="display: none;">'
+                response += cgi.escape(json.dumps(MESSAGE_SERVER.config))
+                response += '</div>'
+            response += "</head><body></body></html>"
+
+        if response is not None:
+            self.set_status(200)
+            self.set_header("Content-Type", content_type)
+            self.write(response)
+
+    def post(self):
+        """Handle POST messages"""
+        import ujson as json
+        try:
+            messages = self.request.body
+            if messages is not None and len(messages):
+                if self.request.uri == '/log':
+                    logging.debug(messages)
+                else:
+                    for line in messages.splitlines():
+                        line = line.strip()
+                        if len(line):
+                            message = json.loads(line)
+                            if 'body' not in message and self.request.uri != '/etw':
+                                message['body'] = None
+                            MESSAGE_SERVER.handle_message(message)
+        except Exception:
+            pass
+        self.set_status(200)
+
+def stop_tornado():
+    """Stop the tornado server"""
+    ioloop = tornado.ioloop.IOLoop.instance()
+    ioloop.add_callback(ioloop.stop)
+
+def start_tornado(message_server):
+    """Start (and run) the tornado server"""
+    global MESSAGE_SERVER
+    MESSAGE_SERVER = message_server
+    application = tornado.web.Application([(r"/.*", TornadoRequestHandler)])
+    application.listen(8888, '127.0.0.1')
+    message_server.is_started.set()
+    tornado.ioloop.IOLoop.instance().start()
diff --git a/ubuntu_install.sh b/ubuntu_install.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 sudo apt-get install -y python2.7 python-pip imagemagick ffmpeg xvfb dbus-x11 cgroup-tools traceroute software-properties-common python-software-properties psmisc
 sudo dbus-uuidgen --ensure
-sudo pip install dnspython monotonic pillow psutil requests ujson xvfbwrapper marionette_driver
+sudo pip install dnspython monotonic pillow psutil requests ujson tornado xvfbwrapper marionette_driver
 curl -sL https://deb.nodesource.com/setup_7.x | sudo bash -
 sudo apt-get install -y nodejs
 sudo npm install -g lighthouse
diff --git a/wptagent.py b/wptagent.py
@@ -77,6 +77,9 @@ def run_testing(self):
         if not self.options.android and not self.options.iOS:
             message_server = MessageServer()
             message_server.start()
+            if not message_server.is_ok():
+                logging.error("Unable to start the local message server")
+                return
         while not self.must_exit:
             try:
                 if os.path.isfile(exit_file):
@@ -293,6 +296,13 @@ def startup(self):
                   "and make sure it is in the path."
             ret = False
 
+        try:
+            import tornado as _
+        except ImportError:
+            from internal.os_util import run_elevated
+            logging.debug('Trying to install tornado...')
+            run_elevated(sys.executable, '-m pip install tornado')
+
         if platform.system() == "Linux":
             try:
                 subprocess.check_output(['traceroute', '--version'])
@@ -366,7 +376,6 @@ def startup(self):
                 ret = False
         return ret
 
-
 class HandleMessage(BaseHTTPRequestHandler):
     """Handle a single message from the extension"""
     # 1.1 keep-alive is pretty buggy, leave this disabled for now
@@ -480,7 +489,8 @@ def __init__(self):
         self.thread = None
         self.messages = Queue.Queue()
         self.config = None
-        self.__is_started = threading.Event()
+        self.is_started = threading.Event()
+        self.using_tornado = False
 
     def get_message(self, timeout):
         """Get a single message from the queue"""
@@ -503,17 +513,21 @@ def handle_message(self, message):
 
     def start(self):
         """Start running the server in a background thread"""
-        self.__is_started.clear()
+        self.is_started.clear()
         self.thread = threading.Thread(target=self.run)
         self.thread.daemon = True
         self.thread.start()
-        self.__is_started.wait(timeout=30)
+        self.is_started.wait(timeout=30)
 
     def stop(self):
         """Stop running the server"""
         logging.debug("Shutting down extension server")
         self.must_exit = True
-        if self.server is not None:
+        if self.using_tornado:
+            from internal.tornado_responder import stop_tornado
+            stop_tornado()
+            self.thread.join()
+        elif self.server is not None:
             try:
                 self.server.shutdown()
             except Exception:
@@ -543,18 +557,22 @@ def run(self):
         handler = handler_template(self)
         logging.debug('Starting extension server on port 8888')
         try:
-            self.server = HTTPServer(('127.0.0.1', 8888), handler)
-            self.server.timeout = 10
-            self.__is_started.set()
-            self.server.serve_forever()
-        except Exception:
-            logging.exception("Message server main loop exception")
-            self.__is_started.set()
-        if self.server is not None:
+            from internal.tornado_responder import start_tornado
+            start_tornado(self)
+        except ImportError:
             try:
-                self.server.socket.close()
+                self.server = HTTPServer(('127.0.0.1', 8888), handler)
+                self.server.timeout = 10
+                self.is_started.set()
+                self.server.serve_forever()
             except Exception:
-                pass
+                logging.exception("Message server main loop exception")
+                self.is_started.set()
+            if self.server is not None:
+                try:
+                    self.server.socket.close()
+                except Exception:
+                    pass
 
 def parse_ini(ini):
     """Parse an ini file and convert it to a dictionary"""
