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