1# Copyright 2014 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
5import collections
6import json
7import os
8import re
9import subprocess
10import sys
11
12from telemetry.core import forwarders
13from telemetry.core import util
14
15NamedPort = collections.namedtuple('NamedPort', ['name', 'port'])
16
17
18class LocalServerBackend(object):
19  def __init__(self):
20    pass
21
22  def StartAndGetNamedPorts(self, args):
23    """Starts the actual server and obtains any sockets on which it
24    should listen.
25
26    Returns a list of NamedPort on which this backend is listening.
27    """
28    raise NotImplementedError()
29
30  def ServeForever(self):
31    raise NotImplementedError()
32
33
34class LocalServer(object):
35  def __init__(self, server_backend_class):
36    assert LocalServerBackend in server_backend_class.__bases__
37    server_module_name = server_backend_class.__module__
38    assert server_module_name in sys.modules, \
39            'The server class\' module must be findable via sys.modules'
40    assert getattr(sys.modules[server_module_name],
41                   server_backend_class.__name__), \
42      'The server class must getattrable from its __module__ by its __name__'
43
44    self._server_backend_class = server_backend_class
45    self._subprocess = None
46    self._devnull = None
47    self._local_server_controller = None
48    self.forwarder = None
49    self.host_ip = None
50
51  def Start(self, local_server_controller):
52    assert self._subprocess == None
53    self._local_server_controller = local_server_controller
54
55    self.host_ip = local_server_controller.host_ip
56
57    server_args = self.GetBackendStartupArgs()
58    server_args_as_json = json.dumps(server_args)
59    server_module_name = self._server_backend_class.__module__
60
61    self._devnull = open(os.devnull, 'w')
62    cmd = [
63        sys.executable, '-m', __name__,
64        'run_backend',
65        server_module_name,
66        self._server_backend_class.__name__,
67        server_args_as_json,
68        ]
69
70    env = os.environ.copy()
71    env['PYTHONPATH'] = os.pathsep.join(sys.path)
72
73    self._subprocess = subprocess.Popen(
74        cmd, cwd=util.GetTelemetryDir(), env=env, stdout=subprocess.PIPE)
75
76    named_ports = self._GetNamedPortsFromBackend()
77    named_port_pair_map = {'http': None, 'https': None, 'dns': None}
78    for name, port in named_ports:
79      assert name in named_port_pair_map, '%s forwarding is unsupported' % name
80      named_port_pair_map[name] = (
81          forwarders.PortPair(port,
82                              local_server_controller.GetRemotePort(port)))
83    self.forwarder = local_server_controller.CreateForwarder(
84        forwarders.PortPairs(**named_port_pair_map))
85
86  def _GetNamedPortsFromBackend(self):
87    named_ports_json = None
88    named_ports_re = re.compile('LocalServerBackend started: (?P<port>.+)')
89    # TODO: This will hang if the subprocess doesn't print the correct output.
90    while self._subprocess.poll() == None:
91      m = named_ports_re.match(self._subprocess.stdout.readline())
92      if m:
93        named_ports_json = m.group('port')
94        break
95
96    if not named_ports_json:
97      raise Exception('Server process died prematurely ' +
98                      'without giving us port pairs.')
99    return [NamedPort(**pair) for pair in json.loads(named_ports_json.lower())]
100
101  @property
102  def is_running(self):
103    return self._subprocess != None
104
105  def __enter__(self):
106    return self
107
108  def __exit__(self, *args):
109    self.Close()
110
111  def __del__(self):
112    self.Close()
113
114  def Close(self):
115    if self.forwarder:
116      self.forwarder.Close()
117      self.forwarder = None
118    if self._subprocess:
119      # TODO(tonyg): Should this block until it goes away?
120      self._subprocess.kill()
121      self._subprocess = None
122    if self._devnull:
123      self._devnull.close()
124      self._devnull = None
125    if self._local_server_controller:
126      self._local_server_controller.ServerDidClose(self)
127      self._local_server_controller = None
128
129  def GetBackendStartupArgs(self):
130    """Returns whatever arguments are required to start up the backend"""
131    raise NotImplementedError()
132
133
134class LocalServerController():
135  """Manages the list of running servers
136
137  This class manages the running servers, but also provides an isolation layer
138  to prevent LocalServer subclasses from accessing the browser backend directly.
139
140  """
141  def __init__(self, browser_backend):
142    self._browser_backend = browser_backend
143    self._local_servers_by_class = {}
144    self.host_ip = self._browser_backend.forwarder_factory.host_ip
145
146  def StartServer(self, server):
147    assert not server.is_running, 'Server already started'
148    assert isinstance(server, LocalServer)
149    if server.__class__ in self._local_servers_by_class:
150      raise Exception(
151          'Canont have two servers of the same class running at once. ' +
152          'Locate the existing one and use it, or call Close() on it.')
153
154    server.Start(self)
155    self._local_servers_by_class[server.__class__] = server
156
157  def GetRunningServer(self, server_class, default_value):
158    return self._local_servers_by_class.get(server_class, default_value)
159
160  @property
161  def local_servers(self):
162    return self._local_servers_by_class.values()
163
164  def Close(self):
165    while len(self._local_servers_by_class):
166      server = self._local_servers_by_class.itervalues().next()
167      try:
168        server.Close()
169      except Exception:
170        import traceback
171        traceback.print_exc()
172
173  def CreateForwarder(self, port_pairs):
174    return self._browser_backend.forwarder_factory.Create(port_pairs)
175
176  def GetRemotePort(self, port):
177    return self._browser_backend.GetRemotePort(port)
178
179  def ServerDidClose(self, server):
180    del self._local_servers_by_class[server.__class__]
181
182
183def _LocalServerBackendMain(args):
184  assert len(args) == 4
185  (cmd, server_module_name,
186   server_backend_class_name, server_args_as_json) = args[:4]
187  assert cmd == 'run_backend'
188  server_module = __import__(server_module_name, fromlist=[True])
189  server_backend_class = getattr(server_module, server_backend_class_name)
190  server = server_backend_class()
191
192  server_args = json.loads(server_args_as_json)
193
194  named_ports = server.StartAndGetNamedPorts(server_args)
195  assert isinstance(named_ports, list)
196  for named_port in named_ports:
197    assert isinstance(named_port, NamedPort)
198
199  # Note: This message is scraped by the parent process'
200  # _GetNamedPortsFromBackend(). Do **not** change it.
201  print 'LocalServerBackend started: %s' % json.dumps(
202      [pair._asdict() for pair in named_ports]) # pylint: disable=W0212
203  sys.stdout.flush()
204
205  return server.ServeForever()
206
207
208if __name__ == '__main__':
209  # This trick is needed because local_server.NamedPort is not the
210  # same as sys.modules['__main__'].NamedPort. The module itself is loaded
211  # twice, basically.
212  from telemetry.core import local_server # pylint: disable=W0406
213  sys.exit(local_server._LocalServerBackendMain( # pylint: disable=W0212
214      sys.argv[1:]))
215