1# Copyright 2014 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
5import contextlib
6import logging
7import socket
8
9from telemetry.internal import forwarders
10
11import py_utils
12
13
14class Error(Exception):
15  """Base class for exceptions in this module."""
16  pass
17
18
19class PortsMismatchError(Error):
20  """Raised when local and remote ports are not equal."""
21  pass
22
23
24class ConnectionError(Error):
25  """Raised when unable to connect to local TCP ports."""
26  pass
27
28
29class DoNothingForwarderFactory(forwarders.ForwarderFactory):
30
31  def Create(self, port_pair):
32    return DoNothingForwarder(port_pair)
33
34
35class DoNothingForwarder(forwarders.Forwarder):
36  """Check that no forwarding is needed for the given port pairs.
37
38  The local and remote ports must be equal. Otherwise, the "do nothing"
39  forwarder does not make sense. (Raises PortsMismatchError.)
40
41  Also, check that all TCP ports support connections.  (Raises ConnectionError.)
42  """
43
44  def __init__(self, port_pair):
45    super(DoNothingForwarder, self).__init__(port_pair)
46    self._CheckPortPair()
47
48  def _CheckPortPair(self):
49    if self._port_pair.local_port != self._port_pair.remote_port:
50      raise PortsMismatchError('Local port forwarding is not supported')
51    try:
52      self._WaitForConnectionEstablished(
53          (self.host_ip, self._port_pair.local_port), timeout=10)
54      logging.debug(
55          'Connection test succeeded for %s:%d',
56          self.host_ip, self._port_pair.local_port)
57    except py_utils.TimeoutException:
58      raise ConnectionError(
59          'Unable to connect to address: %s:%d',
60          self.host_ip, self._port_pair.local_port)
61
62  def _WaitForConnectionEstablished(self, address, timeout):
63    def CanConnect():
64      with contextlib.closing(socket.socket()) as s:
65        return s.connect_ex(address) == 0
66    py_utils.WaitFor(CanConnect, timeout)
67