1# Licensed under the Apache License, Version 2.0 (the "License"); you may
2# not use this file except in compliance with the License. You may obtain
3# a copy of the License at
4#
5# http://www.apache.org/licenses/LICENSE-2.0
6#
7# Unless required by applicable law or agreed to in writing, software
8# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10# License for the specific language governing permissions and limitations
11# under the License.
12
13import contextlib
14import fcntl
15import itertools
16import multiprocessing
17import os
18import pty
19import re
20import signal
21import struct
22import sys
23import tempfile
24import termios
25import time
26import traceback
27
28import types
29from typing import Optional, Generator, Tuple
30import typing
31
32
33(LINES, COLUMNS) = (24, 80)
34
35
36def _open_fifo(path: Optional[str],
37 write: bool) -> Optional[int]:
38 if path is None:
39 return None
40 return os.open(path, os.O_WRONLY if write else os.O_RDONLY)
41
42
43class IsolationEnvironment:
44 def __init__(self, pid: int, ptm_fd: int,
45 stdin_fifo: Optional[str] = None,
46 stdout_fifo: Optional[str] = None,
47 stderr_fifo: Optional[str] = None):
48 self._pid = pid
49
50 self._ptm_fd = ptm_fd
51 self._tty = os.fdopen(ptm_fd, 'r')
52
53 self._stdin_fifo_fd = _open_fifo(stdin_fifo, True)
54 self._stdout_fifo_fd = _open_fifo(stdout_fifo, False)
55 self._stderr_fifo_fd = _open_fifo(stderr_fifo, False)
56
57 self._exit_code: Optional[int] = None
58
59 self._raw_lines: typing.List[str] = []
60 self._lines: typing.List[str] = []
61
62 def interrupt(self) -> None:
63 self.write(b'\x03')
64 time.sleep(0.001)
65
66 def write(self, data: bytes) -> None:
67 os.write(self._ptm_fd, data)
68
69 def readline(self) -> str:
70 line = self._tty.readline()
71 self._raw_lines.append(line)
72 return line
73
74 def error_output(self) -> str:
75 if self._stderr_fifo_fd is None:
76 return ''
77 with os.fdopen(self._stderr_fifo_fd) as errf:
78 self._stderr_fifo_fd = None
79 return errf.read()
80
81 def stdout_pipe(self) -> typing.TextIO:
82 assert self._stdout_fifo_fd is not None
83 stdout = os.fdopen(self._stdout_fifo_fd)
84 self._stdout_fifo_fd = None
85 return stdout
86
87 def stdin_pipe(self) -> typing.TextIO:
88 assert self._stdin_fifo_fd is not None
89 return os.fdopen(self._stdin_fifo_fd, 'w', closefd=False)
90
91 def close(self, get_return_code: typing.Callable[[], int]) -> None:
92 for i in range(100):
93 if os.waitpid(self._pid, os.WNOHANG) != (0, 0):
94 break
95 time.sleep(0.001)
96 else:
97 os.kill(self._pid, signal.SIGTERM)
98 os.waitpid(self._pid, 0)
99 self._exit_code = get_return_code()
100 self._tty.close()
101 for fifo_fd in (self._stdin_fifo_fd,
102 self._stdout_fifo_fd,
103 self._stderr_fifo_fd):
104 if fifo_fd is not None:
105 os.close(fifo_fd)
106
107 def exit_code(self) -> int:
108 assert self._exit_code is not None
109 return self._exit_code
110
111 def record_output(self, line: str) -> None:
112 self._lines.append(line)
113
114 def recorded_output(self) -> typing.Tuple[typing.List[str],
115 typing.List[str]]:
116 return self._raw_lines[:], self._lines[:]
117
118
119class PagerControl:
120 _page_end = None
121 _ctrl_chars = re.compile(r'\x1b\['
122 r'(\?|[0-2]?K|[0-9]*;?[0-9]*H|[0-3]?J|[0-9]*m)')
123
124 def __init__(self, isolation_env: IsolationEnvironment):
125 self.env = isolation_env
126 self._total_lines = 0
127 self._lines = self._iter_lines()
128
129 def _iter_lines(self) -> Generator[typing.Union[str, None],
130 None, None]:
131 def get_content(segment: str) -> Tuple[bool,
132 typing.Union[str, None]]:
133 if not segment:
134 return False, ''
135 if segment == '\x1b[1m~\x1b[0m\n':
136 # Ignore lines that are filling blank vertical space with
137 # '~' after hitting Ctrl-C when the screen is not full
138 return False, ''
139 visible = self._ctrl_chars.sub('', segment)
140 if ((visible.rstrip() == ':' or '(END)' in visible
141 or 'Waiting for data...' in visible)
142 and segment.replace('\x1b[m', '') != visible):
143 return True, self._page_end
144 elif visible.rstrip() or segment == visible:
145 self._total_lines += 1
146 self.env.record_output(visible)
147 return True, visible
148 return False, ''
149
150 while True:
151 line = '\x1b[?'
152 while line.lstrip(' q').startswith('\x1b[?'):
153 rawline = self.env.readline()
154 line = (rawline.replace('\x07', '') # Ignore bell
155 .replace('\x1b[m', '')) # Ignore style reset
156 before, reset, after = line.partition('\x1b[2J')
157
158 valid, content = get_content(before)
159 if valid:
160 yield content
161 if reset and not (valid and (content is self._page_end)):
162 yield self._page_end
163 valid, content = get_content(after)
164 if valid:
165 yield content
166
167 def read_lines(self, count: int) -> typing.Iterator[str]:
168 return itertools.islice((line for line in self._lines
169 if line is not self._page_end), count)
170
171 def _lines_in_page(self) -> int:
172 original_count = self._total_lines
173 try:
174 for line in self._lines:
175 if line is self._page_end:
176 break
177 except IOError:
178 pass
179 return self._total_lines - original_count
180
181 def advance(self) -> int:
182 self.env.write(b' ')
183 return self._lines_in_page()
184
185 def quit(self) -> int:
186 self.env.write(b'q')
187 return self._lines_in_page()
188
189 def total_lines(self) -> int:
190 return self._total_lines
191
192
193@contextlib.contextmanager
194def _fifo(fifo_path: str) -> Generator[str, None, None]:
195 try:
196 os.mkfifo(fifo_path, 0o600)
197 yield fifo_path
198 finally:
199 try:
200 os.unlink(fifo_path)
201 except OSError:
202 pass
203
204
205@contextlib.contextmanager
206def _fifos(*fifo_names: Optional[str]) -> Generator[typing.List[Optional[str]],
207 None, None]:
208 with tempfile.TemporaryDirectory() as directory:
209 with contextlib.ExitStack() as stack:
210 fifos = [stack.enter_context(_fifo(os.path.join(directory,
211 name)))
212 if name is not None else None
213 for name in fifo_names]
214 yield fifos
215
216
217class TestProcessNotComplete(Exception):
218 pass
219
220
221class TimeoutException(Exception):
222 def __init__(self, pid: Optional[int]):
223 super(Exception, self).__init__(f'Test process {pid} timed out')
224
225
226@contextlib.contextmanager
227def isolate(child_function: typing.Callable[[], int],
228 stdin_pipe: bool = False,
229 stdout_pipe: bool = False,
230 stderr_pipe: bool = True,
231 *,
232 lines: int = LINES,
233 columns: int = COLUMNS) -> Generator[IsolationEnvironment,
234 None, None]:
235 with _fifos('stdin' if stdin_pipe else None,
236 'stdout' if stdout_pipe else None,
237 'stderr' if stderr_pipe else None) as fifo_paths:
238 result_r, result_w = os.pipe()
239 env_pid, tty = pty.fork()
240 if env_pid == pty.CHILD:
241 try:
242 os.close(result_r)
243 # Get ttyname from original stdout, even if test runner has
244 # replaced stdout with another file
245 pts = os.ttyname(pty.STDOUT_FILENO)
246
247 ctx = multiprocessing.get_context('spawn')
248 p = ctx.Process(target=_run_test, args=(child_function, pts,
249 *fifo_paths))
250 p.start()
251
252 def handle_terminate(signum: int,
253 frame: Optional[types.FrameType]) -> None:
254 if p.is_alive():
255 p.terminate()
256
257 signal.signal(signal.SIGTERM, handle_terminate)
258 # Signals from the pty are for the test process, not us
259 signal.signal(signal.SIGINT, signal.SIG_IGN)
260
261 p.join(2) # Wait at most 2s
262 if p.exitcode is None:
263 raise TimeoutException(p.pid)
264 except BaseException:
265 try:
266 with os.fdopen(result_w, 'w') as result_writer:
267 traceback.print_exc(file=result_writer)
268 finally:
269 # Prevent blocking in parent process by opening our end of
270 # all FIFOs.
271 for path, write in zip(fifo_paths, [False, True, True]):
272 _open_fifo(path, write)
273 os._exit(1)
274 with os.fdopen(result_w, 'w') as result_writer:
275 result_writer.write(f'{p.exitcode}\n')
276
277 os._exit(0)
278 else:
279 os.close(result_w)
280 fcntl.ioctl(tty, termios.TIOCSWINSZ,
281 struct.pack('HHHH', lines, columns, 0, 0))
282 env = IsolationEnvironment(env_pid, tty, *fifo_paths)
283 time.sleep(0.01)
284 try:
285 try:
286 yield env
287 finally:
288 def get_return_code() -> int:
289 with os.fdopen(result_r) as result_reader:
290 result = result_reader.readline().rstrip()
291 try:
292 return int(result)
293 except ValueError:
294 pass
295 trace = result_reader.read()
296 raise TestProcessNotComplete('\n'.join([result,
297 trace]))
298
299 env.close(get_return_code)
300 except Exception:
301 raw, processed = env.recorded_output()
302 if raw:
303 print(f'Raw output: {repr(raw)}', file=sys.stderr)
304 if processed:
305 print(f'Recorded output: {repr(processed)}',
306 file=sys.stderr)
307 raise
308
309
310def _run_test(test_function: typing.Callable[[], int],
311 pts: str,
312 stdin_fifo: Optional[str],
313 stdout_fifo: Optional[str],
314 stderr_fifo: Optional[str]) -> typing.NoReturn:
315 os.environ['TERM'] = 'xterm-256color'
316
317 tty_fd = os.open(pts, os.O_RDWR)
318 if stdin_fifo is not None:
319 old_stdin, sys.stdin = sys.stdin, open(stdin_fifo)
320 old_stdin.close()
321 else:
322 os.dup2(tty_fd, pty.STDIN_FILENO)
323 if stdout_fifo is not None:
324 old_stdout, sys.stdout = sys.stdout, open(stdout_fifo, 'w')
325 old_stdout.close()
326 else:
327 os.dup2(tty_fd, pty.STDOUT_FILENO)
328 if stderr_fifo is not None:
329 old_stderr, sys.stderr = sys.stderr, open(stderr_fifo, 'w',
330 buffering=1)
331 old_stderr.close()
332 else:
333 os.dup2(tty_fd, pty.STDERR_FILENO)
334 os.close(tty_fd)
335
336 result = test_function()
337 sys.exit(result or 0)
338