src/chrome/test/chromedriver/test/webserver.py

   1 # Copyright 2013 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 import BaseHTTPServer
   6 import os
   7 import threading
   8
   9
  10 class Responder(object):
  11   """Sends a HTTP response. Used with TestWebServer."""
  12
  13   def __init__(self, handler):
  14     self._handler = handler
  15
  16   def SendResponse(self, body):
  17     """Sends OK response with body."""
  18     self.SendHeaders(len(body))
  19     self.SendBody(body)
  20
  21   def SendResponseFromFile(self, path):
  22     """Sends OK response with the given file as the body."""
  23     with open(path, 'r') as f:
  24       self.SendResponse(f.read())
  25
  26   def SendHeaders(self, content_length=None):
  27     """Sends headers for OK response."""
  28     self._handler.send_response(200)
  29     if content_length:
  30       self._handler.send_header('Content-Length', content_length)
  31     self._handler.end_headers()
  32
  33   def SendError(self, code):
  34     """Sends response for the given HTTP error code."""
  35     self._handler.send_error(code)
  36
  37   def SendBody(self, body):
  38     """Just sends the body, no headers."""
  39     self._handler.wfile.write(body)
  40
  41
  42 class Request(object):
  43   """An HTTP request."""
  44
  45   def __init__(self, handler):
  46     self._handler = handler
  47
  48   def GetPath(self):
  49     return self._handler.path
  50
  51
  52 class _BaseServer(BaseHTTPServer.HTTPServer):
  53   """Internal server that throws if timed out waiting for a request."""
  54
  55   def __init__(self, on_request, server_cert_and_key_path=None):
  56     """Starts the server.
  57
  58     It is an HTTP server if parameter server_cert_and_key_path is not provided.
  59     Otherwise, it is an HTTPS server.
  60
  61     Args:
  62       server_cert_and_key_path: path to a PEM file containing the cert and key.
  63                                 if it is None, start the server as an HTTP one.
  64     """
  65     class _Handler(BaseHTTPServer.BaseHTTPRequestHandler):
  66       """Internal handler that just asks the server to handle the request."""
  67
  68       def do_GET(self):
  69         if self.path.endswith('favicon.ico'):
  70           self.send_error(404)
  71           return
  72         on_request(Request(self), Responder(self))
  73
  74       def log_message(self, *args, **kwargs):
  75         """Overriddes base class method to disable logging."""
  76         pass
  77
  78     BaseHTTPServer.HTTPServer.__init__(self, ('127.0.0.1', 0), _Handler)
  79
  80     if server_cert_and_key_path is not None:
  81       self._is_https_enabled = True
  82       self._server.socket = ssl.wrap_socket(
  83           self._server.socket, certfile=server_cert_and_key_path,
  84           server_side=True)
  85     else:
  86       self._is_https_enabled = False
  87
  88   def handle_timeout(self):
  89     """Overridden from SocketServer."""
  90     raise RuntimeError('Timed out waiting for http request')
  91
  92   def GetUrl(self):
  93     """Returns the base URL of the server."""
  94     postfix = '://127.0.0.1:%s' % self.server_port
  95     if self._is_https_enabled:
  96       return 'https' + postfix
  97     return 'http' + postfix
  98
  99
 100 class WebServer(object):
 101   """An HTTP or HTTPS server that serves on its own thread.
 102
 103   Serves files from given directory but may use custom data for specific paths.
 104   """
 105
 106   def __init__(self, root_dir, server_cert_and_key_path=None):
 107     """Starts the server.
 108
 109     It is an HTTP server if parameter server_cert_and_key_path is not provided.
 110     Otherwise, it is an HTTPS server.
 111
 112     Args:
 113       root_dir: root path to serve files from. This parameter is required.
 114       server_cert_and_key_path: path to a PEM file containing the cert and key.
 115                                 if it is None, start the server as an HTTP one.
 116     """
 117     self._root_dir = os.path.abspath(root_dir)
 118     self._server = _BaseServer(self._OnRequest, server_cert_and_key_path)
 119     self._thread = threading.Thread(target=self._server.serve_forever)
 120     self._thread.daemon = True
 121     self._thread.start()
 122     self._path_data_map = {}
 123     self._path_data_lock = threading.Lock()
 124
 125   def _OnRequest(self, request, responder):
 126     path = request.GetPath().split('?')[0]
 127
 128     # Serve from path -> data map.
 129     self._path_data_lock.acquire()
 130     try:
 131       if path in self._path_data_map:
 132         responder.SendResponse(self._path_data_map[path])
 133         return
 134     finally:
 135       self._path_data_lock.release()
 136
 137     # Serve from file.
 138     path = os.path.normpath(
 139         os.path.join(self._root_dir, *path.split('/')))
 140     if not path.startswith(self._root_dir):
 141       responder.SendError(403)
 142       return
 143     if not os.path.exists(path):
 144       responder.SendError(404)
 145       return
 146     responder.SendResponseFromFile(path)
 147
 148   def SetDataForPath(self, path, data):
 149     self._path_data_lock.acquire()
 150     try:
 151       self._path_data_map[path] = data
 152     finally:
 153       self._path_data_lock.release()
 154
 155   def GetUrl(self):
 156     """Returns the base URL of the server."""
 157     return self._server.GetUrl()
 158
 159   def Shutdown(self):
 160     """Shuts down the server synchronously."""
 161     self._server.shutdown()
 162     self._thread.join()
 163
 164
 165 class SyncWebServer(object):
 166   """WebServer for testing.
 167
 168   Incoming requests are blocked until explicitly handled.
 169   This was designed for single thread use. All requests should be handled on
 170   the same thread.
 171   """
 172
 173   def __init__(self):
 174     self._server = _BaseServer(self._OnRequest)
 175     # Recognized by SocketServer.
 176     self._server.timeout = 10
 177     self._on_request = None
 178
 179   def _OnRequest(self, request, responder):
 180     self._on_request(responder)
 181     self._on_request = None
 182
 183   def Respond(self, on_request):
 184     """Blocks until request comes in, then calls given handler function.
 185
 186     Args:
 187       on_request: Function that handles the request. Invoked with single
 188           parameter, an instance of Responder.
 189     """
 190     if self._on_request:
 191       raise RuntimeError('Must handle 1 request at a time.')
 192
 193     self._on_request = on_request
 194     while self._on_request:
 195       # Don't use handle_one_request, because it won't work with the timeout.
 196       self._server.handle_request()
 197
 198   def RespondWithContent(self, content):
 199     """Blocks until request comes in, then handles it with the given content."""
 200     def SendContent(responder):
 201       responder.SendResponse(content)
 202     self.Respond(SendContent)
 203
 204   def GetUrl(self):
 205     return self._server.GetUrl()