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