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