1#!/usr/bin/env python3
2import asyncio
3import base64
4import contextlib
5from collections import deque
6import json
7import logging
8import os
9from pathlib import Path
10import signal
11import socket
12import struct
13import threading
14import time
15import zlib
16
17import gi
18
19gi.require_version("Gst", "1.0")
20gi.require_version("GstSdp", "1.0")
21gi.require_version("GstWebRTC", "1.0")
22gi.require_version("GstRtp", "1.0")
23from gi.repository import GLib, Gst, GstRtp, GstSdp, GstWebRTC
24
25
26MAX_PIXELS = 3840 * 2160
27STUN = os.environ.get("STUDIO_VM_STUN_SERVER", "stun://stun.cloudflare.com:3478")
28SLOTS = asyncio.Semaphore(4)
29GUEST_PROTOCOL = json.loads(Path(os.environ.get("STUDIO_VM_GUEST_PROTOCOL", Path(__file__).resolve().parent.parent / "guest/protocol.json")).read_text())
30GUESTS = {}
31TWCC = "http://www.ietf.org/id/draft-holmer-rmcat-transport-wide-cc-extensions-01"
32
33
34class WindowDesktop:
35 def __init__(self, width, height, windows, previous=None):
36 limits = GUEST_PROTOCOL["limits"]
37 if not 0 < width <= limits["dimension"] or not 0 < height <= limits["dimension"] or width * height > limits["pixels"] or not 0 < len(windows) <= limits["windows"]:
38 raise ValueError("Invalid guest window desktop")
39 self.windows = windows
40 self.surfaces = {}
41 self.width = max(width, *(window["width"] for window in windows))
42 x = y = row_height = 0
43 ids = set()
44 for window in sorted(windows, key=lambda window: window["id"]):
45 w, h = window["width"], window["height"]
46 if not window["id"] or window["id"] in ids or not 0 < w <= limits["dimension"] or not 0 < h <= limits["dimension"] or abs(window["x"]) > limits["dimension"] or abs(window["y"]) > limits["dimension"]:
47 raise ValueError("Invalid guest window bounds")
48 ids.add(window["id"])
49 if x + w > self.width:
50 x, y, row_height = 0, y + row_height, 0
51 window["atlas_x"], window["atlas_y"] = x, y
52 x, row_height = x + w, max(row_height, h)
53 surface = previous.surfaces.get(window["id"]) if previous else None
54 self.surfaces[window["id"]] = surface if surface and (surface[0], surface[1]) == (w, h) else (w, h, None)
55 self.height = y + row_height
56 if self.height > limits["dimension"] or self.width * self.height > limits["pixels"]:
57 raise OverflowError("Guest windows exceed the video atlas")
58 self.repacked = not previous or (self.width, self.height, [(w["id"], w["width"], w["height"]) for w in sorted(windows, key=lambda w: w["id"])]) != (
59 previous.width, previous.height, [(w["id"], w["width"], w["height"]) for w in sorted(previous.windows, key=lambda w: w["id"])])
60 self.pixels = bytearray(self.width * self.height * 4) if self.repacked else previous.pixels
61 for window in windows:
62 surface = self.surfaces[window["id"]]
63 if self.repacked and surface[2] is not None:
64 self.copy(window, 0, 0, window["width"], window["height"], surface[2])
65
66 @property
67 def complete(self):
68 return all(surface[2] is not None for surface in self.surfaces.values())
69
70 def copy(self, window, x, y, width, height, pixels):
71 for row in range(height):
72 offset = ((window["atlas_y"] + y + row) * self.width + window["atlas_x"] + x) * 4
73 self.pixels[offset:offset + width * 4] = pixels[row * width * 4:(row + 1) * width * 4]
74
75 def update(self, values, data):
76 window = next((window for window in self.windows if window["id"] == values["id"]), None)
77 if not window:
78 raise ValueError("Unknown guest window")
79 width, height = window["width"], window["height"]
80 x, y, w, h = values["x"], values["y"], values["rect_width"], values["rect_height"]
81 if (values["width"], values["height"]) != (width, height) or not w or not h or x + w > width or y + h > height or len(data) != w * h * 4:
82 raise ValueError("Invalid guest window rectangle")
83 pixels = self.surfaces[window["id"]][2]
84 if pixels is None:
85 if (x, y, w, h) != (0, 0, width, height):
86 raise ValueError("Guest window needs a complete first frame")
87 pixels = bytearray(data)
88 self.surfaces[window["id"]] = (width, height, pixels)
89 else:
90 for row in range(h):
91 offset = ((y + row) * width + x) * 4
92 pixels[offset:offset + w * 4] = data[row * w * 4:(row + 1) * w * 4]
93 self.copy(window, x, y, w, h, data)
94
95
96def png(width, height, pixels):
97 def chunk(kind, data):
98 return struct.pack("!I", len(data)) + kind + data + struct.pack("!I", zlib.crc32(kind + data))
99 rows = b"".join(b"\0" + pixels[row * width * 4:(row + 1) * width * 4] for row in range(height))
100 data = b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack("!IIBBBBB", width, height, 8, 6, 0, 0, 0)) + chunk(b"IDAT", zlib.compress(rows)) + chunk(b"IEND", b"")
101 return "data:image/png;base64," + base64.b64encode(data).decode()
102
103
104class GuestDesktop:
105 def __init__(self, name):
106 self.name = name
107 self.screens = set()
108 self.pixels = bytearray()
109 self.cursor = []
110 self.pointer = None
111 self.layers = None
112 self.writer = None
113 self.windows_available = None
114 self.stop_supported = None
115 self.stopping = False
116
117 def request(self, refresh=False):
118 if not self.stopping and self.writer and not self.writer.is_closing():
119 if refresh:
120 self.writer.write(struct.pack("!4sII", GUEST_PROTOCOL["magic"].encode(), GUEST_PROTOCOL["packets"]["refresh"]["id"], 0))
121 if self.windows_available is not None:
122 self.writer.write(struct.pack("!4sIII", GUEST_PROTOCOL["magic"].encode(), GUEST_PROTOCOL["packets"]["stream"]["id"], 4, bool(self.screens)))
123 if self.windows_available:
124 enabled = any(screen.experiment for screen in self.screens)
125 self.writer.write(struct.pack("!4sIII", GUEST_PROTOCOL["magic"].encode(), GUEST_PROTOCOL["packets"]["experiment"]["id"], 4, enabled))
126
127 def stop(self):
128 self.stopping = True
129 if self.writer and not self.writer.is_closing():
130 if self.stop_supported:
131 self.writer.write(struct.pack("!4sII", GUEST_PROTOCOL["magic"].encode(), GUEST_PROTOCOL["packets"]["stop"]["id"], 0))
132 elif self.stop_supported is False:
133 self.writer.close()
134
135 def sync(self, screen):
136 screen.send({"type": "capabilities", "windows": bool(self.windows_available)})
137 if self.pixels:
138 if self.pointer:
139 screen.send(self.pointer)
140 for cursor in self.cursor:
141 screen.send(cursor)
142
143 def suspend(self):
144 active = bool(self.pixels)
145 self.pixels = bytearray()
146 self.cursor.clear()
147 self.pointer = None
148 self.layers = None
149 for screen in tuple(self.screens):
150 screen.send({"type": "capabilities", "windows": False})
151 if active:
152 for screen in tuple(self.screens):
153 screen.resize()
154 screen.send(screen.desktop.cursor or {"type": "cursor", "source": "qemu", "frame": 0, "fallback": True, "x": 0, "y": 0})
155 screen.present()
156
157 async def read(self, reader):
158 limits = GUEST_PROTOCOL["limits"]
159 packets = {packet["id"]: (name, packet["fields"]) for name, packet in GUEST_PROTOCOL["packets"].items()}
160 while True:
161 async with asyncio.timeout(4):
162 magic, kind, length = struct.unpack("!4sII", await reader.readexactly(12))
163 if magic != GUEST_PROTOCOL["magic"].encode() or kind not in packets or length > limits["packet_bytes"]:
164 raise ValueError("Invalid guest display packet")
165 name, fields = packets[kind]
166 if name == "capabilities" and length == 4:
167 fields = fields[:1]
168 if length < len(fields) * 4:
169 raise ValueError("Incomplete guest display packet")
170 payload = await reader.readexactly(length)
171 values = dict(zip(fields, struct.unpack("!" + "I" * len(fields), payload[:len(fields) * 4])))
172 data = payload[len(fields) * 4:]
173 if self.stop_supported is None and name != "capabilities":
174 self.stop_supported = False
175 if self.stopping:
176 return
177 if name == "screen":
178 width, height = values["width"], values["height"]
179 x, y, w, h = values["x"], values["y"], values["rect_width"], values["rect_height"]
180 if not 0 < width <= limits["dimension"] or not 0 < height <= limits["dimension"] or width * height > limits["pixels"] or not w or not h or x + w > width or y + h > height or len(data) != w * h * 4:
181 raise ValueError("Invalid guest display rectangle")
182 resized = not self.pixels or (width, height) != (self.width, self.height)
183 if resized:
184 if (x, y, w, h) != (0, 0, width, height):
185 raise ValueError("Guest display needs a complete first frame")
186 self.width, self.height = width, height
187 self.pixels = bytearray(width * height * 4)
188 for row in range(h):
189 offset = ((y + row) * width + x) * 4
190 self.pixels[offset:offset + w * 4] = data[row * w * 4:(row + 1) * w * 4]
191 for screen in tuple(self.screens):
192 if resized:
193 screen.resize()
194 self.sync(screen)
195 if screen.source is self:
196 screen.present()
197 await asyncio.sleep(0)
198 elif name == "cursor":
199 width, height, count = values["width"], values["height"], values["frames"]
200 if not width and not height and not count and not data and not values["hot_x"] and not values["hot_y"]:
201 self.cursor = [{"type": "cursor", "source": "guest", "id": values["id"], "width": 0, "height": 0, "frame": 0, "frames": 0, "fallback": True, "x": 0, "y": 0}]
202 if self.pixels:
203 for screen in tuple(self.screens):
204 screen.send(self.cursor[0])
205 continue
206 pixels = width * height
207 frame_bytes = 4 + pixels * 5
208 if not 0 < width <= limits["cursor_size"] or not 0 < height <= limits["cursor_size"] or not 0 < count <= limits["cursor_frames"] or values["hot_x"] >= width or values["hot_y"] >= height or len(data) != frame_bytes * count:
209 raise ValueError("Invalid guest cursor")
210 self.cursor.clear()
211 for index in range(count):
212 frame = data[index * frame_bytes:(index + 1) * frame_bytes]
213 duration = struct.unpack("!I", frame[:4])[0]
214 if not 1 <= duration <= 60000:
215 raise ValueError("Invalid guest cursor duration")
216 mask = frame[4 + pixels * 4:]
217 if any(value & ~sum(GUEST_PROTOCOL["xor_bits"].values()) for value in mask):
218 raise ValueError("Invalid guest cursor inversion")
219 message = {"type": "cursor", "source": "guest", "id": values["id"], "width": width, "height": height,
220 "x": values["hot_x"], "y": values["hot_y"], "frame": index, "frames": count,
221 "duration": duration, "image": png(width, height, frame[4:4 + pixels * 4])}
222 if any(mask):
223 bits = GUEST_PROTOCOL["xor_bits"]
224 message["invert"] = png(width, height, b"".join(bytes([255 if value & bits["red"] else 0,
225 255 if value & bits["green"] else 0,
226 255 if value & bits["blue"] else 0,
227 255 if value else 0]) for value in mask))
228 self.cursor.append(message)
229 if self.pixels:
230 for screen in tuple(self.screens):
231 screen.send(message)
232 await asyncio.sleep(0)
233 elif name == "pointer":
234 if data or values["visible"] > 1:
235 raise ValueError("Invalid guest pointer")
236 self.pointer = {"type": "cursor-position", "id": values["id"], "visible": bool(values["visible"]),
237 "x": (values["x"] ^ 0x80000000) - 0x80000000, "y": (values["y"] ^ 0x80000000) - 0x80000000}
238 if self.pixels:
239 for screen in tuple(self.screens):
240 screen.send(self.pointer)
241 elif name == "suspend":
242 if data:
243 raise ValueError("Invalid guest display suspension")
244 self.suspend()
245 elif name == "heartbeat":
246 if data:
247 raise ValueError("Invalid guest display heartbeat")
248 elif name == "stop":
249 if data or not self.stopping:
250 raise ValueError("Unexpected guest display stop")
251 return
252 elif name == "capabilities":
253 if data or values["windows"] > 1 or values.get("stop", 0) > 1:
254 raise ValueError("Invalid guest display capabilities")
255 self.windows_available = bool(values["windows"])
256 self.stop_supported = bool(values.get("stop", 0))
257 for screen in tuple(self.screens):
258 screen.send({"type": "capabilities", "windows": bool(self.windows_available)})
259 if self.stopping:
260 self.stop()
261 else:
262 self.request()
263 elif name == "capture_stats":
264 if data or not 0 < values["elapsed_ms"] <= 60000 or values["frames"] > 10000 or any(values[key] > 60000000 for key in ("capture_us", "scan_us", "write_us")):
265 raise ValueError("Invalid guest capture statistics")
266 for screen in tuple(self.screens):
267 screen.send({"type": "capture-stats", **values})
268 elif name == "windows":
269 count = values["count"]
270 fields = GUEST_PROTOCOL["window_fields"]
271 if count > limits["windows"] or len(data) != count * len(fields) * 4:
272 raise ValueError("Invalid guest window list")
273 previous = self.layers
274 if count:
275 records = []
276 for record in struct.iter_unpack("!" + "I" * len(fields), data):
277 window = dict(zip(fields, record))
278 for key in ("x", "y", "caption_left", "caption_top", "caption_right", "caption_bottom"):
279 window[key] = (window[key] ^ 0x80000000) - 0x80000000
280 records.append(window)
281 try:
282 self.layers = WindowDesktop(values["width"], values["height"], records, previous)
283 except OverflowError:
284 self.layers = None
285 else:
286 self.layers = None
287 for screen in tuple(self.screens):
288 if screen.experiment:
289 if self.layers and not self.layers.repacked and self.layers.complete and screen.source is self.layers:
290 screen.send({"type": "window-positions", "windows": self.layers.windows})
291 else:
292 screen.resize()
293 screen.present()
294 screen.send({"type": "window-capture", "available": bool(self.layers)})
295 elif name == "surface":
296 if self.layers:
297 complete = self.layers.complete
298 self.layers.update(values, data)
299 for screen in tuple(self.screens):
300 if screen.experiment and self.layers.complete:
301 if not complete:
302 screen.resize()
303 screen.present()
304 else:
305 raise ValueError("Unexpected guest display request")
306
307 async def run(self):
308 while True:
309 if self.stopping:
310 GUESTS.pop(self.name, None)
311 return
312 writer = None
313 try:
314 reader, writer = await asyncio.wait_for(asyncio.open_unix_connection(
315 os.environ.get("STUDIO_HOST_SOCKET", "/run/studio-host/host.sock")), 10)
316 peer = writer.get_extra_info("socket")
317 if struct.unpack("3i", peer.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))[1] != 0:
318 raise ValueError("Untrusted guest display broker")
319 writer.write(json.dumps({"operation": "vm.guest", "payload": {"name": self.name}}).encode() + b"\n")
320 size = struct.unpack("!I", await asyncio.wait_for(reader.readexactly(4), 70))[0]
321 if size > 65536:
322 raise ValueError("Invalid guest display broker response")
323 response = json.loads(await reader.readexactly(size))
324 if "error" not in response:
325 self.writer = writer
326 writer.write(struct.pack("!4sII", GUEST_PROTOCOL["magic"].encode(), GUEST_PROTOCOL["packets"]["refresh"]["id"], 0))
327 await self.read(reader)
328 except (OSError, ValueError, EOFError, asyncio.IncompleteReadError, TimeoutError) as error:
329 logging.info("Guest display %s: %s", self.name, error)
330 finally:
331 self.writer = None
332 self.windows_available = None
333 self.stop_supported = None
334 self.suspend()
335 if writer:
336 writer.close()
337 with contextlib.suppress(OSError, TimeoutError):
338 await asyncio.wait_for(writer.wait_closed(), 2)
339 if self.stopping or not self.screens:
340 GUESTS.pop(self.name, None)
341 return
342 await asyncio.sleep(2)
343
344
345class Desktop:
346 def __init__(self, reader, writer, notify):
347 self.reader, self.writer, self.notify = reader, writer, notify
348 self.keys = set()
349 self.pointer = (0, 0)
350 self.buttons = 0
351 self.pixels = bytearray()
352 self.cursor = None
353
354 def resize(self, width, height):
355 if not 0 < width <= 4096 or not 0 < height <= 4096 or width * height > MAX_PIXELS:
356 raise ValueError("The VM screen exceeds the supported 4K display size.")
357 self.width, self.height = width, height
358 self.pixels = bytearray(width * height * 4)
359
360 async def start(self):
361 if not (await self.reader.readexactly(12)).startswith(b"RFB 003."):
362 raise ValueError("The VM screen did not respond. Restart the VM and reconnect.")
363 self.writer.write(b"RFB 003.008\n")
364 count = (await self.reader.readexactly(1))[0]
365 security = await self.reader.readexactly(count)
366 if 1 not in security:
367 raise ValueError("Remove the VM's VNC password in its display settings, then reconnect.")
368 self.writer.write(b"\x01")
369 if await self.reader.readexactly(4) != b"\0\0\0\0":
370 raise ValueError("The VM rejected the screen connection. Restart it and reconnect.")
371 self.writer.write(b"\x01")
372 width, height, _, length = struct.unpack("!HH16sI", await self.reader.readexactly(24))
373 if length > 65536:
374 raise ValueError("The VM screen sent an unsupported desktop name.")
375 await self.reader.readexactly(length)
376 self.resize(width, height)
377 self.writer.write(struct.pack("!B3xBBBBHHHBBB3x", 0, 32, 24, 0, 1, 255, 255, 255, 16, 8, 0))
378 encodings = (0, 1, -223, -224, -239)
379 self.writer.write(struct.pack("!BBH" + "i" * len(encodings), 2, 0, len(encodings), *encodings))
380 self.request(False)
381
382 def request(self, incremental=True):
383 self.writer.write(struct.pack("!BBHHHH", 3, incremental, 0, 0, self.width, self.height))
384
385 def input(self, message):
386 kind = message.get("type")
387 if kind == "key":
388 key, down = message.get("key"), message.get("down")
389 if type(key) is not int or not 0 <= key <= 0xFFFFFFFF or type(down) is not bool:
390 raise ValueError("Unsupported key event.")
391 if down:
392 self.keys.add(key)
393 else:
394 self.keys.discard(key)
395 self.writer.write(struct.pack("!BB2xI", 4, down, key))
396 elif kind == "pointer":
397 x, y, buttons = message.get("x"), message.get("y"), message.get("buttons")
398 if any(type(value) is not int for value in (x, y, buttons)) or not 0 <= buttons <= 255:
399 raise ValueError("Unsupported pointer event.")
400 self.pointer = (max(0, min(x, self.width - 1)), max(0, min(y, self.height - 1)))
401 self.buttons = buttons
402 self.writer.write(struct.pack("!BBHH", 5, buttons, *self.pointer))
403 elif kind == "release":
404 for key in self.keys:
405 self.writer.write(struct.pack("!BB2xI", 4, 0, key))
406 self.keys.clear()
407 self.buttons = 0
408 self.writer.write(struct.pack("!BBHH", 5, 0, *self.pointer))
409 elif kind == "clipboard":
410 text = message.get("text")
411 if not isinstance(text, str) or len(text) > 65536:
412 raise ValueError("Shorten the clipboard text, then send it again.")
413 try:
414 data = text.encode("latin1")
415 except UnicodeEncodeError:
416 raise ValueError("Some clipboard characters aren't supported by this guest.") from None
417 self.writer.write(struct.pack("!B3xI", 6, len(data)) + data)
418 else:
419 raise ValueError("Unsupported input event.")
420
421 async def read(self, present, resized):
422 while True:
423 kind = (await self.reader.readexactly(1))[0]
424 if kind == 0:
425 _, count = struct.unpack("!BH", await self.reader.readexactly(3))
426 dirty = False
427 for _ in range(count):
428 x, y, width, height, encoding = struct.unpack("!HHHHi", await self.reader.readexactly(12))
429 if encoding == -224:
430 break
431 if encoding == -223:
432 self.resize(width, height)
433 resized()
434 continue
435 if encoding == -239:
436 if width * height > 256 * 256:
437 raise ValueError("The VM sent an unsupported cursor size.")
438 pixels = bytearray(await self.reader.readexactly(width * height * 4))
439 mask = await self.reader.readexactly(((width + 7) // 8) * height)
440 for row in range(height):
441 for col in range(width):
442 offset = (row * width + col) * 4
443 pixels[offset], pixels[offset + 2] = pixels[offset + 2], pixels[offset]
444 pixels[offset + 3] = 255 if mask[row * ((width + 7) // 8) + col // 8] & (128 >> (col % 8)) else 0
445 self.cursor = {"type": "cursor", "source": "qemu", "width": width, "height": height, "x": x, "y": y,
446 "frame": 0, "frames": 1, "duration": 100,
447 "image": png(width, height, pixels) if width and height else None}
448 self.notify(self.cursor)
449 continue
450 if x + width > self.width or y + height > self.height:
451 raise ValueError("The VM sent a screen update outside its display.")
452 if encoding == 0:
453 data = await self.reader.readexactly(width * height * 4)
454 elif encoding == 1:
455 source_x, source_y = struct.unpack("!HH", await self.reader.readexactly(4))
456 if source_x + width > self.width or source_y + height > self.height:
457 raise ValueError("The VM sent a screen copy outside its display.")
458 data = b"".join(self.pixels[((source_y + row) * self.width + source_x) * 4:
459 ((source_y + row) * self.width + source_x + width) * 4] for row in range(height))
460 else:
461 raise ValueError("The VM sent an unsupported screen encoding.")
462 for row in range(height):
463 start = ((y + row) * self.width + x) * 4
464 self.pixels[start:start + width * 4] = data[row * width * 4:(row + 1) * width * 4]
465 dirty = dirty or bool(width and height)
466 if dirty:
467 present()
468 self.request()
469 await self.writer.drain()
470 elif kind == 2:
471 pass
472 elif kind == 3:
473 _, length = struct.unpack("!3sI", await self.reader.readexactly(7))
474 if length > 65536:
475 raise ValueError("The VM sent too much clipboard text.")
476 self.notify({"type": "clipboard", "text": (await self.reader.readexactly(length)).decode("latin1")})
477 else:
478 raise ValueError("The VM sent an unsupported screen message.")
479
480
481class Screen:
482 def __init__(self, writer, desktop):
483 self.loop = asyncio.get_running_loop()
484 self.writer, self.desktop = writer, desktop
485 self.guest = None
486 self.channel = None
487 self.pointer_channel = None
488 self.pointer_sequence = -1
489 self.experiment = False
490 self.frame_timer = None
491 self.last_frame = 0
492 self.layouts = {}
493 self.encoding_layouts = deque(maxlen=64)
494 self.last_rtp_pts = None
495 self.sent_layout = None
496 self.bwe = None
497 hardware = Gst.ElementFactory.make("nvh264enc")
498 encoder = "nvh264enc name=encoder preset=p4 tune=ultra-low-latency zerolatency=true bframes=0 rc-mode=cbr bitrate=4000 gop-size=60" if hardware else (
499 "x264enc name=encoder tune=zerolatency speed-preset=ultrafast bitrate=4000 key-int-max=60 bframes=0")
500 width, height = self.desktop.width, self.desktop.height
501 ssrc = struct.unpack("!I", os.urandom(4))[0] | 1
502 # TWCC and RTX headers are added after payload packetization.
503 self.pipeline = Gst.parse_launch(
504 f"appsrc name=frames is-live=true format=time max-buffers=1 leaky-type=downstream "
505 f"caps=video/x-raw,format=BGRx,width={width},height={height},framerate=60/1 "
506 f"! videoconvert ! videoscale ! capsfilter name=size caps=video/x-raw,format=NV12,width={width + width % 2},height={height + height % 2} "
507 f"! {encoder} ! video/x-h264,profile=constrained-baseline ! h264parse "
508 f"! rtph264pay name=pay pt=96 ssrc={ssrc} mtu=1190 config-interval=-1 aggregate-mode=zero-latency "
509 f"! capsfilter name=codec caps=application/x-rtp,media=video,encoding-name=H264,clock-rate=90000,ssrc=(uint){ssrc} "
510 "! webrtcbin name=rtc bundle-policy=max-bundle")
511 self.frames = self.pipeline.get_by_name("frames")
512 self.rtc = self.pipeline.get_by_name("rtc")
513 self.ice = self.rtc.get_property("ice-agent")
514 self.ice._ref_sink() # PyGObject sinks the default agent's floating reference, which webrtcbin still owns.
515 self.ice.set_property("min-rtp-port", 50000)
516 self.ice.set_property("max-rtp-port", 50031)
517 if STUN:
518 self.rtc.set_property("stun-server", STUN)
519 self.handlers = [
520 (self.rtc, self.rtc.connect("on-ice-candidate", lambda _, line, candidate: self.loop.call_soon_threadsafe(
521 self.send, {"type": "candidate", "sdpMLineIndex": line, "candidate": candidate}))),
522 (self.rtc, self.rtc.connect("on-data-channel", self.data_channel)),
523 (self.rtc, self.rtc.connect("notify::connection-state", self.connection_state)),
524 (self.rtc, self.rtc.connect("request-aux-sender", lambda *_: self.bwe)),
525 ]
526 self.encoder = self.pipeline.get_by_name("encoder")
527 self.encoder_probe = self.encoder.get_static_pad("sink").add_probe(Gst.PadProbeType.BUFFER, self.encoding)
528 self.rtp_probe = self.pipeline.get_by_name("pay").get_static_pad("src").add_probe(Gst.PadProbeType.BUFFER, self.rtp_layout)
529 self.probe = self.frames.get_static_pad("src").add_probe(Gst.PadProbeType.EVENT_UPSTREAM, self.refresh)
530 self.bus = self.pipeline.get_bus()
531 self.bus.add_signal_watch()
532 self.handlers.append((self.bus, self.bus.connect("message::error", self.pipeline_error)))
533 self.pipeline.set_state(Gst.State.PLAYING)
534 self.send({"type": "display", "width": width, "height": height, "encoder": "nvenc" if hardware else "x264"})
535
536 def close(self):
537 self.writer.close()
538 if self.frame_timer:
539 self.frame_timer.cancel()
540 self.desktop.notify = None
541 if self.guest:
542 self.guest.screens.discard(self)
543 if self.guest.screens:
544 self.guest.request()
545 else:
546 self.guest.stop()
547 for element, handler in self.handlers:
548 element.disconnect(handler)
549 self.handlers.clear()
550 self.frames.get_static_pad("src").remove_probe(self.probe)
551 self.encoder.get_static_pad("sink").remove_probe(self.encoder_probe)
552 self.pipeline.get_by_name("pay").get_static_pad("src").remove_probe(self.rtp_probe)
553 self.bus.remove_signal_watch()
554 self.pipeline.set_state(Gst.State.NULL)
555 self.pipeline = self.encoder = None
556
557 def send(self, message):
558 if not self.writer.is_closing():
559 self.writer.write(json.dumps(message, separators=(",", ":")).encode() + b"\n")
560 if self.writer.transport.get_write_buffer_size() > 262144:
561 self.writer.close()
562
563 def fail(self, message):
564 self.send({"type": "error", "message": message, "fatal": True})
565 self.writer.close()
566
567 def pipeline_error(self, _, message):
568 error, debug = message.parse_error()
569 logging.error("VM screen stream: %s (%s)", error, debug)
570 self.loop.call_soon_threadsafe(self.fail, "The screen stream stopped. Reconnect to try again.")
571
572 def resize(self):
573 if self.writer.is_closing():
574 return
575 width, height = self.source.width, self.source.height
576 self.frames.set_property("caps", Gst.Caps.from_string(
577 f"video/x-raw,format=BGRx,width={width},height={height},framerate=60/1"))
578 self.pipeline.get_by_name("size").set_property("caps", Gst.Caps.from_string(
579 f"video/x-raw,format=NV12,width={width + width % 2},height={height + height % 2}"))
580 self.send({"type": "display", "width": self.guest.width if self.guest and self.guest.pixels else width,
581 "height": self.guest.height if self.guest and self.guest.pixels else height,
582 "source": "guest" if self.guest and self.guest.pixels else "qemu",
583 "encoder": "nvenc" if self.pipeline.get_by_name("encoder").get_factory().get_name() == "nvh264enc" else "x264"})
584
585 def refresh(self, _, info):
586 event = info.get_event().get_structure()
587 if event and event.get_name() == "GstForceKeyUnit":
588 self.loop.call_soon_threadsafe(self.present)
589 return Gst.PadProbeReturn.OK
590
591 def connection_state(self, rtc, _):
592 state = rtc.get_property("connection-state")
593 if state == GstWebRTC.WebRTCPeerConnectionState.CONNECTED:
594 self.loop.call_soon_threadsafe(self.present)
595 self.loop.call_soon_threadsafe(self.desktop.request, False)
596 if self.guest:
597 self.loop.call_soon_threadsafe(self.guest.sync, self)
598 self.loop.call_soon_threadsafe(self.guest.request, True)
599 elif state in (GstWebRTC.WebRTCPeerConnectionState.FAILED, GstWebRTC.WebRTCPeerConnectionState.CLOSED):
600 self.loop.call_soon_threadsafe(self.fail, "The screen disconnected. Reconnect to try again.")
601
602 def data_channel(self, _, channel):
603 label = channel.get_property("label")
604 if label not in ("input", "pointer") or (self.channel if label == "input" else self.pointer_channel):
605 channel.emit("close")
606 return
607 if label == "input":
608 self.channel = channel
609 else:
610 self.pointer_channel = channel
611 self.handlers.append((channel, channel.connect("on-message-string", lambda _, message: self.loop.call_soon_threadsafe(self.input, message, label))))
612
613 def input(self, raw, channel="input"):
614 if self.writer.is_closing():
615 return
616 try:
617 if len(raw) > 131072:
618 raise ValueError("Shorten the clipboard text, then send it again.")
619 message = json.loads(raw)
620 if not isinstance(message, dict):
621 raise ValueError("Unsupported input event.")
622 if channel == "pointer" and message.get("type") != "pointer":
623 raise ValueError("Unsupported pointer event.")
624 if message.get("type") == "experiment":
625 if type(message.get("enabled")) is not bool or not self.guest:
626 raise ValueError("Unsupported window experiment.")
627 self.experiment = message["enabled"] and bool(self.guest.windows_available)
628 self.guest.request()
629 self.resize()
630 self.present()
631 return
632 if message.get("type") in ("pointer", "release") and "sequence" in message:
633 sequence = message["sequence"]
634 if type(sequence) is not int or not 0 <= sequence <= 0x1fffffffffffff:
635 raise ValueError("Unsupported pointer event.")
636 if channel == "pointer" and sequence <= self.pointer_sequence:
637 return
638 self.pointer_sequence = max(sequence, self.pointer_sequence)
639 if channel == "pointer":
640 message["buttons"] = self.desktop.buttons
641 self.desktop.input(message)
642 if self.desktop.writer.transport.get_write_buffer_size() > 262144:
643 self.fail("The VM stopped accepting input. Reconnect to try again.")
644 except (ValueError, TypeError) as error:
645 self.send({"type": "error", "message": str(error)})
646
647 @property
648 def source(self):
649 if self.experiment and self.guest and self.guest.pixels and self.guest.layers and self.guest.layers.complete:
650 return self.guest.layers
651 return self.guest if self.guest and self.guest.pixels else self.desktop
652
653 def present(self):
654 if self.writer.is_closing():
655 return
656 if self.frame_timer is None and self.rtc.get_property("connection-state") == GstWebRTC.WebRTCPeerConnectionState.CONNECTED:
657 self.frame_timer = self.loop.call_later(max(0, self.last_frame + 1 / 60 - time.monotonic()), self.push_frame)
658
659 def push_frame(self):
660 self.frame_timer = None
661 if self.writer.is_closing():
662 return
663 self.last_frame = time.monotonic()
664 source = self.source
665 buffer = Gst.Buffer.new_wrapped(bytes(source.pixels))
666 buffer.pts = self.pipeline.get_clock().get_time() - self.pipeline.get_base_time()
667 buffer.duration = Gst.SECOND // 60
668 layout = {"type": "windows", "windows": source.windows if isinstance(source, WindowDesktop) else [], "width": source.width, "height": source.height}
669 self.layouts[buffer.pts] = layout
670 while len(self.layouts) > 64:
671 del self.layouts[next(iter(self.layouts))]
672 self.frames.emit("push-buffer", buffer)
673
674 def encoding(self, _, info):
675 layout = self.layouts.pop(info.get_buffer().pts, None)
676 if layout is not None:
677 self.encoding_layouts.append(layout)
678 return Gst.PadProbeReturn.OK
679
680 def rtp_layout(self, _, info):
681 buffer = info.get_buffer()
682 if buffer.pts != self.last_rtp_pts and self.encoding_layouts:
683 self.last_rtp_pts = buffer.pts
684 layout = self.encoding_layouts.popleft()
685 if layout != self.sent_layout:
686 self.sent_layout = layout
687 timestamp = struct.unpack("!I", buffer.extract_dup(4, 4))[0]
688 self.loop.call_soon_threadsafe(self.send, {**layout, "rtpTimestamp": timestamp})
689 return Gst.PadProbeReturn.OK
690
691 def signal(self, message):
692 if message.get("type") == "offer" and isinstance(message.get("sdp"), str):
693 if self.rtc.get_property("remote-description"):
694 raise ValueError("Reconnect to start a new screen connection.")
695 result, sdp = GstSdp.SDPMessage.new_from_text(message["sdp"])
696 if result != GstSdp.SDPResult.OK:
697 raise ValueError("Unable to connect the screen. Reconnect to try again.")
698 caps = None
699 for index in range(sdp.medias_len()):
700 media = sdp.get_media(index)
701 if media.get_media() != "video":
702 continue
703 for format_index in range(media.formats_len()):
704 payload = int(media.get_format(format_index))
705 offered = media.get_caps_from_media(payload)
706 if not offered:
707 continue
708 codec = offered.get_structure(0)
709 profile = codec.get_string("profile-level-id")
710 if codec.get_string("encoding-name") == "H264" and codec.get_string("packetization-mode") == "1" and (
711 not profile or profile.startswith("42")):
712 caps = offered
713 break
714 if caps:
715 break
716 if not caps:
717 raise ValueError("This browser doesn't support the VM's H.264 stream. Open it in a current browser.")
718 payload = caps.get_structure(0).get_value("payload")
719 extension = None
720 for index in range(media.attributes_len()):
721 attribute = media.get_attribute(index)
722 if attribute.key == "extmap" and attribute.value.split()[1:2] == [TWCC]:
723 extension = int(attribute.value.split()[0].split("/")[0])
724 if not 1 <= extension <= 255:
725 raise ValueError("Unsupported screen connection extension.")
726 if extension and Gst.ElementFactory.find("rtpgccbwe"):
727 self.bwe = Gst.ElementFactory.make("rtpgccbwe")
728 self.bwe.set_property("min-bitrate", 500000)
729 self.bwe.set_property("max-bitrate", 12000000)
730 self.bwe.set_property("estimated-bitrate", 4000000)
731 self.handlers.append((self.bwe, self.bwe.connect("notify::estimated-bitrate", lambda bwe, _: self.encoder.set_property("bitrate", bwe.get_property("estimated-bitrate") // 1000))))
732 header = GstRtp.RTPHeaderExtension.create_from_uri(TWCC)
733 header.set_id(extension)
734 self.pipeline.get_by_name("pay").emit("add-extension", header)
735 caps = Gst.Caps.from_string(
736 f"application/x-rtp,media=video,encoding-name=H264,clock-rate=90000,payload={payload},"
737 f"ssrc=(uint){self.pipeline.get_by_name('pay').get_property('ssrc')},"
738 "packetization-mode=(string)1,rtcp-fb-nack=(boolean)true,rtcp-fb-nack-pli=(boolean)true")
739 if self.bwe:
740 caps.set_value("rtcp-fb-transport-cc", True)
741 caps.set_value(f"extmap-{extension}", TWCC)
742 self.pipeline.get_by_name("pay").set_property("pt", payload)
743 self.pipeline.get_by_name("codec").set_property("caps", caps)
744 transceiver = self.rtc.get_static_pad("sink_0").get_property("transceiver")
745 preferences = caps.copy()
746 for index in range(media.formats_len()):
747 offered = media.get_caps_from_media(int(media.get_format(index)))
748 if offered:
749 codec = offered.get_structure(0)
750 if codec.get_string("encoding-name") == "RTX" and codec.get_string("apt") == str(payload):
751 preferences.append(offered)
752 transceiver.set_property("codec-preferences", preferences)
753 transceiver.set_property("direction", GstWebRTC.WebRTCRTPTransceiverDirection.SENDONLY)
754 transceiver.set_property("do-nack", True)
755 description = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sdp)
756 promise = Gst.Promise.new_with_change_func(self.remote_description, None, None)
757 self.rtc.emit("set-remote-description", description, promise)
758 elif message.get("type") == "candidate":
759 line, candidate = message.get("sdpMLineIndex"), message.get("candidate")
760 if type(line) is not int or not 0 <= line <= 16 or not isinstance(candidate, str):
761 raise ValueError("Unsupported screen connection candidate.")
762 self.rtc.emit("add-ice-candidate", line, candidate)
763 else:
764 raise ValueError("Unsupported screen connection message.")
765
766 def remote_description(self, promise, *_):
767 if self.writer.is_closing():
768 return
769 reply = promise.get_reply()
770 if reply and reply.has_field("error"):
771 logging.error("VM screen offer: %s", reply.to_string())
772 self.loop.call_soon_threadsafe(self.fail, "Unable to connect the screen. Reconnect to try again.")
773 return
774 self.rtc.emit("create-answer", None, Gst.Promise.new_with_change_func(self.answer, None, None))
775
776 def answer(self, promise, *_):
777 if self.writer.is_closing():
778 return
779 reply = promise.get_reply()
780 if not reply or not reply.has_field("answer"):
781 logging.error("VM screen answer: %s", reply.to_string() if reply else "missing reply")
782 self.loop.call_soon_threadsafe(self.fail, "Unable to connect the screen. Reconnect to try again.")
783 return
784 answer = reply.get_value("answer")
785 self.rtc.emit("set-local-description", answer, Gst.Promise.new())
786 self.loop.call_soon_threadsafe(self.send, {"type": "answer", "sdp": answer.sdp.as_text()})
787
788
789async def serve(reader, writer):
790 desktop = screen = task = None
791 if SLOTS.locked():
792 writer.write(b'{"type":"error","fatal":true,"message":"Four VM screens are already open. Close one, then reconnect."}\n')
793 writer.close()
794 return
795 async with SLOTS:
796 try:
797 peer = writer.get_extra_info("socket")
798 if struct.unpack("3i", peer.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))[1] not in (0, os.getuid()):
799 return
800 name = json.loads(await asyncio.wait_for(reader.readline(), 10))["name"]
801 if not isinstance(name, str) or not name or len(name) > 63:
802 raise ValueError("Choose a VM from the list.")
803 host_reader, host_writer = await asyncio.wait_for(asyncio.open_unix_connection(
804 os.environ.get("STUDIO_HOST_SOCKET", "/run/studio-host/host.sock")), 10)
805 desktop = Desktop(host_reader, host_writer, lambda message: screen.send(message) if message["type"] != "cursor" or screen.source is desktop else None)
806 host_peer = host_writer.get_extra_info("socket")
807 if struct.unpack("3i", host_peer.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))[1] != 0:
808 raise ValueError("The host identity couldn't be verified. Check its configuration.")
809 host_writer.write(json.dumps({"operation": "vm.console", "payload": {"name": name}}).encode() + b"\n")
810 size = struct.unpack("!I", await asyncio.wait_for(host_reader.readexactly(4), 70))[0]
811 if size > 65536:
812 raise ValueError("The host screen response was too large.")
813 result = json.loads(await host_reader.readexactly(size))
814 if "error" in result:
815 raise ValueError(result["error"])
816 await asyncio.wait_for(desktop.start(), 10)
817 screen = Screen(writer, desktop)
818 guest = GUESTS.get(name)
819 if guest and guest.stopping:
820 await asyncio.shield(guest.task)
821 if name not in GUESTS:
822 GUESTS[name] = GuestDesktop(name)
823 GUESTS[name].task = asyncio.create_task(GUESTS[name].run())
824 screen.guest = GUESTS[name]
825 screen.guest.screens.add(screen)
826 if screen.guest.pixels:
827 screen.resize()
828 if screen.guest.pointer:
829 screen.send(screen.guest.pointer)
830 for cursor in screen.guest.cursor:
831 screen.send(cursor)
832 screen.send({"type": "ready", "iceServers": [{"urls": STUN.replace("stun://", "stun:")}] if STUN else []})
833 task = asyncio.create_task(desktop.read(lambda: screen.present() if screen.source is desktop else None,
834 lambda: screen.resize() if screen.source is desktop else None))
835 def captured(done):
836 if done.cancelled():
837 return
838 error = done.exception()
839 if error:
840 logging.error("VM screen capture: %s", error)
841 screen.fail(str(error) if isinstance(error, ValueError) else
842 "The VM screen stopped. Check that the VM is running, then reconnect.")
843 task.add_done_callback(captured)
844 async with asyncio.timeout(8 * 3600):
845 while raw := await reader.readline():
846 message = json.loads(raw)
847 if not isinstance(message, dict):
848 raise ValueError("Unsupported screen connection message.")
849 screen.signal(message)
850 if task.done():
851 task.result()
852 except ValueError as error:
853 if screen:
854 screen.send({"type": "error", "message": str(error), "fatal": True})
855 else:
856 writer.write(json.dumps({"type": "error", "message": str(error), "fatal": True}).encode() + b"\n")
857 except (OSError, EOFError, TimeoutError, GLib.Error, asyncio.IncompleteReadError):
858 logging.exception("VM screen connection stopped")
859 if screen:
860 screen.send({"type": "error", "fatal": True, "message": "Unable to connect the VM screen. Reconnect to try again."})
861 finally:
862 if task:
863 task.cancel()
864 with contextlib.suppress(asyncio.CancelledError, Exception):
865 await task
866 if screen:
867 screen.close()
868 if desktop:
869 with contextlib.suppress(OSError):
870 desktop.input({"type": "release"})
871 desktop.writer.close()
872 try:
873 await asyncio.wait_for(desktop.writer.wait_closed(), 2)
874 except (OSError, TimeoutError):
875 desktop.writer.transport.abort()
876 writer.close()
877 try:
878 await asyncio.wait_for(writer.wait_closed(), 2)
879 except (OSError, TimeoutError):
880 writer.transport.abort()
881
882
883async def main():
884 Gst.init(None)
885 glib = GLib.MainLoop()
886 threading.Thread(target=glib.run, daemon=True).start()
887 stopped = asyncio.Event()
888 loop = asyncio.get_running_loop()
889 for signum in (signal.SIGTERM, signal.SIGINT):
890 loop.add_signal_handler(signum, stopped.set)
891 clients = set()
892 def connected(reader, writer):
893 if stopped.is_set():
894 writer.close()
895 return
896 task = asyncio.create_task(serve(reader, writer))
897 clients.add(task)
898 task.add_done_callback(clients.discard)
899 path = Path(os.environ.get("STUDIO_VM_SCREEN_SOCKET", "/run/studio-vm-screen/screen.sock"))
900 path.unlink(missing_ok=True)
901 server = await asyncio.start_unix_server(connected, path=str(path), limit=131072)
902 path.chmod(0o600)
903 await stopped.wait()
904 server.close()
905 guests = tuple(GUESTS.values())
906 for guest in guests:
907 guest.stop()
908 tasks = tuple(clients)
909 for task in tasks:
910 task.cancel()
911 await asyncio.gather(*tasks, return_exceptions=True)
912 await server.wait_closed()
913 await asyncio.gather(*(guest.task for guest in guests), return_exceptions=True)
914 glib.quit()
915
916
917if __name__ == "__main__":
918 asyncio.run(main())