1# Lint as: python2, python3
2# Copyright (c) 2019 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import logging
7import re
8from six.moves import urllib
9import socket
10import time
11
12import common
13
14from autotest_lib.client.bin import utils
15from autotest_lib.client.common_lib import error
16
17
18def CheckInterfaceForDestination(host, expected_interface,
19                                 family=socket.AF_UNSPEC):
20    """
21    Checks that routes for host go through a given interface.
22
23    The concern here is that our network setup may have gone wrong
24    and our test connections may go over some other network than
25    the one we're trying to test.  So we take all the IP addresses
26    for the supplied host and make sure they go through the given
27    network interface.
28
29    @param host: Destination host
30    @param expected_interface: Expected interface name
31    @raises: error.TestFail if the routes for the given host go through
32            a different interface than the expected one.
33
34    """
35    def _MatchesRoute(address, expected_interface):
36        """
37        Returns whether or not |expected_interface| is used to reach |address|.
38
39        @param address: string containing an IP (v4 or v6) address.
40        @param expected_interface: string containing an interface name.
41
42        """
43        output = utils.run('ip route get %s' % address).stdout
44
45        if re.search(r'unreachable', output):
46            return False
47
48        match = re.search(r'\sdev\s(\S+)', output)
49        if match is None:
50            return False
51        interface = match.group(1)
52
53        logging.info('interface for %s: %s', address, interface)
54        if interface != expected_interface:
55            raise error.TestFail('Target server %s uses interface %s'
56                                 '(%s expected).' %
57                                 (address, interface, expected_interface))
58        return True
59
60    # addrinfo records: (family, type, proto, canonname, (addr, port))
61    server_addresses = [record[4][0]
62                        for record in socket.getaddrinfo(host, 80, family)]
63    for address in server_addresses:
64        # Routes may not always be up by this point. Note that routes for v4 or
65        # v6 may come up before the other, so we simply do this poll for all
66        # addresses.
67        utils.poll_for_condition(
68            condition=lambda: _MatchesRoute(address, expected_interface),
69            exception=error.TestFail('No route to %s' % address),
70            timeout=1)
71
72FETCH_URL_PATTERN_FOR_TEST = \
73    'http://testing-chargen.appspot.com/download?size=%d'
74
75def FetchUrl(url_pattern, bytes_to_fetch=10, fetch_timeout=10):
76    """
77    Fetches a specified number of bytes from a URL.
78
79    @param url_pattern: URL pattern for fetching a specified number of bytes.
80            %d in the pattern is to be filled in with the number of bytes to
81            fetch.
82    @param bytes_to_fetch: Number of bytes to fetch.
83    @param fetch_timeout: Number of seconds to wait for the fetch to complete
84            before it times out.
85    @return: The time in seconds spent for fetching the specified number of
86            bytes.
87    @raises: error.TestError if one of the following happens:
88            - The fetch takes no time.
89            - The number of bytes fetched differs from the specified
90              number.
91
92    """
93    # Limit the amount of bytes to read at a time.
94    _MAX_FETCH_READ_BYTES = 1024 * 1024
95
96    url = url_pattern % bytes_to_fetch
97    logging.info('FetchUrl %s', url)
98    start_time = time.time()
99    result = urllib.request.urlopen(url, timeout=fetch_timeout)
100    bytes_fetched = 0
101    while bytes_fetched < bytes_to_fetch:
102        bytes_left = bytes_to_fetch - bytes_fetched
103        bytes_to_read = min(bytes_left, _MAX_FETCH_READ_BYTES)
104        bytes_read = len(result.read(bytes_to_read))
105        bytes_fetched += bytes_read
106        if bytes_read != bytes_to_read:
107            raise error.TestError('FetchUrl tried to read %d bytes, but got '
108                                  '%d bytes instead.' %
109                                  (bytes_to_read, bytes_read))
110        fetch_time = time.time() - start_time
111        if fetch_time > fetch_timeout:
112            raise error.TestError('FetchUrl exceeded timeout.')
113
114    return fetch_time
115