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 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.check_fail = fail
121
122 if shell and self.terminate:
123 print("# Warning: combining shell and terminate is risky!")
124 print("# SIGTERM may not reach the child on zsh/ksh!")
125
126 def __enter__(self):
127 return self
128
129 def __exit__(self, ex_type, ex_value, ex_tb):
130 return self.process(terminate=self.terminate, fail=self.check_fail)
131
132
133global_defer_queue = []
134
135
136class defer:
137 def __init__(self, func, *args, **kwargs):
138 global global_defer_queue
139
140 if not callable(func):
141 raise Exception("defer created with un-callable object, did you call the function instead of passing its name?")
142
143 self.func = func
144 self.args = args
145 self.kwargs = kwargs
146
147 self._queue = global_defer_queue
148 self._queue.append(self)
149
150 def __enter__(self):
151 return self
152
153 def __exit__(self, ex_type, ex_value, ex_tb):
154 return self.exec()
155
156 def exec_only(self):
157 self.func(*self.args, **self.kwargs)
158
159 def cancel(self):
160 self._queue.remove(self)
161
162 def exec(self):
163 self.cancel()
164 self.exec_only()
165
166
167def tool(name, args, json=None, ns=None, host=None):
168 cmd_str = name + ' '
169 if json:
170 cmd_str += '--json '
171 cmd_str += args
172 cmd_obj = cmd(cmd_str, ns=ns, host=host)
173 if json:
174 return _json.loads(cmd_obj.stdout)
175 return cmd_obj
176
177
178def ip(args, json=None, ns=None, host=None):
179 if ns:
180 args = f'-netns {ns} ' + args
181 return tool('ip', args, json=json, host=host)
182
183
184def ethtool(args, json=None, ns=None, host=None):
185 return tool('ethtool', args, json=json, ns=ns, host=host)
186
187
188def rand_port(type=socket.SOCK_STREAM):
189 """
190 Get a random unprivileged port.
191 """
192 with socket.socket(socket.AF_INET6, type) as s:
193 s.bind(("", 0))
194 return s.getsockname()[1]
195
196
197def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5):
198 end = time.monotonic() + deadline
199
200 pattern = f":{port:04X} .* "
201 if proto == "tcp": # for tcp protocol additionally check the socket state
202 pattern += "0A"
203 pattern = re.compile(pattern)
204
205 while True:
206 data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout
207 for row in data.split("\n"):
208 if pattern.search(row):
209 return
210 if time.monotonic() > end:
211 raise Exception("Waiting for port listen timed out")
212 time.sleep(sleep)