1# SPDX-License-Identifier: GPL-2.0 2 3import json as _json 4import os 5import re 6import select 7import socket 8import subprocess 9import time 10 11 12class CmdExitFailure(Exception): 13 def __init__(self, msg, cmd_obj): 14 super().__init__(msg) 15 self.cmd = cmd_obj 16 17 18def fd_read_timeout(fd, timeout): 19 rlist, _, _ = select.select([fd], [], [], timeout) 20 if rlist: 21 return os.read(fd, 1024) 22 raise TimeoutError("Timeout waiting for fd read") 23 24 25class cmd: 26 """ 27 Execute a command on local or remote host. 28 29 @shell defaults to false, and class will try to split @comm into a list 30 if it's a string with spaces. 31 32 Use bkg() instead to run a command in the background. 33 """ 34 def __init__(self, comm, shell=None, fail=True, ns=None, background=False, 35 host=None, timeout=5, ksft_ready=None, ksft_wait=None): 36 if ns: 37 comm = f'ip netns exec {ns} ' + comm 38 39 self.stdout = None 40 self.stderr = None 41 self.ret = None 42 self.ksft_term_fd = None 43 44 self.comm = comm 45 if host: 46 self.proc = host.cmd(comm) 47 else: 48 # If user doesn't explicitly request shell try to avoid it. 49 if shell is None and isinstance(comm, str) and ' ' in comm: 50 comm = comm.split() 51 52 # ksft_wait lets us wait for the background process to fully start, 53 # we pass an FD to the child process, and wait for it to write back. 54 # Similarly term_fd tells child it's time to exit. 55 pass_fds = [] 56 env = os.environ.copy() 57 if ksft_wait is not None: 58 wait_fd, self.ksft_term_fd = os.pipe() 59 pass_fds.append(wait_fd) 60 env["KSFT_WAIT_FD"] = str(wait_fd) 61 ksft_ready = True # ksft_wait implies ready 62 if ksft_ready is not None: 63 rfd, ready_fd = os.pipe() 64 pass_fds.append(ready_fd) 65 env["KSFT_READY_FD"] = str(ready_fd) 66 67 self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE, 68 stderr=subprocess.PIPE, pass_fds=pass_fds, 69 env=env) 70 if ksft_wait is not None: 71 os.close(wait_fd) 72 if ksft_ready is not None: 73 os.close(ready_fd) 74 msg = fd_read_timeout(rfd, ksft_wait) 75 os.close(rfd) 76 if not msg: 77 raise Exception("Did not receive ready message") 78 if not background: 79 self.process(terminate=False, fail=fail, timeout=timeout) 80 81 def process(self, terminate=True, fail=None, timeout=5): 82 if fail is None: 83 fail = not terminate 84 85 if self.ksft_term_fd: 86 os.write(self.ksft_term_fd, b"1") 87 if terminate: 88 self.proc.terminate() 89 stdout, stderr = self.proc.communicate(timeout) 90 self.stdout = stdout.decode("utf-8") 91 self.stderr = stderr.decode("utf-8") 92 self.proc.stdout.close() 93 self.proc.stderr.close() 94 self.ret = self.proc.returncode 95 96 if self.proc.returncode != 0 and fail: 97 if len(stderr) > 0 and stderr[-1] == "\n": 98 stderr = stderr[:-1] 99 raise CmdExitFailure("Command failed: %s\nSTDOUT: %s\nSTDERR: %s" % 100 (self.proc.args, stdout, stderr), self) 101 102 103class bkg(cmd): 104 """ 105 Run a command in the background. 106 107 Examples usage: 108 109 Run a command on remote host, and wait for it to finish. 110 This is usually paired with wait_port_listen() to make sure 111 the command has initialized: 112 113 with bkg("socat ...", exit_wait=True, host=cfg.remote) as nc: 114 ... 115 116 Run a command and expect it to let us know that it's ready 117 by writing to a special file descriptor passed via KSFT_READY_FD. 118 Command will be terminated when we exit the context manager: 119 120 with bkg("my_binary", ksft_wait=5): 121 """ 122 def __init__(self, comm, shell=None, fail=None, ns=None, host=None, 123 exit_wait=False, ksft_ready=None, ksft_wait=None): 124 super().__init__(comm, background=True, 125 shell=shell, fail=fail, ns=ns, host=host, 126 ksft_ready=ksft_ready, ksft_wait=ksft_wait) 127 self.terminate = not exit_wait and not ksft_wait 128 self._exit_wait = exit_wait 129 self.check_fail = fail 130 131 if shell and self.terminate: 132 print("# Warning: combining shell and terminate is risky!") 133 print("# SIGTERM may not reach the child on zsh/ksh!") 134 135 def __enter__(self): 136 return self 137 138 def __exit__(self, ex_type, ex_value, ex_tb): 139 # Force termination on exception 140 terminate = self.terminate or (self._exit_wait and ex_type) 141 return self.process(terminate=terminate, fail=self.check_fail) 142 143 144GLOBAL_DEFER_QUEUE = [] 145GLOBAL_DEFER_ARMED = False 146 147 148class defer: 149 def __init__(self, func, *args, **kwargs): 150 if not callable(func): 151 raise Exception("defer created with un-callable object, did you call the function instead of passing its name?") 152 153 self.func = func 154 self.args = args 155 self.kwargs = kwargs 156 157 if not GLOBAL_DEFER_ARMED: 158 raise Exception("defer queue not armed, did you use defer() outside of a test case?") 159 self._queue = GLOBAL_DEFER_QUEUE 160 self._queue.append(self) 161 162 def __enter__(self): 163 return self 164 165 def __exit__(self, ex_type, ex_value, ex_tb): 166 return self.exec() 167 168 def exec_only(self): 169 self.func(*self.args, **self.kwargs) 170 171 def cancel(self): 172 self._queue.remove(self) 173 174 def exec(self): 175 self.cancel() 176 self.exec_only() 177 178 179def tool(name, args, json=None, ns=None, host=None): 180 cmd_str = name + ' ' 181 if json: 182 cmd_str += '--json ' 183 cmd_str += args 184 cmd_obj = cmd(cmd_str, ns=ns, host=host) 185 if json: 186 return _json.loads(cmd_obj.stdout) 187 return cmd_obj 188 189 190def bpftool(args, json=None, ns=None, host=None): 191 return tool('bpftool', args, json=json, ns=ns, host=host) 192 193 194def ip(args, json=None, ns=None, host=None): 195 if ns: 196 args = f'-netns {ns} ' + args 197 return tool('ip', args, json=json, host=host) 198 199 200def ethtool(args, json=None, ns=None, host=None): 201 return tool('ethtool', args, json=json, ns=ns, host=host) 202 203 204def bpftrace(expr, json=None, ns=None, host=None, timeout=None): 205 """ 206 Run bpftrace and return map data (if json=True). 207 The output of bpftrace is inconvenient, so the helper converts 208 to a dict indexed by map name, e.g.: 209 { 210 "@": { ... }, 211 "@map2": { ... }, 212 } 213 """ 214 cmd_arr = ['bpftrace'] 215 # Throw in --quiet if json, otherwise the output has two objects 216 if json: 217 cmd_arr += ['-f', 'json', '-q'] 218 if timeout: 219 expr += ' interval:s:' + str(timeout) + ' { exit(); }' 220 cmd_arr += ['-e', expr] 221 cmd_obj = cmd(cmd_arr, ns=ns, host=host, shell=False) 222 if json: 223 # bpftrace prints objects as lines 224 ret = {} 225 for l in cmd_obj.stdout.split('\n'): 226 if not l.strip(): 227 continue 228 one = _json.loads(l) 229 if one.get('type') != 'map': 230 continue 231 for k, v in one["data"].items(): 232 if k.startswith('@'): 233 k = k.lstrip('@') 234 ret[k] = v 235 return ret 236 return cmd_obj 237 238 239def rand_port(stype=socket.SOCK_STREAM): 240 """ 241 Get a random unprivileged port. 242 """ 243 with socket.socket(socket.AF_INET6, stype) as s: 244 s.bind(("", 0)) 245 return s.getsockname()[1] 246 247 248def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5): 249 end = time.monotonic() + deadline 250 251 pattern = f":{port:04X} .* " 252 if proto == "tcp": # for tcp protocol additionally check the socket state 253 pattern += "0A" 254 pattern = re.compile(pattern) 255 256 while True: 257 data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout 258 for row in data.split("\n"): 259 if pattern.search(row): 260 return 261 if time.monotonic() > end: 262 raise Exception("Waiting for port listen timed out") 263 time.sleep(sleep) 264 265 266def wait_file(fname, test_fn, sleep=0.005, deadline=5, encoding='utf-8'): 267 """ 268 Wait for file contents on the local system to satisfy a condition. 269 test_fn() should take one argument (file contents) and return whether 270 condition is met. 271 """ 272 end = time.monotonic() + deadline 273 274 with open(fname, "r", encoding=encoding) as fp: 275 while True: 276 if test_fn(fp.read()): 277 break 278 fp.seek(0) 279 if time.monotonic() > end: 280 raise TimeoutError("Wait for file contents failed", fname) 281 time.sleep(sleep) 282