1#!/usr/bin/env python
2# Copyright (c) 2012 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Constrained Network Server. Serves files with supplied network constraints.
7
8The CNS exposes a web based API allowing network constraints to be imposed on
9file serving.
10
11TODO(dalecurtis): Add some more docs here.
12
13"""
14
15import logging
16from logging import handlers
17import mimetypes
18import optparse
19import os
20import signal
21import sys
22import threading
23import time
24import urllib
25import urllib2
26
27import traffic_control
28
29try:
30  import cherrypy
31except ImportError:
32  print ('CNS requires CherryPy v3 or higher to be installed. Please install\n'
33         'and try again. On Linux: sudo apt-get install python-cherrypy3\n')
34  sys.exit(1)
35
36# Add webm file types to mimetypes map since cherrypy's default type is text.
37mimetypes.types_map['.webm'] = 'video/webm'
38
39# Default logging is ERROR. Use --verbose to enable DEBUG logging.
40_DEFAULT_LOG_LEVEL = logging.ERROR
41
42# Default port to serve the CNS on.
43_DEFAULT_SERVING_PORT = 9000
44
45# Default port range for constrained use.
46_DEFAULT_CNS_PORT_RANGE = (50000, 51000)
47
48# Default number of seconds before a port can be torn down.
49_DEFAULT_PORT_EXPIRY_TIME_SECS = 5 * 60
50
51
52class PortAllocator(object):
53  """Dynamically allocates/deallocates ports with a given set of constraints."""
54
55  def __init__(self, port_range, expiry_time_secs=5 * 60):
56    """Sets up initial state for the Port Allocator.
57
58    Args:
59      port_range: Range of ports available for allocation.
60      expiry_time_secs: Amount of time in seconds before constrained ports are
61          cleaned up.
62    """
63    self._port_range = port_range
64    self._expiry_time_secs = expiry_time_secs
65
66    # Keeps track of ports we've used, the creation key, and the last request
67    # time for the port so they can be cached and cleaned up later.
68    self._ports = {}
69
70    # Locks port creation and cleanup. TODO(dalecurtis): If performance becomes
71    # an issue a per-port based lock system can be used instead.
72    self._port_lock = threading.RLock()
73
74  def Get(self, key, new_port=False, **kwargs):
75    """Sets up a constrained port using the requested parameters.
76
77    Requests for the same key and constraints will result in a cached port being
78    returned if possible, subject to new_port.
79
80    Args:
81      key: Used to cache ports with the given constraints.
82      new_port: Whether to create a new port or use an existing one if possible.
83      **kwargs: Constraints to pass into traffic control.
84
85    Returns:
86      None if no port can be setup or the port number of the constrained port.
87    """
88    with self._port_lock:
89      # Check port key cache to see if this port is already setup. Update the
90      # cache time and return the port if so. Performance isn't a concern here,
91      # so just iterate over ports dict for simplicity.
92      full_key = (key,) + tuple(kwargs.values())
93      if not new_port:
94        for port, status in self._ports.iteritems():
95          if full_key == status['key']:
96            self._ports[port]['last_update'] = time.time()
97            return port
98
99      # Cleanup ports on new port requests. Do it after the cache check though
100      # so we don't erase and then setup the same port.
101      if self._expiry_time_secs > 0:
102        self.Cleanup(all_ports=False)
103
104      # Performance isn't really an issue here, so just iterate over the port
105      # range to find an unused port. If no port is found, None is returned.
106      for port in xrange(self._port_range[0], self._port_range[1]):
107        if port in self._ports:
108          continue
109        if self._SetupPort(port, **kwargs):
110          kwargs['port'] = port
111          self._ports[port] = {'last_update': time.time(), 'key': full_key,
112                               'config': kwargs}
113          return port
114
115  def _SetupPort(self, port, **kwargs):
116    """Setup network constraints on port using the requested parameters.
117
118    Args:
119      port: The port number to setup network constraints on.
120      **kwargs: Network constraints to set up on the port.
121
122    Returns:
123      True if setting the network constraints on the port was successful, false
124      otherwise.
125    """
126    kwargs['port'] = port
127    try:
128      cherrypy.log('Setting up port %d' % port)
129      traffic_control.CreateConstrainedPort(kwargs)
130      return True
131    except traffic_control.TrafficControlError as e:
132      cherrypy.log('Error: %s\nOutput: %s' % (e.msg, e.error))
133      return False
134
135  def Cleanup(self, all_ports, request_ip=None):
136    """Cleans up expired ports, or if all_ports=True, all allocated ports.
137
138    By default, ports which haven't been used for self._expiry_time_secs are
139    torn down. If all_ports=True then they are torn down regardless.
140
141    Args:
142      all_ports: Should all ports be torn down regardless of expiration?
143      request_ip: Tear ports matching the IP address regarless of expiration.
144    """
145    with self._port_lock:
146      now = time.time()
147      # Use .items() instead of .iteritems() so we can delete keys w/o error.
148      for port, status in self._ports.items():
149        expired = now - status['last_update'] > self._expiry_time_secs
150        matching_ip = request_ip and status['key'][0].startswith(request_ip)
151        if all_ports or expired or matching_ip:
152          cherrypy.log('Cleaning up port %d' % port)
153          self._DeletePort(port)
154          del self._ports[port]
155
156  def _DeletePort(self, port):
157    """Deletes network constraints on port.
158
159    Args:
160      port: The port number associated with the network constraints.
161    """
162    try:
163      traffic_control.DeleteConstrainedPort(self._ports[port]['config'])
164    except traffic_control.TrafficControlError as e:
165      cherrypy.log('Error: %s\nOutput: %s' % (e.msg, e.error))
166
167
168class ConstrainedNetworkServer(object):
169  """A CherryPy-based HTTP server for serving files with network constraints."""
170
171  def __init__(self, options, port_allocator):
172    """Sets up initial state for the CNS.
173
174    Args:
175      options: optparse based class returned by ParseArgs()
176      port_allocator: A port allocator instance.
177    """
178    self._options = options
179    self._port_allocator = port_allocator
180
181  @cherrypy.expose
182  def Cleanup(self):
183    """Cleans up all the ports allocated using the request IP address.
184
185    When requesting a constrained port, the cherrypy.request.remote.ip is used
186    as a key for that port (in addition to other request parameters).  Such
187    ports created for the same IP address are removed.
188    """
189    cherrypy.log('Cleaning up ports allocated by %s.' %
190                 cherrypy.request.remote.ip)
191    self._port_allocator.Cleanup(all_ports=False,
192                                 request_ip=cherrypy.request.remote.ip)
193
194  @cherrypy.expose
195  def ServeConstrained(self, f=None, bandwidth=None, latency=None, loss=None,
196                       new_port=False, no_cache=False, **kwargs):
197    """Serves the requested file with the requested constraints.
198
199    Subsequent requests for the same constraints from the same IP will share the
200    previously created port unless new_port equals True. If no constraints
201    are provided the file is served as is.
202
203    Args:
204      f: path relative to http root of file to serve.
205      bandwidth: maximum allowed bandwidth for the provided port (integer
206          in kbit/s).
207      latency: time to add to each packet (integer in ms).
208      loss: percentage of packets to drop (integer, 0-100).
209      new_port: whether to use a new port for this request or not.
210      no_cache: Set reponse's cache-control to no-cache.
211    """
212    if no_cache:
213      response = cherrypy.response
214      response.headers['Pragma'] = 'no-cache'
215      response.headers['Cache-Control'] = 'no-cache'
216
217    # CherryPy is a bit wonky at detecting parameters, so just make them all
218    # optional and validate them ourselves.
219    if not f:
220      raise cherrypy.HTTPError(400, 'Invalid request. File must be specified.')
221
222    # Check existence early to prevent wasted constraint setup.
223    self._CheckRequestedFileExist(f)
224
225    # If there are no constraints, just serve the file.
226    if bandwidth is None and latency is None and loss is None:
227      return self._ServeFile(f)
228
229    constrained_port = self._GetConstrainedPort(
230        f, bandwidth=bandwidth, latency=latency, loss=loss, new_port=new_port,
231        **kwargs)
232
233    # Build constrained URL using the constrained port and original URL
234    # parameters except the network constraints (bandwidth, latency, and loss).
235    constrained_url = self._GetServerURL(f, constrained_port,
236                                         no_cache=no_cache, **kwargs)
237
238    # Redirect request to the constrained port.
239    cherrypy.log('Redirect to %s' % constrained_url)
240    cherrypy.lib.cptools.redirect(constrained_url, internal=False)
241
242  def _CheckRequestedFileExist(self, f):
243    """Checks if the requested file exists, raises HTTPError otherwise."""
244    if self._options.local_server_port:
245      self._CheckFileExistOnLocalServer(f)
246    else:
247      self._CheckFileExistOnServer(f)
248
249  def _CheckFileExistOnServer(self, f):
250    """Checks if requested file f exists to be served by this server."""
251    # Sanitize and check the path to prevent www-root escapes.
252    sanitized_path = os.path.abspath(os.path.join(self._options.www_root, f))
253    if not sanitized_path.startswith(self._options.www_root):
254      raise cherrypy.HTTPError(403, 'Invalid file requested.')
255    if not os.path.exists(sanitized_path):
256      raise cherrypy.HTTPError(404, 'File not found.')
257
258  def _CheckFileExistOnLocalServer(self, f):
259    """Checks if requested file exists on local server hosting files."""
260    test_url = self._GetServerURL(f, self._options.local_server_port)
261    try:
262      cherrypy.log('Check file exist using URL: %s' % test_url)
263      return urllib2.urlopen(test_url) is not None
264    except Exception:
265      raise cherrypy.HTTPError(404, 'File not found on local server.')
266
267  def _ServeFile(self, f):
268    """Serves the file as an http response."""
269    if self._options.local_server_port:
270      redirect_url = self._GetServerURL(f, self._options.local_server_port)
271      cherrypy.log('Redirect to %s' % redirect_url)
272      cherrypy.lib.cptools.redirect(redirect_url, internal=False)
273    else:
274      sanitized_path = os.path.abspath(os.path.join(self._options.www_root, f))
275      return cherrypy.lib.static.serve_file(sanitized_path)
276
277  def _GetServerURL(self, f, port, **kwargs):
278    """Returns a URL for local server to serve the file on given port.
279
280    Args:
281      f: file name to serve on local server. Relative to www_root.
282      port: Local server port (it can be a configured constrained port).
283      kwargs: extra parameteres passed in the URL.
284    """
285    url = '%s?f=%s&' % (cherrypy.url(), f)
286    if self._options.local_server_port:
287      url = '%s/%s?' % (
288          cherrypy.url().replace('ServeConstrained', self._options.www_root), f)
289
290    url = url.replace(':%d' % self._options.port, ':%d' % port)
291    extra_args = urllib.urlencode(kwargs)
292    if extra_args:
293      url += extra_args
294    return url
295
296  def _GetConstrainedPort(self, f=None, bandwidth=None, latency=None, loss=None,
297                          new_port=False, **kwargs):
298    """Creates or gets a port with specified network constraints.
299
300    See ServeConstrained() for more details.
301    """
302    # Validate inputs. isdigit() guarantees a natural number.
303    bandwidth = self._ParseIntParameter(
304        bandwidth, 'Invalid bandwidth constraint.', lambda x: x > 0)
305    latency = self._ParseIntParameter(
306        latency, 'Invalid latency constraint.', lambda x: x >= 0)
307    loss = self._ParseIntParameter(
308        loss, 'Invalid loss constraint.', lambda x: x <= 100 and x >= 0)
309
310    redirect_port = self._options.port
311    if self._options.local_server_port:
312      redirect_port = self._options.local_server_port
313
314    start_time = time.time()
315    # Allocate a port using the given constraints. If a port with the requested
316    # key and kwargs already exist then reuse that port.
317    constrained_port = self._port_allocator.Get(
318        cherrypy.request.remote.ip, server_port=redirect_port,
319        interface=self._options.interface, bandwidth=bandwidth, latency=latency,
320        loss=loss, new_port=new_port, file=f, **kwargs)
321
322    cherrypy.log('Time to set up port %d = %.3fsec.' %
323                 (constrained_port, time.time() - start_time))
324
325    if not constrained_port:
326      raise cherrypy.HTTPError(503, 'Service unavailable. Out of ports.')
327    return constrained_port
328
329  def _ParseIntParameter(self, param, msg, check):
330    """Returns integer value of param and verifies it satisfies the check.
331
332    Args:
333      param: Parameter name to check.
334      msg: Message in error if raised.
335      check: Check to verify the parameter value.
336
337    Returns:
338      None if param is None, integer value of param otherwise.
339
340    Raises:
341      cherrypy.HTTPError if param can not be converted to integer or if it does
342      not satisfy the check.
343    """
344    if param:
345      try:
346        int_value = int(param)
347        if check(int_value):
348          return int_value
349      except:
350        pass
351      raise cherrypy.HTTPError(400, msg)
352
353
354def ParseArgs():
355  """Define and parse the command-line arguments."""
356  parser = optparse.OptionParser()
357
358  parser.add_option('--expiry-time', type='int',
359                    default=_DEFAULT_PORT_EXPIRY_TIME_SECS,
360                    help=('Number of seconds before constrained ports expire '
361                          'and are cleaned up. 0=Disabled. Default: %default'))
362  parser.add_option('--port', type='int', default=_DEFAULT_SERVING_PORT,
363                    help='Port to serve the API on. Default: %default')
364  parser.add_option('--port-range', default=_DEFAULT_CNS_PORT_RANGE,
365                    help=('Range of ports for constrained serving. Specify as '
366                          'a comma separated value pair. Default: %default'))
367  parser.add_option('--interface', default='eth0',
368                    help=('Interface to setup constraints on. Use lo for a '
369                          'local client. Default: %default'))
370  parser.add_option('--socket-timeout', type='int',
371                    default=cherrypy.server.socket_timeout,
372                    help=('Number of seconds before a socket connection times '
373                          'out. Default: %default'))
374  parser.add_option('--threads', type='int',
375                    default=cherrypy._cpserver.Server.thread_pool,
376                    help=('Number of threads in the thread pool. Default: '
377                          '%default'))
378  parser.add_option('--www-root', default='',
379                    help=('Directory root to serve files from. If --local-'
380                          'server-port is used, the path is appended to the '
381                          'redirected URL of local server. Defaults to the '
382                          'current directory (if --local-server-port is not '
383                          'used): %s' % os.getcwd()))
384  parser.add_option('--local-server-port', type='int',
385                    help=('Optional local server port to host files.'))
386  parser.add_option('-v', '--verbose', action='store_true', default=False,
387                    help='Turn on verbose output.')
388
389  options = parser.parse_args()[0]
390
391  # Convert port range into the desired tuple format.
392  try:
393    if isinstance(options.port_range, str):
394      options.port_range = [int(port) for port in options.port_range.split(',')]
395  except ValueError:
396    parser.error('Invalid port range specified.')
397
398  if options.expiry_time < 0:
399    parser.error('Invalid expiry time specified.')
400
401  # Convert the path to an absolute to remove any . or ..
402  if not options.local_server_port:
403    if not options.www_root:
404      options.www_root = os.getcwd()
405    options.www_root = os.path.abspath(options.www_root)
406
407  _SetLogger(options.verbose)
408
409  return options
410
411
412def _SetLogger(verbose):
413  file_handler = handlers.RotatingFileHandler('cns.log', 'a', 10000000, 10)
414  file_handler.setFormatter(logging.Formatter('[%(threadName)s] %(message)s'))
415
416  log_level = _DEFAULT_LOG_LEVEL
417  if verbose:
418    log_level = logging.DEBUG
419  file_handler.setLevel(log_level)
420
421  cherrypy.log.error_log.addHandler(file_handler)
422  cherrypy.log.access_log.addHandler(file_handler)
423
424
425def Main():
426  """Configure and start the ConstrainedNetworkServer."""
427  options = ParseArgs()
428
429  try:
430    traffic_control.CheckRequirements()
431  except traffic_control.TrafficControlError as e:
432    cherrypy.log(e.msg)
433    return
434
435  cherrypy.config.update({'server.socket_host': '::',
436                          'server.socket_port': options.port})
437
438  if options.threads:
439    cherrypy.config.update({'server.thread_pool': options.threads})
440
441  if options.socket_timeout:
442    cherrypy.config.update({'server.socket_timeout': options.socket_timeout})
443
444  # Setup port allocator here so we can call cleanup on failures/exit.
445  pa = PortAllocator(options.port_range, expiry_time_secs=options.expiry_time)
446
447  try:
448    cherrypy.quickstart(ConstrainedNetworkServer(options, pa))
449  finally:
450    # Disable Ctrl-C handler to prevent interruption of cleanup.
451    signal.signal(signal.SIGINT, lambda signal, frame: None)
452    pa.Cleanup(all_ports=True)
453
454
455if __name__ == '__main__':
456  Main()
457