1#!/usr/bin/env python3
2import concurrent.futures
3import importlib
4import json
5import math
6import os
7from pathlib import Path
8import pwd
9import re
10import selectors
11import signal
12import socket
13import struct
14import subprocess
15import threading
16import time
17import vms
18
19dashboard_runs = importlib.import_module("dashboard-run")
20
21
22FIELDS = (
23 "name", "used", "usedbydataset", "usedbysnapshots", "referenced", "available",
24 "compressratio", "logicalused", "compression", "recordsize", "mountpoint",
25 "quota", "origin", "mounted",
26)
27TEXT_FIELDS = {"name", "compression", "mountpoint", "origin", "mounted"}
28MAX_REQUEST = 65536
29MAX_RESPONSE = 16 * 1024 * 1024
30INDEX_SNAPSHOT = r"index-[0-9]+(?:-[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12})?"
31
32
33class Rejected(Exception):
34 pass
35
36
37class Host:
38 def __init__(self, pool):
39 if not re.fullmatch(r"[a-zA-Z][a-zA-Z0-9_.:-]*", pool):
40 raise ValueError("incorrect storage pool")
41 self.pool = pool
42 self.containers = None
43 self.inventory_lock = threading.Lock()
44
45 def dataset(self, value):
46 if (not isinstance(value, str) or len(value) > 255
47 or not re.fullmatch(r"[\w.: ][\w.: -]*(/[\w.: -]+)*", value)
48 or any(part in {".", ".."} for part in value.split("/"))
49 or not (value == self.pool or value.startswith(self.pool + "/"))):
50 raise Rejected("Choose a dataset from this storage pool.")
51 return value
52
53 def snapshot(self, value):
54 if (not isinstance(value, str) or len(value) > 255
55 or not re.fullmatch(r"[\w.: ][\w.: -]*", value)
56 or value in {".", ".."}):
57 raise Rejected("Choose a snapshot from the list.")
58 return value
59
60 def handle(self, request):
61 if not isinstance(request, dict):
62 raise Rejected("Choose a supported host operation.")
63 operation = request.get("operation")
64 parameters = {
65 "storage.datasets": set(),
66 "storage.pool": set(),
67 "storage.mounts": set(),
68 "storage.snapshots": {"dataset"},
69 "storage.reclaim": {"dataset", "from", "to"},
70 "storage.destroy": {"dataset", "from", "to"},
71 "storage.removed": {"dataset", "snapshot"},
72 "files.snapshot": {"datasets"},
73 "files.snapshots": {"datasets"},
74 "files.discard": {"dataset", "snapshot"},
75 "index.snapshot": {"dataset"},
76 "index.snapshots": {"dataset"},
77 "index.diff": {"dataset", "from", "to"},
78 "index.discard": {"dataset", "snapshot"},
79 "host.sample": set(),
80 "host.usage": {"refresh"},
81 **dashboard_runs.FIELDS,
82 **{"vm." + action: set() if fields is None else {"payload"} for action, fields in vms.ACTION_FIELDS.items()},
83 }
84 if not isinstance(operation, str) or operation not in parameters:
85 raise Rejected("Choose a supported host operation.")
86 if set(request) != parameters[operation] | {"operation"}:
87 raise Rejected("Use only the fields required by this host operation.")
88 if operation in dashboard_runs.FIELDS:
89 return dashboard_runs.handle(request)
90 if operation == "host.sample":
91 stat = Path("/proc/stat").read_text().splitlines()
92 # Guest columns are already included in user and nice time.
93 cpu = list(map(int, stat[0].split()[1:9]))
94 total = sum(cpu)
95 memory = {line.split()[0].rstrip(":"): int(line.split()[1]) * 1024
96 for line in Path("/proc/meminfo").read_text().splitlines()}
97 arc = None
98 try:
99 arc = next(int(line.split()[-1]) for line in Path("/proc/spl/kstat/zfs/arcstats").read_text().splitlines() if line.split()[0] == "size")
100 except (OSError, StopIteration):
101 pass
102 temperature = None
103 for root in Path("/sys/class/hwmon").glob("hwmon*"):
104 try:
105 if root.joinpath("name").read_text().strip() in {"k10temp", "coretemp", "zenpower"}:
106 value = float(root.joinpath("temp1_input").read_text()) / 1000
107 if math.isfinite(value) and value > 0:
108 temperature = value
109 break
110 except (OSError, ValueError):
111 pass
112 rx, tx = 0, 0
113 for line in Path("/proc/net/dev").read_text().splitlines():
114 name, separator, values = line.partition(":")
115 if separator and name.strip().startswith(("en", "eth", "wl")):
116 values = values.split()
117 rx += int(values[0])
118 tx += int(values[8])
119 gpu = None
120 try:
121 values = [float(line) for line in command("nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits", timeout=3).splitlines()]
122 values = [value for value in values if math.isfinite(value)]
123 if values:
124 gpu = sum(values) / len(values)
125 except (OSError, ValueError, subprocess.SubprocessError, TimeoutError):
126 pass
127 return {"at": time.monotonic(), "cpu": {"busy": total - cpu[3] - cpu[4], "total": total},
128 "cores": sum(bool(re.match(r"cpu[0-9]+\b", line)) for line in stat),
129 "memory": {"used": memory["MemTotal"] - memory["MemAvailable"], "total": memory["MemTotal"]},
130 "bootedAt": time.time() - float(Path("/proc/uptime").read_text().split()[0]),
131 "load": os.getloadavg()[1], "arc": arc, "temperature": temperature, "gpu": gpu,
132 "network": {"rx": rx, "tx": tx}}
133 if operation == "host.usage":
134 if type(request["refresh"]) is not bool:
135 raise Rejected("Choose whether to refresh the container list.")
136 with self.inventory_lock:
137 if request["refresh"] or not self.containers:
138 self.containers = json.loads(command("podman", "--remote", "--url", "unix:///run/podman/podman.sock", "ps", "--format", "json", timeout=10))
139 rows = []
140 for container in self.containers:
141 pid, identity = container["Pid"], container["Id"]
142 if not isinstance(pid, int) or pid <= 0 or not re.fullmatch(r"[0-9a-f]{64}", identity):
143 continue
144 try:
145 cgroup = next(line[3:] for line in Path(f"/proc/{pid}/cgroup").read_text().splitlines() if line.startswith("0::"))
146 if f"libpod-{identity}.scope" not in Path(cgroup).parts:
147 raise ValueError("container cgroup changed")
148 root = Path("/sys/fs/cgroup") / cgroup.lstrip("/")
149 cpu = next(int(line.split()[1]) for line in root.joinpath("cpu.stat").read_text().splitlines() if line.startswith("usage_usec "))
150 memory = int(root.joinpath("memory.current").read_text())
151 rows.append({"id": identity, "name": container["Names"][0], "cpu": cpu, "memory": memory})
152 except (OSError, ValueError, StopIteration):
153 self.containers = None
154 return {"at": time.monotonic(), "containers": rows}
155 if operation.startswith("vm."):
156 action = operation[3:]
157 payload = request.get("payload")
158 try:
159 vms.validate(action, payload)
160 except ValueError as error:
161 raise Rejected(str(error)) from error
162 args = ["python3", str(Path(__file__).with_name("vms.py")), action]
163 if payload is not None:
164 args.append(json.dumps(payload))
165 try:
166 return json.loads(command(*args))
167 except subprocess.CalledProcessError as error:
168 if error.returncode == 2:
169 raise Rejected(error.stderr) from error
170 raise
171 if operation in {"files.snapshot", "files.snapshots"}:
172 values = request["datasets"]
173 if not isinstance(values, list) or not values:
174 raise Rejected("Choose at least one dataset.")
175 datasets = sorted({self.dataset(value) for value in values})
176 if operation == "files.snapshot":
177 name = f"dash-{time.time_ns()}"
178 command("zfs", "snapshot", *(f"{dataset}@{name}" for dataset in datasets))
179 return name
180 text = command("zfs", "list", "-H", "-t", "snapshot", "-o", "name", "-d", "1", *datasets)
181 return [name for name in text.splitlines() if re.fullmatch(r"dash-[0-9]+", name.partition("@")[2])]
182 dataset = self.dataset(request["dataset"]) if "dataset" in request else None
183 if operation.startswith("index."):
184 if operation == "index.snapshot":
185 snapshot = f"{dataset}@index-{time.time_ns()}"
186 command("zfs", "snapshot", snapshot)
187 return snapshot
188 if operation == "index.snapshots":
189 text = command("zfs", "list", "-H", "-t", "snapshot", "-o", "name", "-d", "1", dataset)
190 return [name for name in text.splitlines() if re.fullmatch(INDEX_SNAPSHOT, name.partition("@")[2])]
191 names = [request["from"], request["to"]] if operation == "index.diff" else [request["snapshot"]]
192 if any(not re.fullmatch(INDEX_SNAPSHOT, self.snapshot(name)) for name in names):
193 raise Rejected("Choose an index snapshot.")
194 snapshots = [f"{dataset}@{name}" for name in names]
195 if operation == "index.diff":
196 return command("zfs", "diff", "-FH", *snapshots)
197 command("zfs", "destroy", "-d", *snapshots)
198 return None
199 if operation == "storage.datasets":
200 text = command("zfs", "list", "-Hp", "-r", "-t", "filesystem",
201 "-o", ",".join(FIELDS), self.pool)
202 rows = []
203 for line in text.splitlines():
204 values = line.split("\t")
205 if len(values) != len(FIELDS):
206 raise ValueError("unexpected ZFS dataset response")
207 row = dict(zip(FIELDS, values))
208 for field in FIELDS:
209 if field not in TEXT_FIELDS:
210 row[field] = numeric(row[field])
211 if row.pop("mounted") != "yes" or row["mountpoint"] == "none":
212 row["mountpoint"] = None
213 if row["origin"] == "-":
214 row["origin"] = None
215 rows.append(row)
216 return rows
217 if operation == "storage.pool":
218 fields = ("size", "allocated", "free", "fragmentation")
219 summary = command("zpool", "list", "-Hp", "-o", ",".join(fields), self.pool)
220 status = json.loads(command("zpool", "status", "-jp", "--json-int", self.pool))
221 return {"summary": dict(zip(fields, map(numeric, summary.strip().split("\t")))),
222 "status": status, "name": self.pool}
223 if operation == "storage.mounts":
224 mounts = json.loads(command("findmnt", "-J", "-l", "-t", "zfs", "-o", "SOURCE,TARGET"))
225 return {"filesystems": [item for item in mounts.get("filesystems", [])
226 if "@" not in item["source"]
227 and (item["source"] == self.pool or item["source"].startswith(self.pool + "/"))]}
228 if operation == "storage.snapshots":
229 text = command("zfs", "list", "-Hp", "-t", "snapshot", "-d", "1", "-s", "creation",
230 "-o", "name,creation,used,referenced,clones", dataset)
231 rows = []
232 for line in text.splitlines():
233 name, creation, used, referenced, clones = line.split("\t")
234 rows.append({"name": name.split("@", 1)[1], "creation": numeric(creation),
235 "used": numeric(used), "referenced": numeric(referenced),
236 "clones": [] if clones == "-" else clones.split(",")})
237 return rows
238 if operation in {"storage.reclaim", "storage.destroy"}:
239 start, end = self.snapshot(request["from"]), self.snapshot(request["to"])
240 flags = ["-nvp"] if operation == "storage.reclaim" else []
241 return command("zfs", "destroy", *flags, f"{dataset}@{start}%{end}")
242 snapshot = self.snapshot(request["snapshot"])
243 if operation == "files.discard":
244 if not re.fullmatch(r"dash-[0-9]+", snapshot):
245 raise Rejected("Choose a file undo snapshot.")
246 command("zfs", "destroy", f"{dataset}@{snapshot}")
247 return None
248 return command("zfs", "diff", "-H", f"{dataset}@{snapshot}", dataset)
249
250
251def numeric(value):
252 try:
253 return float(value.rstrip("x"))
254 except ValueError:
255 return 0
256
257
258def command(*args, timeout=60):
259 with subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
260 stderr=subprocess.PIPE, start_new_session=True) as process:
261 output, errors = bytearray(), bytearray()
262 deadline = time.monotonic() + timeout
263 try:
264 with selectors.DefaultSelector() as selector:
265 selector.register(process.stdout, selectors.EVENT_READ, output)
266 selector.register(process.stderr, selectors.EVENT_READ, errors)
267 while selector.get_map():
268 remaining = deadline - time.monotonic()
269 if remaining <= 0:
270 raise TimeoutError("host command timed out")
271 for key, _ in selector.select(remaining):
272 chunk = os.read(key.fd, 65536)
273 if not chunk:
274 selector.unregister(key.fileobj)
275 continue
276 key.data.extend(chunk)
277 if len(output) + len(errors) > MAX_RESPONSE // 2:
278 raise RuntimeError("The host response is too large. Narrow the selection.")
279 process.wait(timeout=max(0.001, deadline - time.monotonic()))
280 except BaseException:
281 try:
282 os.killpg(process.pid, signal.SIGKILL)
283 except ProcessLookupError:
284 pass
285 raise
286 if process.returncode:
287 raise subprocess.CalledProcessError(process.returncode, args, stderr=errors.decode(errors="replace").strip())
288 return output.decode()
289
290
291def serve_connection(connection, host, allowed_uid):
292 with connection:
293 connection.settimeout(65)
294 _, uid, _ = struct.unpack("3i", connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
295 if uid != allowed_uid:
296 return
297 try:
298 with connection.makefile("rb") as incoming:
299 line = incoming.readline(MAX_REQUEST + 1)
300 if len(line) > MAX_REQUEST or not line.endswith(b"\n"):
301 raise Rejected("The host request is too large or incomplete.")
302 response = {"value": host.handle(json.loads(line))}
303 except Rejected as error:
304 response = {"error": str(error), "status": 400}
305 except dashboard_runs.Error as error:
306 response = {"error": str(error), "status": error.status}
307 except (json.JSONDecodeError, UnicodeDecodeError):
308 response = {"error": "Send the host request as UTF-8 JSON.", "status": 400}
309 except Exception as error:
310 message = error.stderr if isinstance(error, subprocess.CalledProcessError) else str(error)
311 print(f"host operation: {message}", flush=True)
312 if isinstance(error, subprocess.CalledProcessError) and any(line.endswith(": snapshot has dependent clones") for line in message.splitlines()):
313 response = {"error": "A snapshot has clones. Delete its clones before deleting the snapshot.", "status": 409}
314 else:
315 response = {"error": "The host operation couldn't finish. Check its logs, then retry.", "status": 502}
316 payload = json.dumps(response, allow_nan=False).encode()
317 if len(payload) > MAX_RESPONSE:
318 payload = b'{"error":"The host response is too large. Narrow the selection.","status":502}'
319 connection.sendall(struct.pack("!I", len(payload)) + payload)
320
321
322def main():
323 host = Host(os.environ["STUDIO_POOL"])
324 allowed_uid = pwd.getpwnam(os.environ["STUDIO_DASHBOARD_USER"]).pw_uid
325 socket_path = Path(os.environ.get("STUDIO_HOST_SOCKET", "/run/studio-host/host.sock"))
326 socket_path.unlink(missing_ok=True)
327 with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as listener:
328 listener.bind(str(socket_path))
329 socket_path.chmod(0o660)
330 listener.listen(8)
331 if address := os.environ.get("NOTIFY_SOCKET"):
332 with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as notification:
333 notification.connect("\0" + address[1:] if address.startswith("@") else address)
334 notification.sendall(b"READY=1")
335 slots = threading.BoundedSemaphore(8)
336 with concurrent.futures.ThreadPoolExecutor(max_workers=4) as workers:
337 while True:
338 connection, _ = listener.accept()
339 if not slots.acquire(blocking=False):
340 connection.close()
341 continue
342 work = workers.submit(serve_connection, connection, host, allowed_uid)
343 work.add_done_callback(lambda _: slots.release())
344
345
346if __name__ == "__main__":
347 main()