import os import select import selectors import subprocess import threading from common import Results, TimeoutException from time import monotonic as _time _PopenSelector = selectors.PollSelector _PIPE_BUF = getattr(select, 'PIPE_BUF', 512) _DEBUG_LEN = 80 class _processCommunicator(object): """ This class allows greater control over stdin than using Popen.communicate(). Popen.communicate() closes stdin as soon as data is written. This causes TLS clients (OpenSSL derivatives) to shut down before the handshake is complete. To prevent a premature shutdown, we need to wait until the handshake is complete before writing to stdin. To accomplish this we `poll` stdout for a send marker. Once that marker is found, we can write input data to stdin. The benefit of using `poll` and `os.read` is that we get non-blocking IO. Our timeouts are much more reliable, and we don't risk deadlocking on a readline() call which will never complete. Another method is to read stdout line by line. This removes a lot of code that registers and unregisters file descriptors with a selector. It would make reading and writing sequential (as opposed to event based), which can be easier to read and maintain. The downsides with this method exist in the current integration test framework. We rely on sleeps and waits, and still hit hard to debug deadlocks from time to time. """ def __init__(self, proc, name): self.proc = proc self.wait_for_marker = None self.name = name # If the process times out, communicate() is called once more to pick # up any data remaining in stdout/stderr. This flags lets us know if # we need to do initial setup on the file descriptors, or if it was done # during the initial call. self._communication_started = False def wait_for(self, wait_for_marker, timeout=None): """ Wait for a specific marker in stdout. If the marker is not seen, a timeout will be raised. """ self.wait_for_marker = wait_for_marker stdout = None stderr = None try: stdout, stderr = self._communicate(None, timeout=timeout) finally: self._communication_started = True return (stdout, stderr) def communicate(self, input_data=None, send_marker_list=None, close_marker=None, kill_marker=None, send_with_newline=False, timeout=None): """ Communicates with the managed process. If send_marker_list is set, input_data will not be sent until the marker is seen. This method acts very similar to the Popen.communicate method. The only difference is the send_marker_list and close_marker. """ self.wait_for_marker = None stdout = None stderr = None try: stdout, stderr = self._communicate( input_data, send_marker_list, close_marker, kill_marker, send_with_newline, timeout ) finally: self._communication_started = True return (stdout, stderr) def _communicate(self, input_data=None, send_marker_list=None, close_marker=None, kill_marker=None, send_with_newline=False, timeout=None): """ This method will read and write data to a subprocess in a non-blocking manner. The code is heavily based on Popen.communicate. There are a couple differences: * STDIN is not registered for events until the read_to_send marker is found * STDIN is only closed after all registered events have been processed (including pending stdout/stderr events, allowing more data to be stored). """ if input_data is not None and self.proc.stdin: # Flush stdio buffer. This might block, if the user has # been writing to .stdin in an uncontrolled fashion. try: self.proc.stdin.flush() except BrokenPipeError: pass # communicate() must ignore BrokenPipeError. # The process' stdout and stderr are stored in a map, with two variable # pointing to the file objects. This allows us to include stdout/stderr # data in a timeout exception. if not self._communication_started: self._fileobj2output = {} if self.proc.stdout: self._fileobj2output[self.proc.stdout] = [] if self.proc.stderr: self._fileobj2output[self.proc.stderr] = [] stdout = self._fileobj2output[self.proc.stdout] stderr = self._fileobj2output[self.proc.stderr] input_data_len = 0 input_data_offset = 0 input_data_sent = False send_marker = None if send_marker_list: send_marker = send_marker_list.pop(0) # Keeping track of the original timeout value, and the expected end # time of the operation allow us to timeout while reads/writes are # still pending. It also allows us to only wait for the remainder of # the timeout after reads/writes have completed. orig_timeout = timeout if timeout is not None: endtime = _time() + timeout else: endtime = None with _PopenSelector() as selector: if self.proc.stdout and not self.proc.stdout.closed: selector.register(self.proc.stdout, selectors.EVENT_READ) if self.proc.stderr and not self.proc.stderr.closed: selector.register(self.proc.stderr, selectors.EVENT_READ) while selector.get_map(): timeout = self._remaining_time(endtime) if timeout is not None and timeout < 0: self._check_timeout(endtime, orig_timeout, stdout, stderr, skip_check_and_raise=True) raise RuntimeError( # Impossible :) '_check_timeout(..., skip_check_and_raise=True) ' 'failed to raise TimeoutExpired.') ready = selector.select(timeout) self._check_timeout(endtime, orig_timeout, stdout, stderr) for key, events in ready: # STDIN is only registered to receive events after the send_marker is found. if key.fileobj is self.proc.stdin: print(f'{self.name}: stdin available') chunk = input_view[input_data_offset: input_data_offset + _PIPE_BUF] try: input_data_offset += os.write(key.fd, chunk) print(f'{self.name}: sent') except BrokenPipeError: selector.unregister(key.fileobj) else: if input_data_offset >= input_data_len: selector.unregister(key.fileobj) input_data_sent = True input_data_offset = 0 if send_marker_list: send_marker = send_marker_list.pop(0) print(f'{self.name}: next send_marker is {send_marker}') elif key.fileobj in (self.proc.stdout, self.proc.stderr): print(f'{self.name}: stdout available') data = os.read(key.fd, 32768) if not data: selector.unregister(key.fileobj) data_str = str(data) data_debug = data_str[:_DEBUG_LEN] if len(data_str) > _DEBUG_LEN: data_debug += f' ...({len(data_str) - _DEBUG_LEN} more bytes)' # fileobj2output[key.fileobj] is a list of data chunks # that get joined later self._fileobj2output[key.fileobj].append(data) # If we are looking for, and find, the ready-to-send marker, then # register STDIN to receive events. If there is no data to send, # just mark input_send as true so we can close out STDIN. if send_marker: print(f'{self.name}: looking for send_marker {send_marker} in {data_debug}') if send_marker is not None and send_marker in data_str: print(f'{self.name}: found {send_marker}') if self.proc.stdin and input_data: selector.register( self.proc.stdin, selectors.EVENT_WRITE) message = input_data.pop(0) if send_with_newline: message += b'\n' # Data destined for stdin is stored in a memoryview input_view = memoryview(message) input_data_len = len(message) input_data_sent = False print(f'{self.name}: will send {message}') else: input_data_sent = True print(f'{self.name}: will send nothing') if self.wait_for_marker: print(f'{self.name}: looking for wait_for_marker {self.wait_for_marker} in {data_debug}') if self.wait_for_marker is not None and self.wait_for_marker in data_str: selector.unregister(self.proc.stdout) selector.unregister(self.proc.stderr) return None, None if kill_marker: print(f'{self.name}: looking for kill_marker {kill_marker} in {data}') if kill_marker is not None and kill_marker in data: selector.unregister(self.proc.stdout) selector.unregister(self.proc.stderr) self.proc.kill() # If we have finished sending all our input, and have received the # ready-to-send marker, we can close out stdin. if self.proc.stdin and input_data_sent and not input_data: print(f'{self.name}: finished sending') if close_marker: print(f'{self.name}: looking for close_marker {close_marker} in {data_debug}') if close_marker is None or (close_marker and close_marker in data_str): print(f'{self.name}: closing stdin') input_data_sent = None self.proc.stdin.close() self.proc.wait(timeout=self._remaining_time(endtime)) # All data exchanged. Translate lists into strings. if stdout is not None: stdout = b''.join(stdout) if stderr is not None: stderr = b''.join(stderr) return (stdout, stderr) def _remaining_time(self, endtime): """Convenience for _communicate when computing timeouts.""" if endtime is None: return None else: return endtime - _time() def _check_timeout(self, endtime, orig_timeout, stdout_seq, stderr_seq, skip_check_and_raise=False): """ Convenience for checking if a timeout has expired. NOTE: This method is included here to prevent our custom _communicate method from relying on a particular version of Python. """ if endtime is None: return if skip_check_and_raise or _time() > endtime: raise subprocess.TimeoutExpired( self.proc.args, orig_timeout, output=b''.join(stdout_seq) if stdout_seq else None, stderr=b''.join(stderr_seq) if stderr_seq else None) class ManagedProcess(threading.Thread): """ A ManagedProcess is a thread that monitors a subprocess. This class provides a single place to control process timeouts and cleanup. The stdin/stdout/stderr and exist code a monitored and results are made available to the caller. """ def __init__(self, cmd_line, provider_set_ready_condition, wait_for_marker=None, send_marker_list=None, close_marker=None, timeout=5, data_source=None, env_overrides=dict(), expect_stderr=False, kill_marker=None, send_with_newline=False): threading.Thread.__init__(self) proc_env = os.environ.copy() for key in env_overrides: proc_env[key] = env_overrides[key] self.proc_env = proc_env # Command line to execute in the subprocess self.cmd_line = list(map(str, cmd_line)) # Total time to wait until killing the subprocess self.timeout = timeout # Condition variable indicating when results are ready to be collected self.results_condition = threading.Condition() self.results = None # Condition variable indicating when this subprocess has been launched successfully self.ready_condition = threading.Condition() self.process_ready = False self.provider_set_ready_condition = provider_set_ready_condition # Indicates the process has completed some initial setup and is ready for testing self.ready_to_test = wait_for_marker self.close_marker = close_marker self.data_source = data_source self.send_marker_list = send_marker_list self.expect_stderr = expect_stderr self.kill_marker = kill_marker self.send_with_newline = send_with_newline if data_source is not None: if type(data_source) is not list: self.data_source = [data_source] if send_marker_list is not None: if type(send_marker_list) is not list: self.send_marker_list = [send_marker_list] def run(self): with self.results_condition: try: proc = subprocess.Popen(self.cmd_line, env=self.proc_env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True) self.proc = proc except Exception as ex: self.results = Results( None, None, None, ex, self.expect_stderr) raise ex communicator = _processCommunicator(proc, self.cmd_line[0]) if self.ready_to_test is not None: # Some processes won't be ready until they have emitted some string in stdout. communicator.wait_for(self.ready_to_test, timeout=self.timeout) # Let any threads waiting on process launch proceed self.provider_set_ready_condition() proc_results = None try: proc_results = communicator.communicate( input_data=self.data_source, send_marker_list=self.send_marker_list, close_marker=self.close_marker, kill_marker=self.kill_marker, send_with_newline=self.send_with_newline, timeout=self.timeout, ) self.results = Results( proc_results[0], proc_results[1], proc.returncode, None, expect_stderr=self.expect_stderr, expect_nonzero_exit=self.kill_marker is not None ) except subprocess.TimeoutExpired as ex: proc.kill() wrapped_ex = TimeoutException(ex) # Read any remaining output proc_results = communicator.communicate() self.results = Results( proc_results[0], proc_results[1], proc.returncode, wrapped_ex, self.expect_stderr) except Exception as ex: self.results = Results( proc_results[0], proc_results[1], proc.returncode, ex, self.expect_stderr) raise ex finally: # This data is dumped to stdout so we capture this # information no matter where a test fails. print("Command line: {}".format(" ".join(self.cmd_line))) print("Exit code: {}".format(proc.returncode)) print("Stdout: {}".format( proc_results[0].decode("utf-8", "backslashreplace"))) print("Stderr: {}".format( proc_results[1].decode("utf-8", "backslashreplace"))) def kill(self): self.proc.kill() def _process_ready(self): """Condition variable predicate""" return self.process_ready is True def _results_ready(self): """Condition variable predicate""" return self.results is not None def get_cmd_line(self): return self.cmd_line def launch(self): """ This method must be implemented by the subclass. It should call the run function. """ raise NotImplementedError def get_results(self, send_data=None): """ Block until the results are ready, or a timeout is reached. Return the results, or raise the timeout exception. """ with self.results_condition: result = self.results_condition.wait_for( self._results_ready, timeout=self.timeout) if result is False: raise Exception("Timeout") yield self.results