|
| 1 | +# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +"""Run the jailer under a screen session""" |
| 5 | + |
| 6 | +import os |
| 7 | +import re |
| 8 | +import select |
| 9 | +import signal |
| 10 | +import time |
| 11 | +from pathlib import Path |
| 12 | + |
| 13 | +import psutil |
| 14 | +from tenacity import Retrying, retry_if_exception_type, stop_after_attempt, wait_fixed |
| 15 | + |
| 16 | +from framework import utils |
| 17 | + |
| 18 | +from .jailer import JailerContext |
| 19 | + |
| 20 | +FLUSH_CMD = 'screen -S {session} -X colon "logfile flush 0^M"' |
| 21 | + |
| 22 | + |
| 23 | +def start_screen_process(screen_log, session_name, binary_path, binary_params): |
| 24 | + """Start binary process into a screen session.""" |
| 25 | + start_cmd = "screen -L -Logfile {logfile} -dmS {session} {binary} {params}" |
| 26 | + start_cmd = start_cmd.format( |
| 27 | + logfile=screen_log, |
| 28 | + session=session_name, |
| 29 | + binary=binary_path, |
| 30 | + params=" ".join(binary_params), |
| 31 | + ) |
| 32 | + |
| 33 | + utils.check_output(start_cmd) |
| 34 | + |
| 35 | + # Build a regex object to match (number).session_name |
| 36 | + regex_object = re.compile(r"([0-9]+)\.{}".format(session_name)) |
| 37 | + |
| 38 | + # Run 'screen -ls' in a retry loop, 30 times with a 1s delay between calls. |
| 39 | + # If the output of 'screen -ls' matches the regex object, it will return the |
| 40 | + # PID. Otherwise, a RuntimeError will be raised. |
| 41 | + for attempt in Retrying( |
| 42 | + retry=retry_if_exception_type(RuntimeError), |
| 43 | + stop=stop_after_attempt(30), |
| 44 | + wait=wait_fixed(1), |
| 45 | + reraise=True, |
| 46 | + ): |
| 47 | + with attempt: |
| 48 | + screen_pid = utils.search_output_from_cmd( |
| 49 | + cmd="screen -ls", find_regex=regex_object |
| 50 | + ).group(1) |
| 51 | + |
| 52 | + screen_pid = int(screen_pid) |
| 53 | + # Make sure the screen process launched successfully |
| 54 | + # As the parent process for the binary. |
| 55 | + screen_ps = psutil.Process(screen_pid) |
| 56 | + |
| 57 | + for attempt in Retrying( |
| 58 | + stop=stop_after_attempt(5), |
| 59 | + wait=wait_fixed(0.5), |
| 60 | + reraise=True, |
| 61 | + ): |
| 62 | + with attempt: |
| 63 | + assert screen_ps.is_running() |
| 64 | + |
| 65 | + # Configure screen to flush stdout to file. |
| 66 | + utils.check_output(FLUSH_CMD.format(session=session_name)) |
| 67 | + |
| 68 | + return screen_pid |
| 69 | + |
| 70 | + |
| 71 | +class JailerScreen(JailerContext): |
| 72 | + """Spawn Firecracker under screen""" |
| 73 | + |
| 74 | + def __init__(self, *args, **kwargs): |
| 75 | + super().__init__(*args, **kwargs) |
| 76 | + self.daemonize = False |
| 77 | + self.screen_pid = None |
| 78 | + self.expect_kill_by_signal = False |
| 79 | + |
| 80 | + def spawn(self, pre_cmd): |
| 81 | + """Spawn Firecracker under screen""" |
| 82 | + self.screen_pid = None |
| 83 | + cmd = pre_cmd or [] |
| 84 | + cmd += self.construct_param_list() |
| 85 | + # Run Firecracker under screen. This is used when we want to access |
| 86 | + # the serial console. The file will collect the output from |
| 87 | + # 'screen'ed Firecracker. |
| 88 | + self.screen_pid = start_screen_process( |
| 89 | + self.screen_log, |
| 90 | + self.screen_session, |
| 91 | + cmd[0], |
| 92 | + cmd[1:], |
| 93 | + ) |
| 94 | + |
| 95 | + # If `--new-pid-ns` is used, the Firecracker process will detach from |
| 96 | + # the screen and the screen process will exit. We do not want to |
| 97 | + # attempt to kill it in that case to avoid a race condition. |
| 98 | + if self.new_pid_ns: |
| 99 | + self.screen_pid = None |
| 100 | + |
| 101 | + def kill(self): |
| 102 | + """Kill the Firecracker process""" |
| 103 | + if not self.screen_pid: |
| 104 | + raise RuntimeError("screen process not started") |
| 105 | + # Killing screen will send SIGHUP to underlying Firecracker. |
| 106 | + # Needed to avoid false positives in case kill() is called again. |
| 107 | + self.expect_kill_by_signal = True |
| 108 | + os.kill(self.screen_pid, signal.SIGKILL) |
| 109 | + os.kill(self.pid, signal.SIGKILL) |
| 110 | + |
| 111 | + @property |
| 112 | + def console_data(self): |
| 113 | + """Return the output of microVM's console""" |
| 114 | + if self.screen_log is None: |
| 115 | + return None |
| 116 | + file = Path(self.screen_log) |
| 117 | + if not file.exists(): |
| 118 | + return None |
| 119 | + return file.read_text(encoding="utf-8") |
| 120 | + |
| 121 | + @property |
| 122 | + def screen_session(self): |
| 123 | + """The screen session name |
| 124 | +
|
| 125 | + The id of this microVM, which should be unique. |
| 126 | + """ |
| 127 | + return self.jailer_id |
| 128 | + |
| 129 | + @property |
| 130 | + def screen_log(self): |
| 131 | + """Get the screen log file.""" |
| 132 | + return f"/tmp/screen-{self.screen_session}.log" |
| 133 | + |
| 134 | + def serial_input(self, input_string): |
| 135 | + """Send a string to the Firecracker serial console via screen.""" |
| 136 | + input_cmd = f'screen -S {self.screen_session} -p 0 -X stuff "{input_string}"' |
| 137 | + return utils.check_output(input_cmd) |
| 138 | + |
| 139 | + def serial(self): |
| 140 | + """Get a Serial object for this jailer/microvm""" |
| 141 | + return Serial(self) |
| 142 | + |
| 143 | + |
| 144 | +class Serial: |
| 145 | + """Class for serial console communication with a Microvm.""" |
| 146 | + |
| 147 | + RX_TIMEOUT_S = 60 |
| 148 | + |
| 149 | + def __init__(self, screen_jailer): |
| 150 | + """Initialize a new Serial object.""" |
| 151 | + self._poller = None |
| 152 | + self._screen_jailer = screen_jailer |
| 153 | + |
| 154 | + def open(self): |
| 155 | + """Open a serial connection.""" |
| 156 | + # Open the screen log file. |
| 157 | + if self._poller is not None: |
| 158 | + # serial already opened |
| 159 | + return |
| 160 | + |
| 161 | + attempt = 0 |
| 162 | + while not Path(self._screen_jailer.screen_log).exists() and attempt < 5: |
| 163 | + time.sleep(0.2) |
| 164 | + attempt += 1 |
| 165 | + |
| 166 | + screen_log_fd = os.open(self._screen_jailer.screen_log, os.O_RDONLY) |
| 167 | + self._poller = select.poll() |
| 168 | + self._poller.register(screen_log_fd, select.POLLIN | select.POLLHUP) |
| 169 | + |
| 170 | + def tx(self, input_string, end="\n"): |
| 171 | + # pylint: disable=invalid-name |
| 172 | + # No need to have a snake_case naming style for a single word. |
| 173 | + r"""Send a string terminated by an end token (defaulting to "\n").""" |
| 174 | + self._screen_jailer.serial_input(input_string + end) |
| 175 | + |
| 176 | + def rx_char(self): |
| 177 | + """Read a single character.""" |
| 178 | + result = self._poller.poll(0.1) |
| 179 | + |
| 180 | + for fd, flag in result: |
| 181 | + if flag & select.POLLHUP: |
| 182 | + assert False, "Oh! The console vanished before test completed." |
| 183 | + |
| 184 | + if flag & select.POLLIN: |
| 185 | + output_char = str(os.read(fd, 1), encoding="utf-8", errors="ignore") |
| 186 | + return output_char |
| 187 | + |
| 188 | + return "" |
| 189 | + |
| 190 | + def rx(self, token="\n"): |
| 191 | + # pylint: disable=invalid-name |
| 192 | + # No need to have a snake_case naming style for a single word. |
| 193 | + r"""Read a string delimited by an end token (defaults to "\n").""" |
| 194 | + rx_str = "" |
| 195 | + start = time.time() |
| 196 | + while True: |
| 197 | + rx_str += self.rx_char() |
| 198 | + if rx_str.endswith(token): |
| 199 | + break |
| 200 | + if (time.time() - start) >= self.RX_TIMEOUT_S: |
| 201 | + self._screen_jailer.kill() |
| 202 | + assert False |
| 203 | + |
| 204 | + return rx_str |
0 commit comments