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