1# SPDX-License-Identifier: GPL-2.0
2# Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
3
4"""
5Logic to spawn a sub-process and interact with its stdio.
6"""
7
8import io
9import os
10import re
11import pty
12import pytest
13import signal
14import select
15import sys
16import termios
17import time
18import traceback
19
20# Character to send (twice) to exit the terminal
21EXIT_CHAR = 0x1d    # FS (Ctrl + ])
22
23class Timeout(Exception):
24    """An exception sub-class that indicates that a timeout occurred."""
25
26class BootFail(Exception):
27    """An exception sub-class that indicates that a boot failure occurred.
28
29    This is used when a bad pattern is seen when waiting for the boot prompt.
30    It is regarded as fatal, to avoid trying to boot the again and again to no
31    avail.
32    """
33
34class Unexpected(Exception):
35    """An exception sub-class that indicates that unexpected test was seen."""
36
37
38def handle_exception(ubconfig, console, log, err, name, fatal, output=''):
39    """Handle an exception from the console
40
41    Exceptions can occur when there is unexpected output or due to the board
42    crashing or hanging. Some exceptions are likely fatal, where retrying will
43    just chew up time to no available. In those cases it is best to cause
44    further tests be skipped.
45
46    Args:
47        ubconfig (ArbitraryAttributeContainer): ubconfig object
48        log (Logfile): Place to log errors
49        console (ConsoleBase): Console to clean up, if fatal
50        err (Exception): Exception which was thrown
51        name (str): Name of problem, to log
52        fatal (bool): True to abort all tests
53        output (str): Extra output to report on boot failure. This can show the
54           target's console output as it tried to boot
55    """
56    msg = f'{name}: '
57    if fatal:
58        msg += 'Marking connection bad - no other tests will run'
59    else:
60        msg += 'Assuming that lab is healthy'
61    print(msg)
62    log.error(msg)
63    log.error(f'Error: {err}')
64
65    if output:
66        msg += f'; output {output}'
67
68    if fatal:
69        ubconfig.connection_ok = False
70        console.cleanup_spawn()
71        pytest.exit(msg)
72
73
74class Spawn:
75    """Represents the stdio of a freshly created sub-process. Commands may be
76    sent to the process, and responses waited for.
77
78    Members:
79        output: accumulated output from expect()
80    """
81
82    def __init__(self, args, cwd=None, decode_signal=False):
83        """Spawn (fork/exec) the sub-process.
84
85        Args:
86            args: array of processs arguments. argv[0] is the command to
87              execute.
88            cwd: the directory to run the process in, or None for no change.
89            decode_signal (bool): True to indicate the exception number when
90                something goes wrong
91
92        Returns:
93            Nothing.
94        """
95        self.decode_signal = decode_signal
96        self.waited = False
97        self.exit_code = 0
98        self.exit_info = ''
99        self.buf = ''
100        self.output = ''
101        self.logfile_read = None
102        self.before = ''
103        self.after = ''
104        self.timeout = None
105        # http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences
106        self.re_vt100 = re.compile(r'(\x1b\[|\x9b)[^@-_]*[@-_]|\x1b[@-_]', re.I)
107
108        (self.pid, self.fd) = pty.fork()
109        if self.pid == 0:
110            try:
111                # For some reason, SIGHUP is set to SIG_IGN at this point when
112                # run under "go" (www.go.cd). Perhaps this happens under any
113                # background (non-interactive) system?
114                signal.signal(signal.SIGHUP, signal.SIG_DFL)
115                if cwd:
116                    os.chdir(cwd)
117                os.execvp(args[0], args)
118            except:
119                print('CHILD EXECEPTION:')
120                traceback.print_exc()
121            finally:
122                os._exit(255)
123
124        old = None
125        try:
126            isatty = False
127            try:
128                isatty = os.isatty(sys.stdout.fileno())
129
130            # with --capture=tee-sys we cannot call fileno()
131            except io.UnsupportedOperation as exc:
132                pass
133            if isatty:
134                new = termios.tcgetattr(self.fd)
135                old = new
136                new[3] = new[3] & ~(termios.ICANON | termios.ISIG)
137                new[3] = new[3] & ~termios.ECHO
138                new[6][termios.VMIN] = 0
139                new[6][termios.VTIME] = 0
140                termios.tcsetattr(self.fd, termios.TCSANOW, new)
141
142            self.poll = select.poll()
143            self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR |
144                               select.POLLHUP | select.POLLNVAL)
145        except:
146            if old:
147                termios.tcsetattr(self.fd, termios.TCSANOW, old)
148            self.close()
149            raise
150
151    def kill(self, sig):
152        """Send unix signal "sig" to the child process.
153
154        Args:
155            sig: The signal number to send.
156
157        Returns:
158            Nothing.
159        """
160
161        os.kill(self.pid, sig)
162
163    def checkalive(self):
164        """Determine whether the child process is still running.
165
166        Returns:
167            tuple:
168                True if process is alive, else False
169                0 if process is alive, else exit code of process
170                string describing what happened ('' or 'status/signal n')
171        """
172
173        if self.waited:
174            return False, self.exit_code, self.exit_info
175
176        w = os.waitpid(self.pid, os.WNOHANG)
177        if w[0] == 0:
178            return True, 0, 'running'
179        status = w[1]
180
181        if os.WIFEXITED(status):
182            self.exit_code = os.WEXITSTATUS(status)
183            self.exit_info = 'status %d' % self.exit_code
184        elif os.WIFSIGNALED(status):
185            signum = os.WTERMSIG(status)
186            self.exit_code = -signum
187            self.exit_info = 'signal %d (%s)' % (signum, signal.Signals(signum).name)
188        self.waited = True
189        return False, self.exit_code, self.exit_info
190
191    def isalive(self):
192        """Determine whether the child process is still running.
193
194        Args:
195            None.
196
197        Returns:
198            Boolean indicating whether process is alive.
199        """
200        return self.checkalive()[0]
201
202    def send(self, data):
203        """Send data to the sub-process's stdin.
204
205        Args:
206            data: The data to send to the process.
207
208        Returns:
209            Nothing.
210        """
211
212        os.write(self.fd, data.encode(errors='replace'))
213
214    def receive(self, num_bytes):
215        """Receive data from the sub-process's stdin.
216
217        Args:
218            num_bytes (int): Maximum number of bytes to read
219
220        Returns:
221            str: The data received
222
223        Raises:
224            ValueError if U-Boot died
225        """
226        try:
227            c = os.read(self.fd, num_bytes).decode(errors='replace')
228        except OSError as err:
229            # With sandbox, try to detect when U-Boot exits when it
230            # shouldn't and explain why. This is much more friendly than
231            # just dying with an I/O error
232            if self.decode_signal and err.errno == 5:  # I/O error
233                alive, _, info = self.checkalive()
234                if alive:
235                    raise err
236                raise ValueError('U-Boot exited with %s' % info)
237            raise
238        return c
239
240    def expect(self, patterns):
241        """Wait for the sub-process to emit specific data.
242
243        This function waits for the process to emit one pattern from the
244        supplied list of patterns, or for a timeout to occur.
245
246        Args:
247            patterns: A list of strings or regex objects that we expect to
248                see in the sub-process' stdout.
249
250        Returns:
251            The index within the patterns array of the pattern the process
252            emitted.
253
254        Notable exceptions:
255            Timeout, if the process did not emit any of the patterns within
256            the expected time.
257        """
258
259        for pi in range(len(patterns)):
260            if type(patterns[pi]) == type(''):
261                patterns[pi] = re.compile(patterns[pi])
262
263        tstart_s = time.time()
264        try:
265            while True:
266                earliest_m = None
267                earliest_pi = None
268                for pi in range(len(patterns)):
269                    pattern = patterns[pi]
270                    m = pattern.search(self.buf)
271                    if not m:
272                        continue
273                    if earliest_m and m.start() >= earliest_m.start():
274                        continue
275                    earliest_m = m
276                    earliest_pi = pi
277                if earliest_m:
278                    pos = earliest_m.start()
279                    posafter = earliest_m.end()
280                    self.before = self.buf[:pos]
281                    self.after = self.buf[pos:posafter]
282                    self.output += self.buf[:posafter]
283                    self.buf = self.buf[posafter:]
284                    return earliest_pi
285                tnow_s = time.time()
286                if self.timeout:
287                    tdelta_ms = (tnow_s - tstart_s) * 1000
288                    poll_maxwait = self.timeout - tdelta_ms
289                    if tdelta_ms > self.timeout:
290                        raise Timeout()
291                else:
292                    poll_maxwait = None
293                events = self.poll.poll(poll_maxwait)
294                if not events:
295                    raise Timeout()
296                c = self.receive(1024)
297                if self.logfile_read:
298                    self.logfile_read.write(c)
299                self.buf += c
300                # count=0 is supposed to be the default, which indicates
301                # unlimited substitutions, but in practice the version of
302                # Python in Ubuntu 14.04 appears to default to count=2!
303                self.buf = self.re_vt100.sub('', self.buf, count=1000000)
304        finally:
305            if self.logfile_read:
306                self.logfile_read.flush()
307
308    def close(self):
309        """Close the stdio connection to the sub-process.
310
311        This also waits a reasonable time for the sub-process to stop running.
312
313        Args:
314            None.
315
316        Returns:
317            str: Type of closure completed
318        """
319        # For Labgrid-sjg, ask it is exit gracefully, so it can transition the
320        # board to the final state (like 'off') before exiting.
321        if os.environ.get('USE_LABGRID_SJG'):
322            self.send(chr(EXIT_CHAR) * 2)
323
324            # Wait about 10 seconds for Labgrid to close and power off the board
325            for _ in range(100):
326                if not self.isalive():
327                    return 'normal'
328                time.sleep(0.1)
329
330        # That didn't work, so try closing the PTY
331        os.close(self.fd)
332        for _ in range(100):
333            if not self.isalive():
334                return 'break'
335            time.sleep(0.1)
336
337        return 'timeout'
338
339    def get_expect_output(self):
340        """Return the output read by expect()
341
342        Returns:
343            The output processed by expect(), as a string.
344        """
345        return self.output
346