Home | History | Annotate | Download | only in pylib 
      1 # Copyright (c) 2012 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 """Functions that deal with local and device ports."""
      6 
      7 import contextlib
      8 import fcntl
      9 import httplib
     10 import logging
     11 import os
     12 import re
     13 import socket
     14 import traceback
     15 
     16 from pylib import cmd_helper
     17 from pylib import constants
     18 
     19 
     20 # The following two methods are used to allocate the port source for various
     21 # types of test servers. Because some net-related tests can be run on shards at
     22 # same time, it's important to have a mechanism to allocate the port
     23 # process-safe. In here, we implement the safe port allocation by leveraging
     24 # flock.
     25 def ResetTestServerPortAllocation():
     26   """Resets the port allocation to start from TEST_SERVER_PORT_FIRST.
     27 
     28   Returns:
     29     Returns True if reset successes. Otherwise returns False.
     30   """
     31   try:
     32     with open(constants.TEST_SERVER_PORT_FILE, 'w') as fp:
     33       fp.write('%d' % constants.TEST_SERVER_PORT_FIRST)
     34     if os.path.exists(constants.TEST_SERVER_PORT_LOCKFILE):
     35       os.unlink(constants.TEST_SERVER_PORT_LOCKFILE)
     36     return True
     37   except Exception as e:
     38     logging.error(e)
     39   return False
     40 
     41 
     42 def AllocateTestServerPort():
     43   """Allocates a port incrementally.
     44 
     45   Returns:
     46     Returns a valid port which should be in between TEST_SERVER_PORT_FIRST and
     47     TEST_SERVER_PORT_LAST. Returning 0 means no more valid port can be used.
     48   """
     49   port = 0
     50   ports_tried = []
     51   try:
     52     fp_lock = open(constants.TEST_SERVER_PORT_LOCKFILE, 'w')
     53     fcntl.flock(fp_lock, fcntl.LOCK_EX)
     54     # Get current valid port and calculate next valid port.
     55     if not os.path.exists(constants.TEST_SERVER_PORT_FILE):
     56       ResetTestServerPortAllocation()
     57     with open(constants.TEST_SERVER_PORT_FILE, 'r+') as fp:
     58       port = int(fp.read())
     59       ports_tried.append(port)
     60       while IsHostPortUsed(port):
     61         port += 1
     62         ports_tried.append(port)
     63       if (port > constants.TEST_SERVER_PORT_LAST or
     64           port < constants.TEST_SERVER_PORT_FIRST):
     65         port = 0
     66       else:
     67         fp.seek(0, os.SEEK_SET)
     68         fp.write('%d' % (port + 1))
     69   except Exception as e:
     70     logging.info(e)
     71   finally:
     72     if fp_lock:
     73       fcntl.flock(fp_lock, fcntl.LOCK_UN)
     74       fp_lock.close()
     75   if port:
     76     logging.info('Allocate port %d for test server.', port)
     77   else:
     78     logging.error('Could not allocate port for test server. '
     79                   'List of ports tried: %s', str(ports_tried))
     80   return port
     81 
     82 
     83 def IsHostPortUsed(host_port):
     84   """Checks whether the specified host port is used or not.
     85 
     86   Uses -n -P to inhibit the conversion of host/port numbers to host/port names.
     87 
     88   Args:
     89     host_port: Port on host we want to check.
     90 
     91   Returns:
     92     True if the port on host is already used, otherwise returns False.
     93   """
     94   port_info = '(\*)|(127\.0\.0\.1)|(localhost):%d' % host_port
     95   # TODO(jnd): Find a better way to filter the port. Note that connecting to the
     96   # socket and closing it would leave it in the TIME_WAIT state. Setting
     97   # SO_LINGER on it and then closing it makes the Python HTTP server crash.
     98   re_port = re.compile(port_info, re.MULTILINE)
     99   if re_port.search(cmd_helper.GetCmdOutput(['lsof', '-nPi:%d' % host_port])):
    100     return True
    101   return False
    102 
    103 
    104 def IsDevicePortUsed(device, device_port, state=''):
    105   """Checks whether the specified device port is used or not.
    106 
    107   Args:
    108     device: A DeviceUtils instance.
    109     device_port: Port on device we want to check.
    110     state: String of the specified state. Default is empty string, which
    111            means any state.
    112 
    113   Returns:
    114     True if the port on device is already used, otherwise returns False.
    115   """
    116   base_url = '127.0.0.1:%d' % device_port
    117   netstat_results = device.RunShellCommand('netstat')
    118   for single_connect in netstat_results:
    119     # Column 3 is the local address which we want to check with.
    120     connect_results = single_connect.split()
    121     if connect_results[0] != 'tcp':
    122       continue
    123     if len(connect_results) < 6:
    124       raise Exception('Unexpected format while parsing netstat line: ' +
    125                       single_connect)
    126     is_state_match = connect_results[5] == state if state else True
    127     if connect_results[3] == base_url and is_state_match:
    128       return True
    129   return False
    130 
    131 
    132 def IsHttpServerConnectable(host, port, tries=3, command='GET', path='/',
    133                             expected_read='', timeout=2):
    134   """Checks whether the specified http server is ready to serve request or not.
    135 
    136   Args:
    137     host: Host name of the HTTP server.
    138     port: Port number of the HTTP server.
    139     tries: How many times we want to test the connection. The default value is
    140            3.
    141     command: The http command we use to connect to HTTP server. The default
    142              command is 'GET'.
    143     path: The path we use when connecting to HTTP server. The default path is
    144           '/'.
    145     expected_read: The content we expect to read from the response. The default
    146                    value is ''.
    147     timeout: Timeout (in seconds) for each http connection. The default is 2s.
    148 
    149   Returns:
    150     Tuple of (connect status, client error). connect status is a boolean value
    151     to indicate whether the server is connectable. client_error is the error
    152     message the server returns when connect status is false.
    153   """
    154   assert tries >= 1
    155   for i in xrange(0, tries):
    156     client_error = None
    157     try:
    158       with contextlib.closing(httplib.HTTPConnection(
    159           host, port, timeout=timeout)) as http:
    160         # Output some debug information when we have tried more than 2 times.
    161         http.set_debuglevel(i >= 2)
    162         http.request(command, path)
    163         r = http.getresponse()
    164         content = r.read()
    165         if r.status == 200 and r.reason == 'OK' and content == expected_read:
    166           return (True, '')
    167         client_error = ('Bad response: %s %s version %s\n  ' %
    168                         (r.status, r.reason, r.version) +
    169                         '\n  '.join([': '.join(h) for h in r.getheaders()]))
    170     except (httplib.HTTPException, socket.error) as e:
    171       # Probably too quick connecting: try again.
    172       exception_error_msgs = traceback.format_exception_only(type(e), e)
    173       if exception_error_msgs:
    174         client_error = ''.join(exception_error_msgs)
    175   # Only returns last client_error.
    176   return (False, client_error or 'Timeout')
    177