Linux kernel mirror (for testing)
git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git
tjh.dev
kernel
os
linux
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 = []
145
146
147class defer:
148 def __init__(self, func, *args, **kwargs):
149 if not callable(func):
150 raise Exception("defer created with un-callable object, did you call the function instead of passing its name?")
151
152 self.func = func
153 self.args = args
154 self.kwargs = kwargs
155
156 self._queue = global_defer_queue
157 self._queue.append(self)
158
159 def __enter__(self):
160 return self
161
162 def __exit__(self, ex_type, ex_value, ex_tb):
163 return self.exec()
164
165 def exec_only(self):
166 self.func(*self.args, **self.kwargs)
167
168 def cancel(self):
169 self._queue.remove(self)
170
171 def exec(self):
172 self.cancel()
173 self.exec_only()
174
175
176def tool(name, args, json=None, ns=None, host=None):
177 cmd_str = name + ' '
178 if json:
179 cmd_str += '--json '
180 cmd_str += args
181 cmd_obj = cmd(cmd_str, ns=ns, host=host)
182 if json:
183 return _json.loads(cmd_obj.stdout)
184 return cmd_obj
185
186
187def bpftool(args, json=None, ns=None, host=None):
188 return tool('bpftool', args, json=json, ns=ns, host=host)
189
190
191def ip(args, json=None, ns=None, host=None):
192 if ns:
193 args = f'-netns {ns} ' + args
194 return tool('ip', args, json=json, host=host)
195
196
197def ethtool(args, json=None, ns=None, host=None):
198 return tool('ethtool', args, json=json, ns=ns, host=host)
199
200
201def bpftrace(expr, json=None, ns=None, host=None, timeout=None):
202 """
203 Run bpftrace and return map data (if json=True).
204 The output of bpftrace is inconvenient, so the helper converts
205 to a dict indexed by map name, e.g.:
206 {
207 "@": { ... },
208 "@map2": { ... },
209 }
210 """
211 cmd_arr = ['bpftrace']
212 # Throw in --quiet if json, otherwise the output has two objects
213 if json:
214 cmd_arr += ['-f', 'json', '-q']
215 if timeout:
216 expr += ' interval:s:' + str(timeout) + ' { exit(); }'
217 cmd_arr += ['-e', expr]
218 cmd_obj = cmd(cmd_arr, ns=ns, host=host, shell=False)
219 if json:
220 # bpftrace prints objects as lines
221 ret = {}
222 for l in cmd_obj.stdout.split('\n'):
223 if not l.strip():
224 continue
225 one = _json.loads(l)
226 if one.get('type') != 'map':
227 continue
228 for k, v in one["data"].items():
229 if k.startswith('@'):
230 k = k.lstrip('@')
231 ret[k] = v
232 return ret
233 return cmd_obj
234
235
236def rand_port(stype=socket.SOCK_STREAM):
237 """
238 Get a random unprivileged port.
239 """
240 with socket.socket(socket.AF_INET6, stype) as s:
241 s.bind(("", 0))
242 return s.getsockname()[1]
243
244
245def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5):
246 end = time.monotonic() + deadline
247
248 pattern = f":{port:04X} .* "
249 if proto == "tcp": # for tcp protocol additionally check the socket state
250 pattern += "0A"
251 pattern = re.compile(pattern)
252
253 while True:
254 data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout
255 for row in data.split("\n"):
256 if pattern.search(row):
257 return
258 if time.monotonic() > end:
259 raise Exception("Waiting for port listen timed out")
260 time.sleep(sleep)
261
262
263def wait_file(fname, test_fn, sleep=0.005, deadline=5, encoding='utf-8'):
264 """
265 Wait for file contents on the local system to satisfy a condition.
266 test_fn() should take one argument (file contents) and return whether
267 condition is met.
268 """
269 end = time.monotonic() + deadline
270
271 with open(fname, "r", encoding=encoding) as fp:
272 while True:
273 if test_fn(fp.read()):
274 break
275 fp.seek(0)
276 if time.monotonic() > end:
277 raise TimeoutError("Wait for file contents failed", fname)
278 time.sleep(sleep)