michael@0: # Copyright (c) 2012 The Chromium Authors. All rights reserved.
michael@0: # Use of this source code is governed by a BSD-style license that can be
michael@0: # found in the LICENSE file.
michael@0: 
michael@0: """Functions that deals with local and device ports."""
michael@0: 
michael@0: import contextlib
michael@0: import fcntl
michael@0: import httplib
michael@0: import logging
michael@0: import os
michael@0: import re
michael@0: import socket
michael@0: import traceback
michael@0: 
michael@0: import cmd_helper
michael@0: import constants
michael@0: 
michael@0: 
michael@0: #The following two methods are used to allocate the port source for various
michael@0: # types of test servers. Because some net relates tests can be run on shards
michael@0: # at same time, it's important to have a mechanism to allocate the port process
michael@0: # safe. In here, we implement the safe port allocation by leveraging flock.
michael@0: def ResetTestServerPortAllocation():
michael@0:   """Reset the port allocation to start from TEST_SERVER_PORT_FIRST.
michael@0: 
michael@0:   Returns:
michael@0:     Returns True if reset successes. Otherwise returns False.
michael@0:   """
michael@0:   try:
michael@0:     with open(constants.TEST_SERVER_PORT_FILE, 'w') as fp:
michael@0:       fp.write('%d' % constants.TEST_SERVER_PORT_FIRST)
michael@0:     if os.path.exists(constants.TEST_SERVER_PORT_LOCKFILE):
michael@0:       os.unlink(constants.TEST_SERVER_PORT_LOCKFILE)
michael@0:     return True
michael@0:   except Exception as e:
michael@0:     logging.error(e)
michael@0:   return False
michael@0: 
michael@0: 
michael@0: def AllocateTestServerPort():
michael@0:   """Allocate a port incrementally.
michael@0: 
michael@0:   Returns:
michael@0:     Returns a valid port which should be in between TEST_SERVER_PORT_FIRST and
michael@0:     TEST_SERVER_PORT_LAST. Returning 0 means no more valid port can be used.
michael@0:   """
michael@0:   port = 0
michael@0:   ports_tried = []
michael@0:   try:
michael@0:     fp_lock = open(constants.TEST_SERVER_PORT_LOCKFILE, 'w')
michael@0:     fcntl.flock(fp_lock, fcntl.LOCK_EX)
michael@0:     # Get current valid port and calculate next valid port.
michael@0:     assert os.path.exists(constants.TEST_SERVER_PORT_FILE)
michael@0:     with open(constants.TEST_SERVER_PORT_FILE, 'r+') as fp:
michael@0:       port = int(fp.read())
michael@0:       ports_tried.append(port)
michael@0:       while IsHostPortUsed(port):
michael@0:         port += 1
michael@0:         ports_tried.append(port)
michael@0:       if (port > constants.TEST_SERVER_PORT_LAST or
michael@0:           port < constants.TEST_SERVER_PORT_FIRST):
michael@0:         port = 0
michael@0:       else:
michael@0:         fp.seek(0, os.SEEK_SET)
michael@0:         fp.write('%d' % (port + 1))
michael@0:   except Exception as e:
michael@0:     logging.info(e)
michael@0:   finally:
michael@0:     if fp_lock:
michael@0:       fcntl.flock(fp_lock, fcntl.LOCK_UN)
michael@0:       fp_lock.close()
michael@0:   if port:
michael@0:     logging.info('Allocate port %d for test server.', port)
michael@0:   else:
michael@0:     logging.error('Could not allocate port for test server. '
michael@0:                   'List of ports tried: %s', str(ports_tried))
michael@0:   return port
michael@0: 
michael@0: 
michael@0: def IsHostPortUsed(host_port):
michael@0:   """Checks whether the specified host port is used or not.
michael@0: 
michael@0:   Uses -n -P to inhibit the conversion of host/port numbers to host/port names.
michael@0: 
michael@0:   Args:
michael@0:     host_port: Port on host we want to check.
michael@0: 
michael@0:   Returns:
michael@0:     True if the port on host is already used, otherwise returns False.
michael@0:   """
michael@0:   port_info = '(127\.0\.0\.1)|(localhost)\:%d' % host_port
michael@0:   # TODO(jnd): Find a better way to filter the port.
michael@0:   re_port = re.compile(port_info, re.MULTILINE)
michael@0:   if re_port.findall(cmd_helper.GetCmdOutput(['lsof', '-nPi:%d' % host_port])):
michael@0:     return True
michael@0:   return False
michael@0: 
michael@0: 
michael@0: def IsDevicePortUsed(adb, device_port, state=''):
michael@0:   """Checks whether the specified device port is used or not.
michael@0: 
michael@0:   Args:
michael@0:     adb: Instance of AndroidCommands for talking to the device.
michael@0:     device_port: Port on device we want to check.
michael@0:     state: String of the specified state. Default is empty string, which
michael@0:            means any state.
michael@0: 
michael@0:   Returns:
michael@0:     True if the port on device is already used, otherwise returns False.
michael@0:   """
michael@0:   base_url = '127.0.0.1:%d' % device_port
michael@0:   netstat_results = adb.RunShellCommand('netstat', log_result=False)
michael@0:   for single_connect in netstat_results:
michael@0:     # Column 3 is the local address which we want to check with.
michael@0:     connect_results = single_connect.split()
michael@0:     is_state_match = connect_results[5] == state if state else True
michael@0:     if connect_results[3] == base_url and is_state_match:
michael@0:       return True
michael@0:   return False
michael@0: 
michael@0: 
michael@0: def IsHttpServerConnectable(host, port, tries=3, command='GET', path='/',
michael@0:                             expected_read='', timeout=2):
michael@0:   """Checks whether the specified http server is ready to serve request or not.
michael@0: 
michael@0:   Args:
michael@0:     host: Host name of the HTTP server.
michael@0:     port: Port number of the HTTP server.
michael@0:     tries: How many times we want to test the connection. The default value is
michael@0:            3.
michael@0:     command: The http command we use to connect to HTTP server. The default
michael@0:              command is 'GET'.
michael@0:     path: The path we use when connecting to HTTP server. The default path is
michael@0:           '/'.
michael@0:     expected_read: The content we expect to read from the response. The default
michael@0:                    value is ''.
michael@0:     timeout: Timeout (in seconds) for each http connection. The default is 2s.
michael@0: 
michael@0:   Returns:
michael@0:     Tuple of (connect status, client error). connect status is a boolean value
michael@0:     to indicate whether the server is connectable. client_error is the error
michael@0:     message the server returns when connect status is false.
michael@0:   """
michael@0:   assert tries >= 1
michael@0:   for i in xrange(0, tries):
michael@0:     client_error = None
michael@0:     try:
michael@0:       with contextlib.closing(httplib.HTTPConnection(
michael@0:           host, port, timeout=timeout)) as http:
michael@0:         # Output some debug information when we have tried more than 2 times.
michael@0:         http.set_debuglevel(i >= 2)
michael@0:         http.request(command, path)
michael@0:         r = http.getresponse()
michael@0:         content = r.read()
michael@0:         if r.status == 200 and r.reason == 'OK' and content == expected_read:
michael@0:           return (True, '')
michael@0:         client_error = ('Bad response: %s %s version %s\n  ' %
michael@0:                         (r.status, r.reason, r.version) +
michael@0:                         '\n  '.join([': '.join(h) for h in r.getheaders()]))
michael@0:     except (httplib.HTTPException, socket.error) as e:
michael@0:       # Probably too quick connecting: try again.
michael@0:       exception_error_msgs = traceback.format_exception_only(type(e), e)
michael@0:       if exception_error_msgs:
michael@0:         client_error = ''.join(exception_error_msgs)
michael@0:   # Only returns last client_error.
michael@0:   return (False, client_error or 'Timeout')