• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2016 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Various utility functions."""
16
17import errno
18import functools
19import os
20import signal
21import subprocess
22import sys
23import tempfile
24import time
25
26_path = os.path.realpath(__file__ + '/../..')
27if sys.path[0] != _path:
28    sys.path.insert(0, _path)
29del _path
30
31# pylint: disable=wrong-import-position
32import rh.shell
33import rh.signals
34
35
36def timedelta_str(delta):
37    """A less noisy timedelta.__str__.
38
39    The default timedelta stringification contains a lot of leading zeros and
40    uses microsecond resolution.  This makes for noisy output.
41    """
42    total = delta.total_seconds()
43    hours, rem = divmod(total, 3600)
44    mins, secs = divmod(rem, 60)
45    ret = f'{int(secs)}.{delta.microseconds // 1000:03}s'
46    if mins:
47        ret = f'{int(mins)}m{ret}'
48    if hours:
49        ret = f'{int(hours)}h{ret}'
50    return ret
51
52
53class CompletedProcess(subprocess.CompletedProcess):
54    """An object to store various attributes of a child process.
55
56    This is akin to subprocess.CompletedProcess.
57    """
58
59    def __init__(self, args=None, returncode=None, stdout=None, stderr=None):
60        super().__init__(
61            args=args, returncode=returncode, stdout=stdout, stderr=stderr)
62
63    @property
64    def cmd(self):
65        """Alias to self.args to better match other subprocess APIs."""
66        return self.args
67
68    @property
69    def cmdstr(self):
70        """Return self.cmd as a nicely formatted string (useful for logs)."""
71        return rh.shell.cmd_to_str(self.cmd)
72
73
74class CalledProcessError(subprocess.CalledProcessError):
75    """Error caught in run() function.
76
77    This is akin to subprocess.CalledProcessError.  We do not support |output|,
78    only |stdout|.
79
80    Attributes:
81      returncode: The exit code of the process.
82      cmd: The command that triggered this exception.
83      msg: Short explanation of the error.
84    """
85
86    def __init__(self, returncode, cmd, stdout=None, stderr=None, msg=None):
87        super().__init__(returncode, cmd, stdout, stderr=stderr)
88
89        # The parent class will set |output|, so delete it. If Python ever drops
90        # this output/stdout compat logic, we can drop this to match.
91        del self.output
92        self._stdout = stdout
93
94        self.msg = msg
95
96    @property
97    def stdout(self):
98        """Override parent's usage of .output"""
99        return self._stdout
100
101    @stdout.setter
102    def stdout(self, value):
103        """Override parent's usage of .output"""
104        self._stdout = value
105
106    @property
107    def cmdstr(self):
108        """Return self.cmd as a well shell-quoted string for debugging."""
109        return '' if self.cmd is None else rh.shell.cmd_to_str(self.cmd)
110
111    def stringify(self, stdout=True, stderr=True):
112        """Custom method for controlling what is included in stringifying this.
113
114        Args:
115          stdout: Whether to include captured stdout in the return value.
116          stderr: Whether to include captured stderr in the return value.
117
118        Returns:
119          A summary string for this result.
120        """
121        items = [
122            f'return code: {self.returncode}; command: {self.cmdstr}',
123        ]
124        if stderr and self.stderr:
125            items.append(self.stderr)
126        if stdout and self.stdout:
127            items.append(self.stdout)
128        if self.msg:
129            items.append(self.msg)
130        return '\n'.join(items)
131
132    def __str__(self):
133        return self.stringify()
134
135
136class TerminateCalledProcessError(CalledProcessError):
137    """We were signaled to shutdown while running a command.
138
139    Client code shouldn't generally know, nor care about this class.  It's
140    used internally to suppress retry attempts when we're signaled to die.
141    """
142
143
144def _kill_child_process(proc, int_timeout, kill_timeout, cmd, original_handler,
145                        signum, frame):
146    """Used as a signal handler by RunCommand.
147
148    This is internal to Runcommand.  No other code should use this.
149    """
150    if signum:
151        # If we've been invoked because of a signal, ignore delivery of that
152        # signal from this point forward.  The invoking context of this func
153        # restores signal delivery to what it was prior; we suppress future
154        # delivery till then since this code handles SIGINT/SIGTERM fully
155        # including delivering the signal to the original handler on the way
156        # out.
157        signal.signal(signum, signal.SIG_IGN)
158
159    # Do not trust Popen's returncode alone; we can be invoked from contexts
160    # where the Popen instance was created, but no process was generated.
161    if proc.returncode is None and proc.pid is not None:
162        try:
163            while proc.poll_lock_breaker() is None and int_timeout >= 0:
164                time.sleep(0.1)
165                int_timeout -= 0.1
166
167            proc.terminate()
168            while proc.poll_lock_breaker() is None and kill_timeout >= 0:
169                time.sleep(0.1)
170                kill_timeout -= 0.1
171
172            if proc.poll_lock_breaker() is None:
173                # Still doesn't want to die.  Too bad, so sad, time to die.
174                proc.kill()
175        except EnvironmentError as e:
176            print(f'Ignoring unhandled exception in _kill_child_process: {e}',
177                  file=sys.stderr)
178
179        # Ensure our child process has been reaped, but don't wait forever.
180        proc.wait_lock_breaker(timeout=60)
181
182    if not rh.signals.relay_signal(original_handler, signum, frame):
183        # Mock up our own, matching exit code for signaling.
184        raise TerminateCalledProcessError(
185            signum << 8, cmd, msg=f'Received signal {signum}')
186
187
188class _Popen(subprocess.Popen):
189    """subprocess.Popen derivative customized for our usage.
190
191    Specifically, we fix terminate/send_signal/kill to work if the child process
192    was a setuid binary; on vanilla kernels, the parent can wax the child
193    regardless, on goobuntu this apparently isn't allowed, thus we fall back
194    to the sudo machinery we have.
195
196    While we're overriding send_signal, we also suppress ESRCH being raised
197    if the process has exited, and suppress signaling all together if the
198    process has knowingly been waitpid'd already.
199    """
200
201    # pylint: disable=arguments-differ,arguments-renamed
202    def send_signal(self, signum):
203        if self.returncode is not None:
204            # The original implementation in Popen allows signaling whatever
205            # process now occupies this pid, even if the Popen object had
206            # waitpid'd.  Since we can escalate to sudo kill, we do not want
207            # to allow that.  Fixing this addresses that angle, and makes the
208            # API less sucky in the process.
209            return
210
211        try:
212            os.kill(self.pid, signum)
213        except EnvironmentError as e:
214            if e.errno == errno.ESRCH:
215                # Since we know the process is dead, reap it now.
216                # Normally Popen would throw this error- we suppress it since
217                # frankly that's a misfeature and we're already overriding
218                # this method.
219                self.poll()
220            else:
221                raise
222
223    def _lock_breaker(self, func, *args, **kwargs):
224        """Helper to manage the waitpid lock.
225
226        Workaround https://bugs.python.org/issue25960.
227        """
228        # If the lock doesn't exist, or is not locked, call the func directly.
229        lock = getattr(self, '_waitpid_lock', None)
230        if lock is not None and lock.locked():
231            try:
232                lock.release()
233                return func(*args, **kwargs)
234            finally:
235                if not lock.locked():
236                    lock.acquire()
237        else:
238            return func(*args, **kwargs)
239
240    def poll_lock_breaker(self, *args, **kwargs):
241        """Wrapper around poll() to break locks if needed."""
242        return self._lock_breaker(self.poll, *args, **kwargs)
243
244    def wait_lock_breaker(self, *args, **kwargs):
245        """Wrapper around wait() to break locks if needed."""
246        return self._lock_breaker(self.wait, *args, **kwargs)
247
248
249# We use the keyword arg |input| which trips up pylint checks.
250# pylint: disable=redefined-builtin
251def run(cmd, redirect_stdout=False, redirect_stderr=False, cwd=None, input=None,
252        shell=False, env=None, extra_env=None, combine_stdout_stderr=False,
253        check=True, int_timeout=1, kill_timeout=1, capture_output=False):
254    """Runs a command.
255
256    Args:
257      cmd: cmd to run.  Should be input to subprocess.Popen.  If a string, shell
258          must be true.  Otherwise the command must be an array of arguments,
259          and shell must be false.
260      redirect_stdout: Returns the stdout.
261      redirect_stderr: Holds stderr output until input is communicated.
262      cwd: The working directory to run this cmd.
263      input: The data to pipe into this command through stdin.  If a file object
264          or file descriptor, stdin will be connected directly to that.
265      shell: Controls whether we add a shell as a command interpreter.  See cmd
266          since it has to agree as to the type.
267      env: If non-None, this is the environment for the new process.
268      extra_env: If set, this is added to the environment for the new process.
269          This dictionary is not used to clear any entries though.
270      combine_stdout_stderr: Combines stdout and stderr streams into stdout.
271      check: Whether to raise an exception when command returns a non-zero exit
272          code, or return the CompletedProcess object containing the exit code.
273          Note: will still raise an exception if the cmd file does not exist.
274      int_timeout: If we're interrupted, how long (in seconds) should we give
275          the invoked process to clean up before we send a SIGTERM.
276      kill_timeout: If we're interrupted, how long (in seconds) should we give
277          the invoked process to shutdown from a SIGTERM before we SIGKILL it.
278      capture_output: Set |redirect_stdout| and |redirect_stderr| to True.
279
280    Returns:
281      A CompletedProcess object.
282
283    Raises:
284      CalledProcessError: Raises exception on error.
285    """
286    if capture_output:
287        redirect_stdout, redirect_stderr = True, True
288
289    # Set default for variables.
290    popen_stdout = None
291    popen_stderr = None
292    stdin = None
293    result = CompletedProcess()
294
295    # Force the timeout to float; in the process, if it's not convertible,
296    # a self-explanatory exception will be thrown.
297    kill_timeout = float(kill_timeout)
298
299    def _get_tempfile():
300        try:
301            return tempfile.TemporaryFile(buffering=0)
302        except EnvironmentError as e:
303            if e.errno != errno.ENOENT:
304                raise
305            # This can occur if we were pointed at a specific location for our
306            # TMP, but that location has since been deleted.  Suppress that
307            # issue in this particular case since our usage gurantees deletion,
308            # and since this is primarily triggered during hard cgroups
309            # shutdown.
310            return tempfile.TemporaryFile(dir='/tmp', buffering=0)
311
312    # Modify defaults based on parameters.
313    # Note that tempfiles must be unbuffered else attempts to read
314    # what a separate process did to that file can result in a bad
315    # view of the file.
316    # The Popen API accepts either an int or a file handle for stdout/stderr.
317    # pylint: disable=redefined-variable-type
318    if redirect_stdout:
319        popen_stdout = _get_tempfile()
320
321    if combine_stdout_stderr:
322        popen_stderr = subprocess.STDOUT
323    elif redirect_stderr:
324        popen_stderr = _get_tempfile()
325    # pylint: enable=redefined-variable-type
326
327    # If subprocesses have direct access to stdout or stderr, they can bypass
328    # our buffers, so we need to flush to ensure that output is not interleaved.
329    if popen_stdout is None or popen_stderr is None:
330        sys.stdout.flush()
331        sys.stderr.flush()
332
333    # If input is a string, we'll create a pipe and send it through that.
334    # Otherwise we assume it's a file object that can be read from directly.
335    if isinstance(input, str):
336        stdin = subprocess.PIPE
337        input = input.encode('utf-8')
338    elif input is not None:
339        stdin = input
340        input = None
341
342    if isinstance(cmd, str):
343        if not shell:
344            raise Exception('Cannot run a string command without a shell')
345        cmd = ['/bin/bash', '-c', cmd]
346        shell = False
347    elif shell:
348        raise Exception('Cannot run an array command with a shell')
349
350    # If we are using enter_chroot we need to use enterchroot pass env through
351    # to the final command.
352    env = env.copy() if env is not None else os.environ.copy()
353    env.update(extra_env if extra_env else {})
354
355    def ensure_text(s):
356        """Make sure |s| is a string if it's bytes."""
357        if isinstance(s, bytes):
358            s = s.decode('utf-8', 'replace')
359        return s
360
361    result.args = cmd
362
363    proc = None
364    try:
365        proc = _Popen(cmd, cwd=cwd, stdin=stdin, stdout=popen_stdout,
366                      stderr=popen_stderr, shell=False, env=env,
367                      close_fds=True)
368
369        old_sigint = signal.getsignal(signal.SIGINT)
370        handler = functools.partial(_kill_child_process, proc, int_timeout,
371                                    kill_timeout, cmd, old_sigint)
372        # We have to ignore ValueError in case we're run from a thread.
373        try:
374            signal.signal(signal.SIGINT, handler)
375        except ValueError:
376            old_sigint = None
377
378        old_sigterm = signal.getsignal(signal.SIGTERM)
379        handler = functools.partial(_kill_child_process, proc, int_timeout,
380                                    kill_timeout, cmd, old_sigterm)
381        try:
382            signal.signal(signal.SIGTERM, handler)
383        except ValueError:
384            old_sigterm = None
385
386        try:
387            (result.stdout, result.stderr) = proc.communicate(input)
388        finally:
389            if old_sigint is not None:
390                signal.signal(signal.SIGINT, old_sigint)
391            if old_sigterm is not None:
392                signal.signal(signal.SIGTERM, old_sigterm)
393
394            if popen_stdout:
395                # The linter is confused by how stdout is a file & an int.
396                # pylint: disable=maybe-no-member,no-member
397                popen_stdout.seek(0)
398                result.stdout = popen_stdout.read()
399                popen_stdout.close()
400
401            if popen_stderr and popen_stderr != subprocess.STDOUT:
402                # The linter is confused by how stderr is a file & an int.
403                # pylint: disable=maybe-no-member,no-member
404                popen_stderr.seek(0)
405                result.stderr = popen_stderr.read()
406                popen_stderr.close()
407
408        result.returncode = proc.returncode
409
410        if check and proc.returncode:
411            msg = f'cwd={cwd}'
412            if extra_env:
413                msg += f', extra env={extra_env}'
414            raise CalledProcessError(
415                result.returncode, result.cmd, msg=msg,
416                stdout=ensure_text(result.stdout),
417                stderr=ensure_text(result.stderr))
418    except OSError as e:
419        # Avoid leaking tempfiles.
420        if popen_stdout is not None and not isinstance(popen_stdout, int):
421            popen_stdout.close()
422        if popen_stderr is not None and not isinstance(popen_stderr, int):
423            popen_stderr.close()
424
425        estr = str(e)
426        if e.errno == errno.EACCES:
427            estr += '; does the program need `chmod a+x`?'
428        if not check:
429            result = CompletedProcess(args=cmd, stderr=estr, returncode=255)
430        else:
431            raise CalledProcessError(
432                result.returncode, result.cmd, msg=estr,
433                stdout=ensure_text(result.stdout),
434                stderr=ensure_text(result.stderr)) from e
435    finally:
436        if proc is not None:
437            # Ensure the process is dead.
438            # Some pylint3 versions are confused here.
439            # pylint: disable=too-many-function-args
440            _kill_child_process(proc, int_timeout, kill_timeout, cmd, None,
441                                None, None)
442
443    # Make sure output is returned as a string rather than bytes.
444    result.stdout = ensure_text(result.stdout)
445    result.stderr = ensure_text(result.stderr)
446
447    return result
448# pylint: enable=redefined-builtin
449