test/py/u_boot_spawn.py

   1 # SPDX-License-Identifier: GPL-2.0
   2 # Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
   3
   4 """
   5 Logic to spawn a sub-process and interact with its stdio.
   6 """
   7
   8 import os
   9 import re
  10 import pty
  11 import signal
  12 import select
  13 import time
  14 import traceback
  15
  16 class Timeout(Exception):
  17     """An exception sub-class that indicates that a timeout occurred."""
  18
  19 class Spawn:
  20     """Represents the stdio of a freshly created sub-process. Commands may be
  21     sent to the process, and responses waited for.
  22
  23     Members:
  24         output: accumulated output from expect()
  25     """
  26
  27     def __init__(self, args, cwd=None):
  28         """Spawn (fork/exec) the sub-process.
  29
  30         Args:
  31             args: array of processs arguments. argv[0] is the command to
  32               execute.
  33             cwd: the directory to run the process in, or None for no change.
  34
  35         Returns:
  36             Nothing.
  37         """
  38
  39         self.waited = False
  40         self.exit_code = 0
  41         self.exit_info = ''
  42         self.buf = ''
  43         self.output = ''
  44         self.logfile_read = None
  45         self.before = ''
  46         self.after = ''
  47         self.timeout = None
  48         # http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences
  49         self.re_vt100 = re.compile(r'(\x1b\[|\x9b)[^@-_]*[@-_]|\x1b[@-_]', re.I)
  50
  51         (self.pid, self.fd) = pty.fork()
  52         if self.pid == 0:
  53             try:
  54                 # For some reason, SIGHUP is set to SIG_IGN at this point when
  55                 # run under "go" (www.go.cd). Perhaps this happens under any
  56                 # background (non-interactive) system?
  57                 signal.signal(signal.SIGHUP, signal.SIG_DFL)
  58                 if cwd:
  59                     os.chdir(cwd)
  60                 os.execvp(args[0], args)
  61             except:
  62                 print('CHILD EXECEPTION:')
  63                 traceback.print_exc()
  64             finally:
  65                 os._exit(255)
  66
  67         try:
  68             self.poll = select.poll()
  69             self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR |
  70                                select.POLLHUP | select.POLLNVAL)
  71         except:
  72             self.close()
  73             raise
  74
  75     def kill(self, sig):
  76         """Send unix signal "sig" to the child process.
  77
  78         Args:
  79             sig: The signal number to send.
  80
  81         Returns:
  82             Nothing.
  83         """
  84
  85         os.kill(self.pid, sig)
  86
  87     def checkalive(self):
  88         """Determine whether the child process is still running.
  89
  90         Returns:
  91             tuple:
  92                 True if process is alive, else False
  93                 0 if process is alive, else exit code of process
  94                 string describing what happened ('' or 'status/signal n')
  95         """
  96
  97         if self.waited:
  98             return False, self.exit_code, self.exit_info
  99
 100         w = os.waitpid(self.pid, os.WNOHANG)
 101         if w[0] == 0:
 102             return True, 0, 'running'
 103         status = w[1]
 104
 105         if os.WIFEXITED(status):
 106             self.exit_code = os.WEXITSTATUS(status)
 107             self.exit_info = 'status %d' % self.exit_code
 108         elif os.WIFSIGNALED(status):
 109             signum = os.WTERMSIG(status)
 110             self.exit_code = -signum
 111             self.exit_info = 'signal %d (%s)' % (signum, signal.Signals(signum).name)
 112         self.waited = True
 113         return False, self.exit_code, self.exit_info
 114
 115     def isalive(self):
 116         """Determine whether the child process is still running.
 117
 118         Args:
 119             None.
 120
 121         Returns:
 122             Boolean indicating whether process is alive.
 123         """
 124         return self.checkalive()[0]
 125
 126     def send(self, data):
 127         """Send data to the sub-process's stdin.
 128
 129         Args:
 130             data: The data to send to the process.
 131
 132         Returns:
 133             Nothing.
 134         """
 135
 136         os.write(self.fd, data.encode(errors='replace'))
 137
 138     def expect(self, patterns):
 139         """Wait for the sub-process to emit specific data.
 140
 141         This function waits for the process to emit one pattern from the
 142         supplied list of patterns, or for a timeout to occur.
 143
 144         Args:
 145             patterns: A list of strings or regex objects that we expect to
 146                 see in the sub-process' stdout.
 147
 148         Returns:
 149             The index within the patterns array of the pattern the process
 150             emitted.
 151
 152         Notable exceptions:
 153             Timeout, if the process did not emit any of the patterns within
 154             the expected time.
 155         """
 156
 157         for pi in range(len(patterns)):
 158             if type(patterns[pi]) == type(''):
 159                 patterns[pi] = re.compile(patterns[pi])
 160
 161         tstart_s = time.time()
 162         try:
 163             while True:
 164                 earliest_m = None
 165                 earliest_pi = None
 166                 for pi in range(len(patterns)):
 167                     pattern = patterns[pi]
 168                     m = pattern.search(self.buf)
 169                     if not m:
 170                         continue
 171                     if earliest_m and m.start() >= earliest_m.start():
 172                         continue
 173                     earliest_m = m
 174                     earliest_pi = pi
 175                 if earliest_m:
 176                     pos = earliest_m.start()
 177                     posafter = earliest_m.end()
 178                     self.before = self.buf[:pos]
 179                     self.after = self.buf[pos:posafter]
 180                     self.output += self.buf[:posafter]
 181                     self.buf = self.buf[posafter:]
 182                     return earliest_pi
 183                 tnow_s = time.time()
 184                 if self.timeout:
 185                     tdelta_ms = (tnow_s - tstart_s) * 1000
 186                     poll_maxwait = self.timeout - tdelta_ms
 187                     if tdelta_ms > self.timeout:
 188                         raise Timeout()
 189                 else:
 190                     poll_maxwait = None
 191                 events = self.poll.poll(poll_maxwait)
 192                 if not events:
 193                     raise Timeout()
 194                 try:
 195                     c = os.read(self.fd, 1024).decode(errors='replace')
 196                 except OSError as err:
 197                     # With sandbox, try to detect when U-Boot exits when it
 198                     # shouldn't and explain why. This is much more friendly than
 199                     # just dying with an I/O error
 200                     if err.errno == 5:  # Input/output error
 201                         alive, _, info = self.checkalive()
 202                         if alive:
 203                             raise err
 204                         raise ValueError('U-Boot exited with %s' % info)
 205                     raise err
 206                 if self.logfile_read:
 207                     self.logfile_read.write(c)
 208                 self.buf += c
 209                 # count=0 is supposed to be the default, which indicates
 210                 # unlimited substitutions, but in practice the version of
 211                 # Python in Ubuntu 14.04 appears to default to count=2!
 212                 self.buf = self.re_vt100.sub('', self.buf, count=1000000)
 213         finally:
 214             if self.logfile_read:
 215                 self.logfile_read.flush()
 216
 217     def close(self):
 218         """Close the stdio connection to the sub-process.
 219
 220         This also waits a reasonable time for the sub-process to stop running.
 221
 222         Args:
 223             None.
 224
 225         Returns:
 226             Nothing.
 227         """
 228
 229         os.close(self.fd)
 230         for _ in range(100):
 231             if not self.isalive():
 232                 break
 233             time.sleep(0.1)
 234
 235     def get_expect_output(self):
 236         """Return the output read by expect()
 237
 238         Returns:
 239             The output processed by expect(), as a string.
 240         """
 241         return self.output