| # SPDX-License-Identifier: GPL-2.0 | |
| import json as _json | |
| import os | |
| import re | |
| import select | |
| import socket | |
| import subprocess | |
| import time | |
| class CmdInitFailure(Exception): | |
| """ Command failed to start. Only raised by bkg(). """ | |
| def __init__(self, msg, cmd_obj): | |
| super().__init__(msg + "\n" + repr(cmd_obj)) | |
| self.cmd = cmd_obj | |
| class CmdExitFailure(Exception): | |
| """ Command failed (returned non-zero exit code). """ | |
| def __init__(self, msg, cmd_obj): | |
| super().__init__(msg + "\n" + repr(cmd_obj)) | |
| self.cmd = cmd_obj | |
| class CmdExitZeroFailure(CmdExitFailure): | |
| """ Command succeeded (returned zero exit code), but expected failure. """ | |
| def fd_read_timeout(fd, timeout): | |
| rlist, _, _ = select.select([fd], [], [], timeout) | |
| if rlist: | |
| return os.read(fd, 1024) | |
| raise TimeoutError("Timeout waiting for fd read") | |
| class cmd: | |
| """ | |
| Execute a command on local or remote host. | |
| @shell defaults to false, and class will try to split @comm into a list | |
| if it's a string with spaces. | |
| Use bkg() instead to run a command in the background. | |
| """ | |
| def __init__(self, comm, shell=None, fail=True, expect_fail=False, ns=None, | |
| background=False, host=None, timeout=20, ksft_ready=None, | |
| ksft_wait=None): | |
| if ns: | |
| if hasattr(ns, 'user_ns_path'): | |
| comm = (f'nsenter --user={ns.user_ns_path} ' | |
| f'--net={ns.net_ns_path} --setuid=0 --setgid=0 -- ' | |
| + comm) | |
| else: | |
| comm = f'ip netns exec {ns} ' + comm | |
| self.stdout = None | |
| self.stderr = None | |
| self.ret = None | |
| self.ksft_term_fd = None | |
| self.host = host | |
| self.comm = comm | |
| if host: | |
| self.proc = host.cmd(comm) | |
| else: | |
| # If user doesn't explicitly request shell try to avoid it. | |
| if shell is None and isinstance(comm, str) and ' ' in comm: | |
| comm = comm.split() | |
| # ksft_wait lets us wait for the background process to fully start, | |
| # we pass an FD to the child process, and wait for it to write back. | |
| # Similarly term_fd tells child it's time to exit. | |
| pass_fds = [] | |
| env = os.environ.copy() | |
| if ksft_wait is not None: | |
| wait_fd, self.ksft_term_fd = os.pipe() | |
| pass_fds.append(wait_fd) | |
| env["KSFT_WAIT_FD"] = str(wait_fd) | |
| ksft_ready = True # ksft_wait implies ready | |
| if ksft_ready is not None: | |
| rfd, ready_fd = os.pipe() | |
| pass_fds.append(ready_fd) | |
| env["KSFT_READY_FD"] = str(ready_fd) | |
| self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, pass_fds=pass_fds, | |
| env=env) | |
| if ksft_wait is not None: | |
| os.close(wait_fd) | |
| if ksft_ready is not None: | |
| os.close(ready_fd) | |
| msg = fd_read_timeout(rfd, ksft_wait) | |
| os.close(rfd) | |
| if not msg: | |
| terminate = self.proc.poll() is None | |
| self._process_terminate(terminate=terminate, timeout=1) | |
| raise CmdInitFailure("Did not receive ready message", self) | |
| if not background: | |
| self.process(terminate=False, fail=fail, expect_fail=expect_fail, | |
| timeout=timeout) | |
| def _process_terminate(self, terminate, timeout): | |
| if terminate: | |
| self.proc.terminate() | |
| stdout, stderr = self.proc.communicate(timeout=timeout) | |
| self.stdout = stdout.decode("utf-8") | |
| self.stderr = stderr.decode("utf-8") | |
| self.proc.stdout.close() | |
| self.proc.stderr.close() | |
| self.ret = self.proc.returncode | |
| return stdout, stderr | |
| def process(self, terminate=True, fail=None, expect_fail=False, timeout=20): | |
| if fail is None: | |
| fail = not terminate | |
| if self.ksft_term_fd: | |
| os.write(self.ksft_term_fd, b"1") | |
| stdout, stderr = self._process_terminate(terminate=terminate, | |
| timeout=timeout) | |
| # Fail on unexpected test failure if fail. | |
| # Fail on unexpected test success if expect_fail. | |
| # Fail on negative returncode if either: | |
| # Set by subprocess on crash or signal, this is never expected failure. | |
| if (self.proc.returncode != 0 and fail or | |
| (self.proc.returncode < 0 and expect_fail)): | |
| if len(stderr) > 0 and stderr[-1] == "\n": | |
| stderr = stderr[:-1] | |
| raise CmdExitFailure("Command failed", self) | |
| elif self.proc.returncode == 0 and expect_fail: | |
| raise CmdExitZeroFailure("Command succeeded (expected fail)", self) | |
| def __repr__(self): | |
| def str_fmt(name, s): | |
| name += ': ' | |
| return (name + s.strip().replace('\n', '\n' + ' ' * len(name))) | |
| ret = "CMD" | |
| if self.host: | |
| ret += "[remote]" | |
| if self.ret is None: | |
| ret += f" (unterminated): {self.comm}\n" | |
| elif self.ret == 0: | |
| ret += f" (success): {self.comm}\n" | |
| else: | |
| ret += f": {self.comm}\n" | |
| ret += f" EXIT: {self.ret}\n" | |
| if self.stdout: | |
| ret += str_fmt(" STDOUT", self.stdout) + "\n" | |
| if self.stderr: | |
| ret += str_fmt(" STDERR", self.stderr) + "\n" | |
| return ret.strip() | |
| class bkg(cmd): | |
| """ | |
| Run a command in the background. | |
| Examples usage: | |
| Run a command on remote host, and wait for it to finish. | |
| This is usually paired with wait_port_listen() to make sure | |
| the command has initialized: | |
| with bkg("socat ...", exit_wait=True, host=cfg.remote) as nc: | |
| ... | |
| Run a command and expect it to let us know that it's ready | |
| by writing to a special file descriptor passed via KSFT_READY_FD. | |
| Command will be terminated when we exit the context manager: | |
| with bkg("my_binary", ksft_wait=5): | |
| """ | |
| def __init__(self, comm, shell=None, fail=None, expect_fail=None, | |
| ns=None, host=None, exit_wait=False, ksft_ready=None, | |
| ksft_wait=None): | |
| super().__init__(comm, background=True, | |
| shell=shell, fail=fail, expect_fail=expect_fail, | |
| ns=ns, host=host, ksft_ready=ksft_ready, | |
| ksft_wait=ksft_wait) | |
| self.terminate = not exit_wait and not ksft_wait | |
| self._exit_wait = exit_wait | |
| self.check_fail = fail | |
| self.expect_fail = expect_fail | |
| if shell and self.terminate: | |
| print("# Warning: combining shell and terminate is risky!") | |
| print("# SIGTERM may not reach the child on zsh/ksh!") | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, ex_type, ex_value, ex_tb): | |
| terminate = self.terminate | |
| # Force termination on exception, but only if bkg() didn't already exit | |
| # since forcing termination silences failures with fail=None | |
| if self.proc.poll() is None: | |
| terminate = terminate or (self._exit_wait and ex_type is not None) | |
| return self.process(terminate=terminate, fail=self.check_fail, | |
| expect_fail=self.expect_fail) | |
| GLOBAL_DEFER_QUEUE = [] | |
| GLOBAL_DEFER_ARMED = False | |
| class defer: | |
| def __init__(self, func, *args, **kwargs): | |
| if not callable(func): | |
| raise Exception("defer created with un-callable object, did you call the function instead of passing its name?") | |
| self.func = func | |
| self.args = args | |
| self.kwargs = kwargs | |
| if not GLOBAL_DEFER_ARMED: | |
| raise Exception("defer queue not armed, did you use defer() outside of a test case?") | |
| self._queue = GLOBAL_DEFER_QUEUE | |
| self._queue.append(self) | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, ex_type, ex_value, ex_tb): | |
| return self.exec() | |
| def exec_only(self): | |
| self.func(*self.args, **self.kwargs) | |
| def cancel(self): | |
| self._queue.remove(self) | |
| def exec(self): | |
| self.cancel() | |
| self.exec_only() | |
| def tool(name, args, json=None, ns=None, host=None): | |
| cmd_str = name + ' ' | |
| if json: | |
| if name == 'tc': | |
| cmd_str += '-json ' | |
| else: | |
| cmd_str += '--json ' | |
| cmd_str += args | |
| cmd_obj = cmd(cmd_str, ns=ns, host=host) | |
| if json: | |
| return _json.loads(cmd_obj.stdout) | |
| return cmd_obj | |
| def bpftool(args, json=None, ns=None, host=None): | |
| return tool('bpftool', args, json=json, ns=ns, host=host) | |
| def ip(args, json=None, ns=None, host=None): | |
| if ns: | |
| args = f'-netns {ns} ' + args | |
| return tool('ip', args, json=json, host=host) | |
| def tc(args, json=None, ns=None, host=None): | |
| """ Helper to call tc with standard set of optional args. """ | |
| if ns: | |
| args = f'-netns {ns} ' + args | |
| return tool('tc', args, json=json, host=host) | |
| def ethtool(args, json=None, ns=None, host=None): | |
| return tool('ethtool', args, json=json, ns=ns, host=host) | |
| def bpftrace(expr, json=None, ns=None, host=None, timeout=None): | |
| """ | |
| Run bpftrace and return map data (if json=True). | |
| The output of bpftrace is inconvenient, so the helper converts | |
| to a dict indexed by map name, e.g.: | |
| { | |
| "@": { ... }, | |
| "@map2": { ... }, | |
| } | |
| """ | |
| cmd_arr = ['bpftrace'] | |
| # Throw in --quiet if json, otherwise the output has two objects | |
| if json: | |
| cmd_arr += ['-f', 'json', '-q'] | |
| if timeout: | |
| expr += ' interval:s:' + str(timeout) + ' { exit(); }' | |
| timeout += 20 | |
| cmd_arr += ['-e', expr] | |
| cmd_obj = cmd(cmd_arr, ns=ns, host=host, shell=False, timeout=timeout) | |
| if json: | |
| # bpftrace prints objects as lines | |
| ret = {} | |
| for l in cmd_obj.stdout.split('\n'): | |
| if not l.strip(): | |
| continue | |
| one = _json.loads(l) | |
| if one.get('type') != 'map': | |
| continue | |
| for k, v in one["data"].items(): | |
| if k.startswith('@'): | |
| k = k.lstrip('@') | |
| ret[k] = v | |
| return ret | |
| return cmd_obj | |
| def rand_port(stype=socket.SOCK_STREAM): | |
| """ | |
| Get a random unprivileged port. | |
| """ | |
| return rand_ports(1, stype)[0] | |
| def rand_ports(count, stype=socket.SOCK_STREAM): | |
| """ | |
| Get a unique set of random unprivileged ports. | |
| """ | |
| sockets = [] | |
| ports = [] | |
| try: | |
| for _ in range(count): | |
| s = socket.socket(socket.AF_INET6, stype) | |
| sockets.append(s) | |
| s.bind(("", 0)) | |
| ports.append(s.getsockname()[1]) | |
| finally: | |
| for s in sockets: | |
| s.close() | |
| return ports | |
| def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5): | |
| end = time.monotonic() + deadline | |
| pattern = f":{port:04X} .* " | |
| if proto == "tcp": # for tcp protocol additionally check the socket state | |
| pattern += "0A" | |
| pattern = re.compile(pattern) | |
| while True: | |
| data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout | |
| for row in data.split("\n"): | |
| if pattern.search(row): | |
| return | |
| if time.monotonic() > end: | |
| raise Exception("Waiting for port listen timed out") | |
| time.sleep(sleep) | |
| def wait_file(fname, test_fn, sleep=0.005, deadline=5, encoding='utf-8'): | |
| """ | |
| Wait for file contents on the local system to satisfy a condition. | |
| test_fn() should take one argument (file contents) and return whether | |
| condition is met. | |
| """ | |
| end = time.monotonic() + deadline | |
| with open(fname, "r", encoding=encoding) as fp: | |
| while True: | |
| if test_fn(fp.read()): | |
| break | |
| fp.seek(0) | |
| if time.monotonic() > end: | |
| raise TimeoutError("Wait for file contents failed", fname) | |
| time.sleep(sleep) |