1#!/usr/bin/env python3
2import concurrent.futures
3import contextlib
4import fcntl
5import importlib
6import json
7import math
8import os
9from pathlib import Path
10import pwd
11import re
12import selectors
13import select
14import stat
15import signal
16import socket
17import struct
18import subprocess
19import threading
20import time
21import vms
22import tty
23
24dashboard_runs = importlib.import_module("dashboard-run")
25
26
27FIELDS = (
28 "name", "used", "usedbydataset", "usedbysnapshots", "referenced", "available",
29 "compressratio", "logicalused", "compression", "recordsize", "mountpoint",
30 "quota", "origin", "mounted",
31)
32TEXT_FIELDS = {"name", "compression", "mountpoint", "origin", "mounted"}
33MAX_REQUEST = 65536
34MAX_RESPONSE = 16 * 1024 * 1024
35STREAM_SLOTS = threading.BoundedSemaphore(4)
36GUEST_SLOTS = threading.BoundedSemaphore(4)
37VM_PREPARATION_SLOT = threading.BoundedSemaphore(1)
38INDEX_SNAPSHOT = r"index-[0-9]+(?:-[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12})?"
39
40
41class Rejected(Exception):
42 pass
43
44
45class Host:
46 def __init__(self, pool):
47 if not re.fullmatch(r"[a-zA-Z][a-zA-Z0-9_.:-]*", pool):
48 raise ValueError("incorrect storage pool")
49 self.pool = pool
50 self.containers = None
51 self.inventory_lock = threading.Lock()
52
53 def dataset(self, value):
54 if (not isinstance(value, str) or len(value) > 255
55 or not re.fullmatch(r"[\w.: ][\w.: -]*(/[\w.: -]+)*", value)
56 or any(part in {".", ".."} for part in value.split("/"))
57 or not (value == self.pool or value.startswith(self.pool + "/"))):
58 raise Rejected("Choose a dataset from this storage pool.")
59 return value
60
61 def snapshot(self, value):
62 if (not isinstance(value, str) or len(value) > 255
63 or not re.fullmatch(r"[\w.: ][\w.: -]*", value)
64 or value in {".", ".."}):
65 raise Rejected("Choose a snapshot from the list.")
66 return value
67
68 def handle(self, request):
69 if not isinstance(request, dict):
70 raise Rejected("Choose a supported host operation.")
71 operation = request.get("operation")
72 parameters = {
73 "storage.datasets": set(),
74 "storage.pool": set(),
75 "storage.mounts": set(),
76 "storage.snapshots": {"dataset"},
77 "storage.reclaim": {"dataset", "from", "to"},
78 "storage.destroy": {"dataset", "from", "to"},
79 "storage.removed": {"dataset", "snapshot"},
80 "files.snapshot": {"datasets"},
81 "files.snapshots": {"datasets"},
82 "files.discard": {"dataset", "snapshot"},
83 "index.snapshot": {"dataset"},
84 "index.snapshots": {"dataset"},
85 "index.diff": {"dataset", "from", "to"},
86 "index.discard": {"dataset", "snapshot"},
87 "host.sample": set(),
88 "host.usage": {"refresh"},
89 **dashboard_runs.FIELDS,
90 **{"vm." + action: set() if fields is None else {"payload"} for action, fields in vms.ACTION_FIELDS.items()},
91 }
92 if not isinstance(operation, str) or operation not in parameters:
93 raise Rejected("Choose a supported host operation.")
94 if set(request) != parameters[operation] | {"operation"}:
95 raise Rejected("Use only the fields required by this host operation.")
96 if operation in dashboard_runs.FIELDS:
97 return dashboard_runs.handle(request)
98 if operation == "host.sample":
99 stat = Path("/proc/stat").read_text().splitlines()
100 # Guest columns are already included in user and nice time.
101 cpu = list(map(int, stat[0].split()[1:9]))
102 total = sum(cpu)
103 memory = {line.split()[0].rstrip(":"): int(line.split()[1]) * 1024
104 for line in Path("/proc/meminfo").read_text().splitlines()}
105 arc = None
106 try:
107 arc = next(int(line.split()[-1]) for line in Path("/proc/spl/kstat/zfs/arcstats").read_text().splitlines() if line.split()[0] == "size")
108 except (OSError, StopIteration):
109 pass
110 temperature = None
111 for root in Path("/sys/class/hwmon").glob("hwmon*"):
112 try:
113 if root.joinpath("name").read_text().strip() in {"k10temp", "coretemp", "zenpower"}:
114 value = float(root.joinpath("temp1_input").read_text()) / 1000
115 if math.isfinite(value) and value > 0:
116 temperature = value
117 break
118 except (OSError, ValueError):
119 pass
120 rx, tx = 0, 0
121 for line in Path("/proc/net/dev").read_text().splitlines():
122 name, separator, values = line.partition(":")
123 if separator and name.strip().startswith(("en", "eth", "wl")):
124 values = values.split()
125 rx += int(values[0])
126 tx += int(values[8])
127 gpu = None
128 try:
129 values = [float(line) for line in command("nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits", timeout=3).splitlines()]
130 values = [value for value in values if math.isfinite(value)]
131 if values:
132 gpu = sum(values) / len(values)
133 except (OSError, ValueError, subprocess.SubprocessError, TimeoutError):
134 pass
135 return {"at": time.monotonic(), "cpu": {"busy": total - cpu[3] - cpu[4], "total": total},
136 "cores": sum(bool(re.match(r"cpu[0-9]+\b", line)) for line in stat),
137 "memory": {"used": memory["MemTotal"] - memory["MemAvailable"], "total": memory["MemTotal"]},
138 "bootedAt": time.time() - float(Path("/proc/uptime").read_text().split()[0]),
139 "load": os.getloadavg()[1], "arc": arc, "temperature": temperature, "gpu": gpu,
140 "network": {"rx": rx, "tx": tx}}
141 if operation == "host.usage":
142 if type(request["refresh"]) is not bool:
143 raise Rejected("Choose whether to refresh the container list.")
144 with self.inventory_lock:
145 if request["refresh"] or not self.containers:
146 self.containers = json.loads(command("podman", "--remote", "--url", "unix:///run/podman/podman.sock", "ps", "--format", "json", timeout=10))
147 rows = []
148 for container in self.containers:
149 pid, identity = container["Pid"], container["Id"]
150 if not isinstance(pid, int) or pid <= 0 or not re.fullmatch(r"[0-9a-f]{64}", identity):
151 continue
152 try:
153 cgroup = next(line[3:] for line in Path(f"/proc/{pid}/cgroup").read_text().splitlines() if line.startswith("0::"))
154 if f"libpod-{identity}.scope" not in Path(cgroup).parts:
155 raise ValueError("container cgroup changed")
156 root = Path("/sys/fs/cgroup") / cgroup.lstrip("/")
157 cpu = next(int(line.split()[1]) for line in root.joinpath("cpu.stat").read_text().splitlines() if line.startswith("usage_usec "))
158 memory = int(root.joinpath("memory.current").read_text())
159 rows.append({"id": identity, "name": container["Names"][0], "cpu": cpu, "memory": memory})
160 except (OSError, ValueError, StopIteration):
161 self.containers = None
162 return {"at": time.monotonic(), "containers": rows}
163 if operation.startswith("vm."):
164 action = operation[3:]
165 payload = request.get("payload")
166 try:
167 vms.validate(action, payload)
168 except ValueError as error:
169 raise Rejected(str(error)) from error
170 args = ["python3", str(Path(__file__).with_name("vms.py")), action]
171 if payload is not None:
172 args.append(json.dumps(payload))
173 try:
174 if action in {"create", "media", "preset"}:
175 if not VM_PREPARATION_SLOT.acquire(blocking=False):
176 raise Rejected("A VM image is being prepared. Wait for it to finish, then retry.")
177 try:
178 return json.loads(command(*args, timeout=900))
179 finally:
180 VM_PREPARATION_SLOT.release()
181 return json.loads(command(*args))
182 except subprocess.CalledProcessError as error:
183 if error.returncode == 2:
184 raise Rejected(error.stderr) from error
185 raise
186 if operation in {"files.snapshot", "files.snapshots"}:
187 values = request["datasets"]
188 if not isinstance(values, list) or not values:
189 raise Rejected("Choose at least one dataset.")
190 datasets = sorted({self.dataset(value) for value in values})
191 if operation == "files.snapshot":
192 name = f"dash-{time.time_ns()}"
193 command("zfs", "snapshot", *(f"{dataset}@{name}" for dataset in datasets))
194 return name
195 text = command("zfs", "list", "-H", "-t", "snapshot", "-o", "name", "-d", "1", *datasets)
196 return [name for name in text.splitlines() if re.fullmatch(r"dash-[0-9]+", name.partition("@")[2])]
197 dataset = self.dataset(request["dataset"]) if "dataset" in request else None
198 if operation.startswith("index."):
199 if operation == "index.snapshot":
200 snapshot = f"{dataset}@index-{time.time_ns()}"
201 command("zfs", "snapshot", snapshot)
202 return snapshot
203 if operation == "index.snapshots":
204 text = command("zfs", "list", "-H", "-t", "snapshot", "-o", "name", "-d", "1", dataset)
205 return [name for name in text.splitlines() if re.fullmatch(INDEX_SNAPSHOT, name.partition("@")[2])]
206 names = [request["from"], request["to"]] if operation == "index.diff" else [request["snapshot"]]
207 if any(not re.fullmatch(INDEX_SNAPSHOT, self.snapshot(name)) for name in names):
208 raise Rejected("Choose an index snapshot.")
209 snapshots = [f"{dataset}@{name}" for name in names]
210 if operation == "index.diff":
211 return command("zfs", "diff", "-FH", *snapshots)
212 command("zfs", "destroy", "-d", *snapshots)
213 return None
214 if operation == "storage.datasets":
215 text = command("zfs", "list", "-Hp", "-r", "-t", "filesystem",
216 "-o", ",".join(FIELDS), self.pool)
217 rows = []
218 for line in text.splitlines():
219 values = line.split("\t")
220 if len(values) != len(FIELDS):
221 raise ValueError("unexpected ZFS dataset response")
222 row = dict(zip(FIELDS, values))
223 for field in FIELDS:
224 if field not in TEXT_FIELDS:
225 row[field] = numeric(row[field])
226 if row.pop("mounted") != "yes" or row["mountpoint"] == "none":
227 row["mountpoint"] = None
228 if row["origin"] == "-":
229 row["origin"] = None
230 rows.append(row)
231 return rows
232 if operation == "storage.pool":
233 fields = ("size", "allocated", "free", "fragmentation")
234 summary = command("zpool", "list", "-Hp", "-o", ",".join(fields), self.pool)
235 status = json.loads(command("zpool", "status", "-jp", "--json-int", self.pool))
236 return {"summary": dict(zip(fields, map(numeric, summary.strip().split("\t")))),
237 "status": status, "name": self.pool}
238 if operation == "storage.mounts":
239 mounts = json.loads(command("findmnt", "-J", "-l", "-t", "zfs", "-o", "SOURCE,TARGET"))
240 return {"filesystems": [item for item in mounts.get("filesystems", [])
241 if "@" not in item["source"]
242 and (item["source"] == self.pool or item["source"].startswith(self.pool + "/"))]}
243 if operation == "storage.snapshots":
244 text = command("zfs", "list", "-Hp", "-t", "snapshot", "-d", "1", "-s", "creation",
245 "-o", "name,creation,used,referenced,clones", dataset)
246 rows = []
247 for line in text.splitlines():
248 name, creation, used, referenced, clones = line.split("\t")
249 rows.append({"name": name.split("@", 1)[1], "creation": numeric(creation),
250 "used": numeric(used), "referenced": numeric(referenced),
251 "clones": [] if clones == "-" else clones.split(",")})
252 return rows
253 if operation in {"storage.reclaim", "storage.destroy"}:
254 start, end = self.snapshot(request["from"]), self.snapshot(request["to"])
255 flags = ["-nvp"] if operation == "storage.reclaim" else []
256 return command("zfs", "destroy", *flags, f"{dataset}@{start}%{end}")
257 snapshot = self.snapshot(request["snapshot"])
258 if operation == "files.discard":
259 if not re.fullmatch(r"dash-[0-9]+", snapshot):
260 raise Rejected("Choose a file undo snapshot.")
261 command("zfs", "destroy", f"{dataset}@{snapshot}")
262 return None
263 return command("zfs", "diff", "-H", f"{dataset}@{snapshot}", dataset)
264
265
266def numeric(value):
267 try:
268 return float(value.rstrip("x"))
269 except ValueError:
270 return 0
271
272
273def command(*args, timeout=60):
274 with subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
275 stderr=subprocess.PIPE, start_new_session=True) as process:
276 output, errors = bytearray(), bytearray()
277 deadline = time.monotonic() + timeout
278 try:
279 with selectors.DefaultSelector() as selector:
280 selector.register(process.stdout, selectors.EVENT_READ, output)
281 selector.register(process.stderr, selectors.EVENT_READ, errors)
282 while selector.get_map():
283 remaining = deadline - time.monotonic()
284 if remaining <= 0:
285 raise TimeoutError("host command timed out")
286 for key, _ in selector.select(remaining):
287 chunk = os.read(key.fd, 65536)
288 if not chunk:
289 selector.unregister(key.fileobj)
290 continue
291 key.data.extend(chunk)
292 if len(output) + len(errors) > MAX_RESPONSE // 2:
293 raise RuntimeError("The host response is too large. Narrow the selection.")
294 process.wait(timeout=max(0.001, deadline - time.monotonic()))
295 except BaseException:
296 try:
297 os.killpg(process.pid, signal.SIGKILL)
298 except ProcessLookupError:
299 pass
300 raise
301 if process.returncode:
302 raise subprocess.CalledProcessError(process.returncode, args, stderr=errors.decode(errors="replace").strip())
303 return output.decode()
304
305
306def send_response(connection, response):
307 payload = json.dumps(response, allow_nan=False).encode()
308 if len(payload) > MAX_RESPONSE:
309 payload = b'{"error":"The host response is too large. Narrow the selection.","status":502}'
310 connection.sendall(struct.pack("!I", len(payload)) + payload)
311
312
313def stream_console(connection, host, request):
314 details = host.handle(request)
315 with contextlib.ExitStack() as resources:
316 if request["operation"] == "vm.serial":
317 descriptor = os.open(details["path"], os.O_RDWR | os.O_NOCTTY | os.O_NOFOLLOW | os.O_NONBLOCK)
318 screen = resources.enter_context(os.fdopen(descriptor, "r+b", buffering=0))
319 if not stat.S_ISCHR(os.fstat(descriptor).st_mode):
320 raise Rejected("This VM's serial console is unavailable. Restart the VM and reconnect.")
321 try:
322 fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
323 except BlockingIOError:
324 raise Rejected("This console is already open. Close its other console tab, then reconnect.") from None
325 tty.setraw(descriptor)
326 elif request["operation"] == "vm.guest":
327 screen = resources.enter_context(socket.socket(socket.AF_UNIX, socket.SOCK_STREAM))
328 screen.settimeout(10)
329 screen.connect(details["path"])
330 screen.setblocking(False)
331 else:
332 identity = details["uuid"]
333 deadline = time.monotonic() + 5
334 while True:
335 try:
336 screen = resources.enter_context(socket.create_connection(("127.0.0.1", details["port"]), timeout=10))
337 break
338 except ConnectionRefusedError:
339 if time.monotonic() >= deadline:
340 raise Rejected("This VM's screen isn't ready. Wait a moment, then reconnect.") from None
341 time.sleep(0.1)
342 details = host.handle(request)
343 if details["uuid"] != identity:
344 raise Rejected("The VM changed while its screen was opening. Open it again.")
345 screen.setblocking(False)
346 send_response(connection, {"value": None})
347 try:
348 connection.setblocking(False)
349 with selectors.DefaultSelector() as selector:
350 selector.register(connection, selectors.EVENT_READ, screen)
351 selector.register(screen, selectors.EVENT_READ, connection)
352 deadline = time.monotonic() + 8 * 3600
353 while time.monotonic() < deadline:
354 ready = selector.select(300)
355 if not ready:
356 continue
357 for key, _ in ready:
358 try:
359 data = os.read(key.fd, 65536)
360 except BlockingIOError:
361 continue
362 if not data:
363 return
364 remaining = memoryview(data)
365 stalled = time.monotonic() + 30
366 while remaining:
367 try:
368 written = os.write(key.data.fileno(), remaining)
369 if written == 0:
370 return
371 remaining = remaining[written:]
372 except BlockingIOError:
373 if not select.select([], [key.data], [], max(0, stalled - time.monotonic()))[1]:
374 return
375 except OSError:
376 return
377
378
379def stream_upload(connection, request):
380 if set(request) != {"operation", "payload"}:
381 raise Rejected("Use only the fields required to upload an ISO.")
382 spec = request["payload"]
383 try:
384 vms.validate("upload", spec)
385 except ValueError as error:
386 raise Rejected(str(error)) from error
387 directory = vms.UPLOADS
388 vms.managed_directory(directory)
389 target = directory / spec["volume"]
390 if target.exists() or target.is_symlink():
391 raise Rejected("An ISO already has this filename. Rename the file before uploading.")
392 import tempfile
393 fd, temporary = tempfile.mkstemp(prefix=".upload-", dir=directory)
394 try:
395 send_response(connection, {"value": None})
396 connection.settimeout(120)
397 deadline = time.monotonic() + 2 * 3600
398 with os.fdopen(fd, "wb") as outgoing:
399 remaining = spec["size"]
400 while remaining:
401 if time.monotonic() > deadline:
402 raise TimeoutError("ISO upload timed out. Try uploading again.")
403 data = connection.recv(min(1024 * 1024, remaining))
404 if not data:
405 raise Rejected("The ISO upload was interrupted. Upload the file again.")
406 outgoing.write(data)
407 remaining -= len(data)
408 outgoing.flush()
409 os.fsync(outgoing.fileno())
410 # Never overwrite a file that another upload published in the meantime.
411 os.link(temporary, target)
412 send_response(connection, {"value": {"volume": vms.image_id(target)}})
413 finally:
414 Path(temporary).unlink(missing_ok=True)
415
416
417def serve_connection(connection, host, allowed_uid):
418 with connection:
419 connection.settimeout(65)
420 _, uid, _ = struct.unpack("3i", connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
421 if uid != allowed_uid:
422 return
423 try:
424 # No buffered reader: bytes after the header belong to the stream.
425 line = bytearray()
426 while len(line) <= MAX_REQUEST:
427 byte = connection.recv(1)
428 if not byte:
429 break
430 line.extend(byte)
431 if byte == b"\n":
432 break
433 if len(line) > MAX_REQUEST or not line.endswith(b"\n"):
434 raise Rejected("The host request is too large or incomplete.")
435 request = json.loads(line)
436 if isinstance(request, dict) and request.get("operation") in {"vm.console", "vm.serial", "vm.guest", "vm.upload"}:
437 streams = GUEST_SLOTS if request["operation"] == "vm.guest" else STREAM_SLOTS
438 if not streams.acquire(blocking=False):
439 raise Rejected("Four VM connections or uploads are open. Close one and try again.")
440 try:
441 if request["operation"] in {"vm.console", "vm.serial", "vm.guest"}:
442 stream_console(connection, host, request)
443 else:
444 stream_upload(connection, request)
445 return
446 finally:
447 streams.release()
448 response = {"value": host.handle(request)}
449 except Rejected as error:
450 response = {"error": str(error), "status": 400}
451 except dashboard_runs.Error as error:
452 response = {"error": str(error), "status": error.status}
453 except (json.JSONDecodeError, UnicodeDecodeError):
454 response = {"error": "Send the host request as UTF-8 JSON.", "status": 400}
455 except Exception as error:
456 message = error.stderr if isinstance(error, subprocess.CalledProcessError) else str(error)
457 print(f"host operation: {message}", flush=True)
458 if isinstance(error, subprocess.CalledProcessError) and any(line.endswith(": snapshot has dependent clones") for line in message.splitlines()):
459 response = {"error": "A snapshot has clones. Delete its clones before deleting the snapshot.", "status": 409}
460 else:
461 response = {"error": "The host operation couldn't finish. Check its logs, then retry.", "status": 502}
462 try:
463 send_response(connection, response)
464 except (BrokenPipeError, ConnectionResetError):
465 pass
466
467
468def main():
469 host = Host(os.environ["STUDIO_POOL"])
470 allowed_uid = pwd.getpwnam(os.environ["STUDIO_DASHBOARD_USER"]).pw_uid
471 socket_path = Path(os.environ.get("STUDIO_HOST_SOCKET", "/run/studio-host/host.sock"))
472 socket_path.unlink(missing_ok=True)
473 with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as listener:
474 listener.bind(str(socket_path))
475 socket_path.chmod(0o660)
476 listener.listen(8)
477 if address := os.environ.get("NOTIFY_SOCKET"):
478 with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as notification:
479 notification.connect("\0" + address[1:] if address.startswith("@") else address)
480 notification.sendall(b"READY=1")
481 slots = threading.BoundedSemaphore(16)
482 with concurrent.futures.ThreadPoolExecutor(max_workers=12) as workers:
483 while True:
484 connection, _ = listener.accept()
485 if not slots.acquire(blocking=False):
486 connection.close()
487 continue
488 work = workers.submit(serve_connection, connection, host, allowed_uid)
489 work.add_done_callback(lambda _: slots.release())
490
491
492if __name__ == "__main__":
493 main()