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