1#!/usr/bin/env python3
2#
3#   Copyright 2018 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import logging
18import os
19import subprocess
20import socket
21import threading
22
23from acts import context
24from acts import utils
25from acts.controllers.adb_lib.error import AdbCommandError
26from acts.controllers.android_device import AndroidDevice
27from acts.controllers.fuchsia_lib.ssh import SSHProvider
28from acts.controllers.iperf_server import _AndroidDeviceBridge
29from acts.controllers.utils_lib.ssh import connection
30from acts.controllers.utils_lib.ssh import settings
31from acts.libs.proc import job
32from paramiko.buffered_pipe import PipeTimeout
33from paramiko.ssh_exception import SSHException
34
35MOBLY_CONTROLLER_CONFIG_NAME = 'IPerfClient'
36ACTS_CONTROLLER_REFERENCE_NAME = 'iperf_clients'
37
38
39class IPerfError(Exception):
40    """Raised on execution errors of iPerf."""
41
42
43def create(configs):
44    """Factory method for iperf clients.
45
46    The function creates iperf clients based on at least one config.
47    If configs contain ssh settings or and AndroidDevice, remote iperf clients
48    will be started on those devices, otherwise, a the client will run on the
49    local machine.
50
51    Args:
52        configs: config parameters for the iperf server
53    """
54    results = []
55    for c in configs:
56        if type(c) is dict and 'AndroidDevice' in c:
57            results.append(
58                IPerfClientOverAdb(c['AndroidDevice'],
59                                   test_interface=c.get('test_interface')))
60        elif type(c) is dict and 'ssh_config' in c:
61            results.append(
62                IPerfClientOverSsh(c['ssh_config'],
63                                   test_interface=c.get('test_interface')))
64        else:
65            results.append(IPerfClient())
66    return results
67
68
69def get_info(iperf_clients):
70    """Placeholder for info about iperf clients
71
72    Returns:
73        None
74    """
75    return None
76
77
78def destroy(_):
79    # No cleanup needed.
80    pass
81
82
83class IPerfClientBase(object):
84    """The Base class for all IPerfClients.
85
86    This base class is responsible for synchronizing the logging to prevent
87    multiple IPerfClients from writing results to the same file, as well
88    as providing the interface for IPerfClient objects.
89    """
90    # Keeps track of the number of IPerfClient logs to prevent file name
91    # collisions.
92    __log_file_counter = 0
93
94    __log_file_lock = threading.Lock()
95
96    @staticmethod
97    def _get_full_file_path(tag=''):
98        """Returns the full file path for the IPerfClient log file.
99
100        Note: If the directory for the file path does not exist, it will be
101        created.
102
103        Args:
104            tag: The tag passed in to the server run.
105        """
106        current_context = context.get_current_context()
107        full_out_dir = os.path.join(current_context.get_full_output_path(),
108                                    'iperf_client_files')
109
110        with IPerfClientBase.__log_file_lock:
111            os.makedirs(full_out_dir, exist_ok=True)
112            tags = ['IPerfClient', tag, IPerfClientBase.__log_file_counter]
113            out_file_name = '%s.log' % (','.join(
114                [str(x) for x in tags if x != '' and x is not None]))
115            IPerfClientBase.__log_file_counter += 1
116
117        return os.path.join(full_out_dir, out_file_name)
118
119    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
120        """Starts iperf client, and waits for completion.
121
122        Args:
123            ip: iperf server ip address.
124            iperf_args: A string representing arguments to start iperf
125                client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
126            tag: A string to further identify iperf results file
127            timeout: the maximum amount of time the iperf client can run.
128            iperf_binary: Location of iperf3 binary. If none, it is assumed the
129                the binary is in the path.
130
131        Returns:
132            full_out_path: iperf result path.
133        """
134        raise NotImplementedError('start() must be implemented.')
135
136
137class IPerfClient(IPerfClientBase):
138    """Class that handles iperf3 client operations."""
139
140    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
141        """Starts iperf client, and waits for completion.
142
143        Args:
144            ip: iperf server ip address.
145            iperf_args: A string representing arguments to start iperf
146            client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
147            tag: tag to further identify iperf results file
148            timeout: unused.
149            iperf_binary: Location of iperf3 binary. If none, it is assumed the
150                the binary is in the path.
151
152        Returns:
153            full_out_path: iperf result path.
154        """
155        if not iperf_binary:
156            logging.debug('No iperf3 binary specified.  '
157                          'Assuming iperf3 is in the path.')
158            iperf_binary = 'iperf3'
159        else:
160            logging.debug('Using iperf3 binary located at %s' % iperf_binary)
161        iperf_cmd = [str(iperf_binary), '-c', ip] + iperf_args.split(' ')
162        full_out_path = self._get_full_file_path(tag)
163
164        with open(full_out_path, 'w') as out_file:
165            subprocess.call(iperf_cmd, stdout=out_file)
166
167        return full_out_path
168
169
170class IPerfClientOverSsh(IPerfClientBase):
171    """Class that handles iperf3 client operations on remote machines."""
172
173    def __init__(self,
174                 ssh_config: str,
175                 test_interface: str = None,
176                 ssh_provider: SSHProvider = None):
177        self._ssh_provider = ssh_provider
178        if not self._ssh_provider:
179            self._ssh_settings = settings.from_config(ssh_config)
180            if not (utils.is_valid_ipv4_address(self._ssh_settings.hostname) or
181                    utils.is_valid_ipv6_address(self._ssh_settings.hostname)):
182                mdns_ip = utils.get_fuchsia_mdns_ipv6_address(
183                    self._ssh_settings.hostname)
184                if mdns_ip:
185                    self._ssh_settings.hostname = mdns_ip
186        self._ssh_session = None
187        self.start_ssh()
188
189        self.test_interface = test_interface
190
191    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
192        """Starts iperf client, and waits for completion.
193
194        Args:
195            ip: iperf server ip address.
196            iperf_args: A string representing arguments to start iperf
197            client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
198            tag: tag to further identify iperf results file
199            timeout: the maximum amount of time to allow the iperf client to run
200            iperf_binary: Location of iperf3 binary. If none, it is assumed the
201                the binary is in the path.
202
203        Returns:
204            full_out_path: iperf result path.
205        """
206        if not iperf_binary:
207            logging.debug('No iperf3 binary specified.  '
208                          'Assuming iperf3 is in the path.')
209            iperf_binary = 'iperf3'
210        else:
211            logging.debug('Using iperf3 binary located at %s' % iperf_binary)
212        iperf_cmd = '{} -c {} {}'.format(iperf_binary, ip, iperf_args)
213        full_out_path = self._get_full_file_path(tag)
214
215        try:
216            self.start_ssh()
217            if self._ssh_provider:
218                iperf_process = self._ssh_provider.run(iperf_cmd,
219                                                       timeout_sec=timeout)
220            else:
221                iperf_process = self._ssh_session.run(iperf_cmd,
222                                                      timeout=timeout)
223            iperf_output = iperf_process.stdout
224            with open(full_out_path, 'w') as out_file:
225                out_file.write(iperf_output)
226        except PipeTimeout:
227            raise TimeoutError('Paramiko PipeTimeout. Timed out waiting for '
228                               'iperf client to finish.')
229        except socket.timeout:
230            raise TimeoutError('Socket timeout. Timed out waiting for iperf '
231                               'client to finish.')
232        except SSHException as err:
233            raise ConnectionError('SSH connection failed: {}'.format(err))
234        except Exception as err:
235            logging.exception('iperf run failed: {}'.format(err))
236
237        return full_out_path
238
239    def start_ssh(self):
240        """Starts an ssh session to the iperf client."""
241        if self._ssh_provider:
242            # SSH sessions are created by the provider.
243            return
244        if not self._ssh_session:
245            self._ssh_session = connection.SshConnection(self._ssh_settings)
246
247    def close_ssh(self):
248        """Closes the ssh session to the iperf client, if one exists, preventing
249        connection reset errors when rebooting client device.
250        """
251        if self._ssh_session:
252            self._ssh_session.close()
253            self._ssh_session = None
254
255
256class IPerfClientOverAdb(IPerfClientBase):
257    """Class that handles iperf3 operations over ADB devices."""
258
259    def __init__(self, android_device_or_serial, test_interface=None):
260        """Creates a new IPerfClientOverAdb object.
261
262        Args:
263            android_device_or_serial: Either an AndroidDevice object, or the
264                serial that corresponds to the AndroidDevice. Note that the
265                serial must be present in an AndroidDevice entry in the ACTS
266                config.
267            test_interface: The network interface that will be used to send
268                traffic to the iperf server.
269        """
270        self._android_device_or_serial = android_device_or_serial
271        self.test_interface = test_interface
272
273    @property
274    def _android_device(self):
275        if isinstance(self._android_device_or_serial, AndroidDevice):
276            return self._android_device_or_serial
277        else:
278            return _AndroidDeviceBridge.android_devices()[
279                self._android_device_or_serial]
280
281    def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
282        """Starts iperf client, and waits for completion.
283
284        Args:
285            ip: iperf server ip address.
286            iperf_args: A string representing arguments to start iperf
287            client. Eg: iperf_args = "-t 10 -p 5001 -w 512k/-u -b 200M -J".
288            tag: tag to further identify iperf results file
289            timeout: the maximum amount of time to allow the iperf client to run
290            iperf_binary: Location of iperf3 binary. If none, it is assumed the
291                the binary is in the path.
292
293        Returns:
294            The iperf result file path.
295        """
296        clean_out = ''
297        try:
298            if not iperf_binary:
299                logging.debug('No iperf3 binary specified.  '
300                              'Assuming iperf3 is in the path.')
301                iperf_binary = 'iperf3'
302            else:
303                logging.debug('Using iperf3 binary located at %s' %
304                              iperf_binary)
305            iperf_cmd = '{} -c {} {}'.format(iperf_binary, ip, iperf_args)
306            out = self._android_device.adb.shell(str(iperf_cmd),
307                                                 timeout=timeout)
308            clean_out = out.split('\n')
309            if 'error' in clean_out[0].lower():
310                raise IPerfError(clean_out)
311        except (job.TimeoutError, AdbCommandError):
312            logging.warning('TimeoutError: Iperf measurement failed.')
313
314        full_out_path = self._get_full_file_path(tag)
315        with open(full_out_path, 'w') as out_file:
316            out_file.write('\n'.join(clean_out))
317
318        return full_out_path
319