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