1# SPDX-License-Identifier: GPL-2.0
2# Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.
3
4"""
5Utility code shared across multiple tests.
6"""
7
8import hashlib
9import inspect
10import os
11import os.path
12import pathlib
13import signal
14import sys
15import time
16import re
17import pytest
18
19def md5sum_data(data):
20    """Calculate the MD5 hash of some data.
21
22    Args:
23        data: The data to hash.
24
25    Returns:
26        The hash of the data, as a binary string.
27    """
28
29    h = hashlib.md5()
30    h.update(data)
31    return h.digest()
32
33def md5sum_file(fn, max_length=None):
34    """Calculate the MD5 hash of the contents of a file.
35
36    Args:
37        fn: The filename of the file to hash.
38        max_length: The number of bytes to hash. If the file has more
39            bytes than this, they will be ignored. If None or omitted, the
40            entire file will be hashed.
41
42    Returns:
43        The hash of the file content, as a binary string.
44    """
45
46    with open(fn, 'rb') as fh:
47        if max_length:
48            params = [max_length]
49        else:
50            params = []
51        data = fh.read(*params)
52    return md5sum_data(data)
53
54class PersistentRandomFile:
55    """Generate and store information about a persistent file containing
56    random data."""
57
58    def __init__(self, u_boot_console, fn, size):
59        """Create or process the persistent file.
60
61        If the file does not exist, it is generated.
62
63        If the file does exist, its content is hashed for later comparison.
64
65        These files are always located in the "persistent data directory" of
66        the current test run.
67
68        Args:
69            u_boot_console: A console connection to U-Boot.
70            fn: The filename (without path) to create.
71            size: The desired size of the file in bytes.
72
73        Returns:
74            Nothing.
75        """
76
77        self.fn = fn
78
79        self.abs_fn = u_boot_console.config.persistent_data_dir + '/' + fn
80
81        if os.path.exists(self.abs_fn):
82            u_boot_console.log.action('Persistent data file ' + self.abs_fn +
83                ' already exists')
84            self.content_hash = md5sum_file(self.abs_fn)
85        else:
86            u_boot_console.log.action('Generating ' + self.abs_fn +
87                ' (random, persistent, %d bytes)' % size)
88            data = os.urandom(size)
89            with open(self.abs_fn, 'wb') as fh:
90                fh.write(data)
91            self.content_hash = md5sum_data(data)
92
93def attempt_to_open_file(fn):
94    """Attempt to open a file, without throwing exceptions.
95
96    Any errors (exceptions) that occur during the attempt to open the file
97    are ignored. This is useful in order to test whether a file (in
98    particular, a device node) exists and can be successfully opened, in order
99    to poll for e.g. USB enumeration completion.
100
101    Args:
102        fn: The filename to attempt to open.
103
104    Returns:
105        An open file handle to the file, or None if the file could not be
106            opened.
107    """
108
109    try:
110        return open(fn, 'rb')
111    except:
112        return None
113
114def wait_until_open_succeeds(fn):
115    """Poll until a file can be opened, or a timeout occurs.
116
117    Continually attempt to open a file, and return when this succeeds, or
118    raise an exception after a timeout.
119
120    Args:
121        fn: The filename to attempt to open.
122
123    Returns:
124        An open file handle to the file.
125    """
126
127    for i in range(100):
128        fh = attempt_to_open_file(fn)
129        if fh:
130            return fh
131        time.sleep(0.1)
132    raise Exception('File could not be opened')
133
134def wait_until_file_open_fails(fn, ignore_errors):
135    """Poll until a file cannot be opened, or a timeout occurs.
136
137    Continually attempt to open a file, and return when this fails, or
138    raise an exception after a timeout.
139
140    Args:
141        fn: The filename to attempt to open.
142        ignore_errors: Indicate whether to ignore timeout errors. If True, the
143            function will simply return if a timeout occurs, otherwise an
144            exception will be raised.
145
146    Returns:
147        Nothing.
148    """
149
150    for _ in range(100):
151        fh = attempt_to_open_file(fn)
152        if not fh:
153            return
154        fh.close()
155        time.sleep(0.1)
156    if ignore_errors:
157        return
158    raise Exception('File can still be opened')
159
160def run_and_log(u_boot_console, cmd, ignore_errors=False, stdin=None, env=None):
161    """Run a command and log its output.
162
163    Args:
164        u_boot_console: A console connection to U-Boot.
165        cmd: The command to run, as an array of argv[], or a string.
166            If a string, note that it is split up so that quoted spaces
167            will not be preserved. E.g. "fred and" becomes ['"fred', 'and"']
168        ignore_errors: Indicate whether to ignore errors. If True, the function
169            will simply return if the command cannot be executed or exits with
170            an error code, otherwise an exception will be raised if such
171            problems occur.
172        stdin: Input string to pass to the command as stdin (or None)
173            env: Environment to use, or None to use the current one
174
175    Returns:
176        The output as a string.
177    """
178    if isinstance(cmd, str):
179        cmd = cmd.split()
180    runner = u_boot_console.log.get_runner(cmd[0], sys.stdout)
181    output = runner.run(cmd, ignore_errors=ignore_errors, stdin=stdin, env=env)
182    runner.close()
183    return output
184
185def run_and_log_expect_exception(u_boot_console, cmd, retcode, msg):
186    """Run a command that is expected to fail.
187
188    This runs a command and checks that it fails with the expected return code
189    and exception method. If not, an exception is raised.
190
191    Args:
192        u_boot_console: A console connection to U-Boot.
193        cmd: The command to run, as an array of argv[].
194        retcode: Expected non-zero return code from the command.
195        msg: String that should be contained within the command's output.
196    """
197    try:
198        runner = u_boot_console.log.get_runner(cmd[0], sys.stdout)
199        runner.run(cmd)
200    except Exception:
201        assert retcode == runner.exit_status
202        assert msg in runner.output
203    else:
204        raise Exception("Expected an exception with retcode %d message '%s',"
205                        "but it was not raised" % (retcode, msg))
206    finally:
207        runner.close()
208
209ram_base = None
210def find_ram_base(u_boot_console):
211    """Find the running U-Boot's RAM location.
212
213    Probe the running U-Boot to determine the address of the first bank
214    of RAM. This is useful for tests that test reading/writing RAM, or
215    load/save files that aren't associated with some standard address
216    typically represented in an environment variable such as
217    ${kernel_addr_r}. The value is cached so that it only needs to be
218    actively read once.
219
220    Args:
221        u_boot_console: A console connection to U-Boot.
222
223    Returns:
224        The address of U-Boot's first RAM bank, as an integer.
225    """
226
227    global ram_base
228    if u_boot_console.config.buildconfig.get('config_cmd_bdi', 'n') != 'y':
229        pytest.skip('bdinfo command not supported')
230    if ram_base == -1:
231        pytest.skip('Previously failed to find RAM bank start')
232    if ram_base is not None:
233        return ram_base
234
235    with u_boot_console.log.section('find_ram_base'):
236        response = u_boot_console.run_command('bdinfo')
237        for l in response.split('\n'):
238            if '-> start' in l or 'memstart    =' in l:
239                ram_base = int(l.split('=')[1].strip(), 16)
240                break
241        if ram_base is None:
242            ram_base = -1
243            raise Exception('Failed to find RAM bank start in `bdinfo`')
244
245    # We don't want ram_base to be zero as some functions test if the given
246    # address is NULL (0). Besides, on some RISC-V targets the low memory
247    # is protected that prevents S-mode U-Boot from access.
248    # Let's add 2MiB then (size of an ARM LPAE/v8 section).
249
250    ram_base += 1024 * 1024 * 2
251
252    return ram_base
253
254class PersistentFileHelperCtxMgr(object):
255    """A context manager for Python's "with" statement, which ensures that any
256    generated file is deleted (and hence regenerated) if its mtime is older
257    than the mtime of the Python module which generated it, and gets an mtime
258    newer than the mtime of the Python module which generated after it is
259    generated. Objects of this type should be created by factory function
260    persistent_file_helper rather than directly."""
261
262    def __init__(self, log, filename):
263        """Initialize a new object.
264
265        Args:
266            log: The Logfile object to log to.
267            filename: The filename of the generated file.
268
269        Returns:
270            Nothing.
271        """
272
273        self.log = log
274        self.filename = filename
275
276    def __enter__(self):
277        frame = inspect.stack()[1]
278        module = inspect.getmodule(frame[0])
279        self.module_filename = module.__file__
280        self.module_timestamp = os.path.getmtime(self.module_filename)
281
282        if os.path.exists(self.filename):
283            filename_timestamp = os.path.getmtime(self.filename)
284            if filename_timestamp < self.module_timestamp:
285                self.log.action('Removing stale generated file ' +
286                    self.filename)
287                pathlib.Path(self.filename).unlink()
288
289    def __exit__(self, extype, value, traceback):
290        if extype:
291            try:
292                pathlib.Path(self.filename).unlink()
293            except Exception:
294                pass
295            return
296        logged = False
297        for _ in range(20):
298            filename_timestamp = os.path.getmtime(self.filename)
299            if filename_timestamp > self.module_timestamp:
300                break
301            if not logged:
302                self.log.action(
303                    'Waiting for generated file timestamp to increase')
304                logged = True
305            os.utime(self.filename)
306            time.sleep(0.1)
307
308def persistent_file_helper(u_boot_log, filename):
309    """Manage the timestamps and regeneration of a persistent generated
310    file. This function creates a context manager for Python's "with"
311    statement
312
313    Usage:
314        with persistent_file_helper(u_boot_console.log, filename):
315            code to generate the file, if it's missing.
316
317    Args:
318        u_boot_log: u_boot_console.log.
319        filename: The filename of the generated file.
320
321    Returns:
322        A context manager object.
323    """
324
325    return PersistentFileHelperCtxMgr(u_boot_log, filename)
326
327def crc32(u_boot_console, address, count):
328    """Helper function used to compute the CRC32 value of a section of RAM.
329
330    Args:
331        u_boot_console: A U-Boot console connection.
332        address: Address where data starts.
333        count: Amount of data to use for calculation.
334
335    Returns:
336        CRC32 value
337    """
338
339    bcfg = u_boot_console.config.buildconfig
340    has_cmd_crc32 = bcfg.get('config_cmd_crc32', 'n') == 'y'
341    assert has_cmd_crc32, 'Cannot compute crc32 without CONFIG_CMD_CRC32.'
342    output = u_boot_console.run_command('crc32 %08x %x' % (address, count))
343
344    m = re.search('==> ([0-9a-fA-F]{8})$', output)
345    assert m, 'CRC32 operation failed.'
346
347    return m.group(1)
348
349def waitpid(pid, timeout=60, kill=False):
350    """Wait a process to terminate by its PID
351
352    This is an alternative to a os.waitpid(pid, 0) call that works on
353    processes that aren't children of the python process.
354
355    Args:
356        pid: PID of a running process.
357        timeout: Time in seconds to wait.
358        kill: Whether to forcibly kill the process after timeout.
359
360    Returns:
361        True, if the process ended on its own.
362        False, if the process was killed by this function.
363
364    Raises:
365        TimeoutError, if the process is still running after timeout.
366    """
367    try:
368        for _ in range(timeout):
369            os.kill(pid, 0)
370            time.sleep(1)
371
372        if kill:
373            os.kill(pid, signal.SIGKILL)
374            return False
375
376    except ProcessLookupError:
377        return True
378
379    raise TimeoutError(
380        "Process with PID {} did not terminate after {} seconds."
381        .format(pid, timeout)
382    )
383