1#!/usr/bin/env python3
2import asyncio
3import base64
4import contextlib
5import json
6import logging
7import os
8from pathlib import Path
9import socket
10import struct
11import threading
12import zlib
13
14import gi
15
16gi.require_version("Gst", "1.0")
17gi.require_version("GstSdp", "1.0")
18gi.require_version("GstWebRTC", "1.0")
19from gi.repository import GLib, Gst, GstSdp, GstWebRTC
20
21
22MAX_PIXELS = 3840 * 2160
23STUN = os.environ.get("STUDIO_VM_STUN_SERVER", "stun://stun.cloudflare.com:3478")
24SLOTS = asyncio.Semaphore(4)
25GUEST_PROTOCOL = json.loads(Path(os.environ.get("STUDIO_VM_GUEST_PROTOCOL", Path(__file__).resolve().parent.parent / "guest/protocol.json")).read_text())
26GUESTS = {}
27
28
29def png(width, height, pixels):
30 def chunk(kind, data):
31 return struct.pack("!I", len(data)) + kind + data + struct.pack("!I", zlib.crc32(kind + data))
32 rows = b"".join(b"\0" + pixels[row * width * 4:(row + 1) * width * 4] for row in range(height))
33 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"")
34 return "data:image/png;base64," + base64.b64encode(data).decode()
35
36
37class GuestDesktop:
38 def __init__(self, name):
39 self.name = name
40 self.screens = set()
41 self.pixels = bytearray()
42 self.cursor = []
43 self.pointer = None
44
45 def suspend(self):
46 active = bool(self.pixels)
47 self.pixels = bytearray()
48 self.cursor.clear()
49 self.pointer = None
50 if active:
51 for screen in self.screens:
52 screen.resize()
53 screen.send(screen.desktop.cursor or {"type": "cursor", "source": "qemu", "frame": 0, "fallback": True, "x": 0, "y": 0})
54 screen.present()
55
56 async def read(self, reader):
57 limits = GUEST_PROTOCOL["limits"]
58 packets = {packet["id"]: (name, packet["fields"]) for name, packet in GUEST_PROTOCOL["packets"].items()}
59 while True:
60 async with asyncio.timeout(4):
61 magic, kind, length = struct.unpack("!4sII", await reader.readexactly(12))
62 if magic != GUEST_PROTOCOL["magic"].encode() or kind not in packets or length > limits["packet_bytes"]:
63 raise ValueError("Invalid guest display packet")
64 name, fields = packets[kind]
65 if length < len(fields) * 4:
66 raise ValueError("Incomplete guest display packet")
67 payload = await reader.readexactly(length)
68 values = dict(zip(fields, struct.unpack("!" + "I" * len(fields), payload[:len(fields) * 4])))
69 data = payload[len(fields) * 4:]
70 if name == "screen":
71 width, height = values["width"], values["height"]
72 x, y, w, h = values["x"], values["y"], values["rect_width"], values["rect_height"]
73 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:
74 raise ValueError("Invalid guest display rectangle")
75 resized = not self.pixels or (width, height) != (self.width, self.height)
76 if resized:
77 if (x, y, w, h) != (0, 0, width, height):
78 raise ValueError("Guest display needs a complete first frame")
79 self.width, self.height = width, height
80 self.pixels = bytearray(width * height * 4)
81 for row in range(h):
82 offset = ((y + row) * width + x) * 4
83 self.pixels[offset:offset + w * 4] = data[row * w * 4:(row + 1) * w * 4]
84 for screen in tuple(self.screens):
85 if resized:
86 screen.resize()
87 if self.pointer:
88 screen.send(self.pointer)
89 for cursor in self.cursor:
90 screen.send(cursor)
91 screen.present()
92 await asyncio.sleep(0)
93 elif name == "cursor":
94 width, height, count = values["width"], values["height"], values["frames"]
95 if not width and not height and not count and not data and not values["hot_x"] and not values["hot_y"]:
96 self.cursor = [{"type": "cursor", "source": "guest", "id": values["id"], "width": 0, "height": 0, "frame": 0, "frames": 0, "fallback": True, "x": 0, "y": 0}]
97 if self.pixels:
98 for screen in self.screens:
99 screen.send(self.cursor[0])
100 continue
101 pixels = width * height
102 frame_bytes = 4 + pixels * 5
103 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:
104 raise ValueError("Invalid guest cursor")
105 self.cursor.clear()
106 for index in range(count):
107 frame = data[index * frame_bytes:(index + 1) * frame_bytes]
108 duration = struct.unpack("!I", frame[:4])[0]
109 if not 1 <= duration <= 60000:
110 raise ValueError("Invalid guest cursor duration")
111 mask = frame[4 + pixels * 4:]
112 if any(value & ~sum(GUEST_PROTOCOL["xor_bits"].values()) for value in mask):
113 raise ValueError("Invalid guest cursor inversion")
114 message = {"type": "cursor", "source": "guest", "id": values["id"], "width": width, "height": height,
115 "x": values["hot_x"], "y": values["hot_y"], "frame": index, "frames": count,
116 "duration": duration, "image": png(width, height, frame[4:4 + pixels * 4])}
117 if any(mask):
118 bits = GUEST_PROTOCOL["xor_bits"]
119 message["invert"] = png(width, height, b"".join(bytes([255 if value & bits["red"] else 0,
120 255 if value & bits["green"] else 0,
121 255 if value & bits["blue"] else 0,
122 255 if value else 0]) for value in mask))
123 self.cursor.append(message)
124 if self.pixels:
125 for screen in tuple(self.screens):
126 screen.send(message)
127 await asyncio.sleep(0)
128 elif name == "pointer":
129 if data or values["visible"] > 1:
130 raise ValueError("Invalid guest pointer")
131 self.pointer = {"type": "cursor-position", "id": values["id"], "visible": bool(values["visible"]),
132 "x": (values["x"] ^ 0x80000000) - 0x80000000, "y": (values["y"] ^ 0x80000000) - 0x80000000}
133 if self.pixels:
134 for screen in self.screens:
135 screen.send(self.pointer)
136 elif name == "suspend":
137 if data:
138 raise ValueError("Invalid guest display suspension")
139 self.suspend()
140 elif name == "heartbeat":
141 if data:
142 raise ValueError("Invalid guest display heartbeat")
143 else:
144 raise ValueError("Unexpected guest display request")
145
146 async def run(self):
147 while True:
148 writer = None
149 try:
150 reader, writer = await asyncio.wait_for(asyncio.open_unix_connection(
151 os.environ.get("STUDIO_HOST_SOCKET", "/run/studio-host/host.sock")), 10)
152 peer = writer.get_extra_info("socket")
153 if struct.unpack("3i", peer.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))[1] != 0:
154 raise ValueError("Untrusted guest display broker")
155 writer.write(json.dumps({"operation": "vm.guest", "payload": {"name": self.name}}).encode() + b"\n")
156 size = struct.unpack("!I", await asyncio.wait_for(reader.readexactly(4), 70))[0]
157 if size > 65536:
158 raise ValueError("Invalid guest display broker response")
159 response = json.loads(await reader.readexactly(size))
160 if "error" not in response:
161 writer.write(struct.pack("!4sII", GUEST_PROTOCOL["magic"].encode(), GUEST_PROTOCOL["packets"]["refresh"]["id"], 0))
162 await self.read(reader)
163 except (OSError, ValueError, EOFError, asyncio.IncompleteReadError, TimeoutError) as error:
164 logging.info("Guest display %s: %s", self.name, error)
165 finally:
166 self.suspend()
167 if writer:
168 writer.close()
169 with contextlib.suppress(OSError, TimeoutError):
170 await asyncio.wait_for(writer.wait_closed(), 2)
171 await asyncio.sleep(2)
172
173
174class Desktop:
175 def __init__(self, reader, writer, notify):
176 self.reader, self.writer, self.notify = reader, writer, notify
177 self.keys = set()
178 self.pointer = (0, 0)
179 self.pixels = bytearray()
180 self.cursor = None
181
182 def resize(self, width, height):
183 if not 0 < width <= 4096 or not 0 < height <= 4096 or width * height > MAX_PIXELS:
184 raise ValueError("The VM screen exceeds the supported 4K display size.")
185 self.width, self.height = width, height
186 self.pixels = bytearray(width * height * 4)
187
188 async def start(self):
189 if not (await self.reader.readexactly(12)).startswith(b"RFB 003."):
190 raise ValueError("The VM screen did not respond. Restart the VM and reconnect.")
191 self.writer.write(b"RFB 003.008\n")
192 count = (await self.reader.readexactly(1))[0]
193 security = await self.reader.readexactly(count)
194 if 1 not in security:
195 raise ValueError("Remove the VM's VNC password in its display settings, then reconnect.")
196 self.writer.write(b"\x01")
197 if await self.reader.readexactly(4) != b"\0\0\0\0":
198 raise ValueError("The VM rejected the screen connection. Restart it and reconnect.")
199 self.writer.write(b"\x01")
200 width, height, _, length = struct.unpack("!HH16sI", await self.reader.readexactly(24))
201 if length > 65536:
202 raise ValueError("The VM screen sent an unsupported desktop name.")
203 await self.reader.readexactly(length)
204 self.resize(width, height)
205 self.writer.write(struct.pack("!B3xBBBBHHHBBB3x", 0, 32, 24, 0, 1, 255, 255, 255, 16, 8, 0))
206 encodings = (0, 1, -223, -224, -239)
207 self.writer.write(struct.pack("!BBH" + "i" * len(encodings), 2, 0, len(encodings), *encodings))
208 self.request(False)
209
210 def request(self, incremental=True):
211 self.writer.write(struct.pack("!BBHHHH", 3, incremental, 0, 0, self.width, self.height))
212
213 def input(self, message):
214 kind = message.get("type")
215 if kind == "key":
216 key, down = message.get("key"), message.get("down")
217 if type(key) is not int or not 0 <= key <= 0xFFFFFFFF or type(down) is not bool:
218 raise ValueError("Unsupported key event.")
219 if down:
220 self.keys.add(key)
221 else:
222 self.keys.discard(key)
223 self.writer.write(struct.pack("!BB2xI", 4, down, key))
224 elif kind == "pointer":
225 x, y, buttons = message.get("x"), message.get("y"), message.get("buttons")
226 if any(type(value) is not int for value in (x, y, buttons)) or not 0 <= buttons <= 255:
227 raise ValueError("Unsupported pointer event.")
228 self.pointer = (max(0, min(x, self.width - 1)), max(0, min(y, self.height - 1)))
229 self.writer.write(struct.pack("!BBHH", 5, buttons, *self.pointer))
230 elif kind == "release":
231 for key in self.keys:
232 self.writer.write(struct.pack("!BB2xI", 4, 0, key))
233 self.keys.clear()
234 self.writer.write(struct.pack("!BBHH", 5, 0, *self.pointer))
235 elif kind == "clipboard":
236 text = message.get("text")
237 if not isinstance(text, str) or len(text) > 65536:
238 raise ValueError("Shorten the clipboard text, then send it again.")
239 try:
240 data = text.encode("latin1")
241 except UnicodeEncodeError:
242 raise ValueError("Some clipboard characters aren't supported by this guest.") from None
243 self.writer.write(struct.pack("!B3xI", 6, len(data)) + data)
244 else:
245 raise ValueError("Unsupported input event.")
246
247 async def read(self, present, resized):
248 while True:
249 kind = (await self.reader.readexactly(1))[0]
250 if kind == 0:
251 _, count = struct.unpack("!BH", await self.reader.readexactly(3))
252 dirty = False
253 for _ in range(count):
254 x, y, width, height, encoding = struct.unpack("!HHHHi", await self.reader.readexactly(12))
255 if encoding == -224:
256 break
257 if encoding == -223:
258 self.resize(width, height)
259 resized()
260 continue
261 if encoding == -239:
262 if width * height > 256 * 256:
263 raise ValueError("The VM sent an unsupported cursor size.")
264 pixels = bytearray(await self.reader.readexactly(width * height * 4))
265 mask = await self.reader.readexactly(((width + 7) // 8) * height)
266 for row in range(height):
267 for col in range(width):
268 offset = (row * width + col) * 4
269 pixels[offset], pixels[offset + 2] = pixels[offset + 2], pixels[offset]
270 pixels[offset + 3] = 255 if mask[row * ((width + 7) // 8) + col // 8] & (128 >> (col % 8)) else 0
271 self.cursor = {"type": "cursor", "source": "qemu", "width": width, "height": height, "x": x, "y": y,
272 "frame": 0, "frames": 1, "duration": 100,
273 "image": png(width, height, pixels) if width and height else None}
274 self.notify(self.cursor)
275 continue
276 if x + width > self.width or y + height > self.height:
277 raise ValueError("The VM sent a screen update outside its display.")
278 if encoding == 0:
279 data = await self.reader.readexactly(width * height * 4)
280 elif encoding == 1:
281 source_x, source_y = struct.unpack("!HH", await self.reader.readexactly(4))
282 if source_x + width > self.width or source_y + height > self.height:
283 raise ValueError("The VM sent a screen copy outside its display.")
284 data = b"".join(self.pixels[((source_y + row) * self.width + source_x) * 4:
285 ((source_y + row) * self.width + source_x + width) * 4] for row in range(height))
286 else:
287 raise ValueError("The VM sent an unsupported screen encoding.")
288 for row in range(height):
289 start = ((y + row) * self.width + x) * 4
290 self.pixels[start:start + width * 4] = data[row * width * 4:(row + 1) * width * 4]
291 dirty = dirty or bool(width and height)
292 if dirty:
293 present()
294 self.request()
295 await self.writer.drain()
296 elif kind == 2:
297 pass
298 elif kind == 3:
299 _, length = struct.unpack("!3sI", await self.reader.readexactly(7))
300 if length > 65536:
301 raise ValueError("The VM sent too much clipboard text.")
302 self.notify({"type": "clipboard", "text": (await self.reader.readexactly(length)).decode("latin1")})
303 else:
304 raise ValueError("The VM sent an unsupported screen message.")
305
306
307class Screen:
308 def __init__(self, writer, desktop):
309 self.loop = asyncio.get_running_loop()
310 self.writer, self.desktop = writer, desktop
311 self.guest = None
312 self.channel = None
313 hardware = Gst.ElementFactory.make("nvh264enc")
314 encoder = "nvh264enc name=encoder preset=p4 tune=ultra-low-latency zerolatency=true bframes=0 bitrate=6000 gop-size=30" if hardware else (
315 "x264enc name=encoder tune=zerolatency speed-preset=ultrafast bitrate=6000 key-int-max=30 bframes=0")
316 width, height = self.desktop.width, self.desktop.height
317 self.pipeline = Gst.parse_launch(
318 f"appsrc name=frames is-live=true format=time do-timestamp=true max-buffers=1 leaky-type=downstream "
319 f"caps=video/x-raw,format=BGRx,width={width},height={height},framerate=30/1 "
320 f"! videoconvert ! videoscale ! capsfilter name=size caps=video/x-raw,format=NV12,width={width + width % 2},height={height + height % 2} "
321 f"! {encoder} ! video/x-h264,profile=constrained-baseline ! h264parse "
322 "! rtph264pay name=pay pt=96 mtu=1200 config-interval=-1 aggregate-mode=zero-latency "
323 "! capsfilter name=codec caps=application/x-rtp,media=video,encoding-name=H264,clock-rate=90000 "
324 "! webrtcbin name=rtc bundle-policy=max-bundle")
325 self.frames = self.pipeline.get_by_name("frames")
326 self.rtc = self.pipeline.get_by_name("rtc")
327 self.ice = self.rtc.get_property("ice-agent")
328 self.ice._ref_sink() # PyGObject sinks the default agent's floating reference, which webrtcbin still owns.
329 self.ice.set_property("min-rtp-port", 50000)
330 self.ice.set_property("max-rtp-port", 50031)
331 if STUN:
332 self.rtc.set_property("stun-server", STUN)
333 self.handlers = [
334 (self.rtc, self.rtc.connect("on-ice-candidate", lambda _, line, candidate: self.loop.call_soon_threadsafe(
335 self.send, {"type": "candidate", "sdpMLineIndex": line, "candidate": candidate}))),
336 (self.rtc, self.rtc.connect("on-data-channel", self.data_channel)),
337 (self.rtc, self.rtc.connect("notify::connection-state", self.connection_state)),
338 ]
339 self.probe = self.frames.get_static_pad("src").add_probe(Gst.PadProbeType.EVENT_UPSTREAM, self.refresh)
340 self.bus = self.pipeline.get_bus()
341 self.bus.add_signal_watch()
342 self.handlers.append((self.bus, self.bus.connect("message::error", self.pipeline_error)))
343 self.pipeline.set_state(Gst.State.PLAYING)
344 self.send({"type": "display", "width": width, "height": height, "encoder": "nvenc" if hardware else "x264"})
345
346 def close(self):
347 self.desktop.notify = None
348 if self.guest:
349 self.guest.screens.discard(self)
350 if not self.guest.screens:
351 GUESTS.pop(self.guest.name, None)
352 self.guest.task.cancel()
353 for element, handler in self.handlers:
354 element.disconnect(handler)
355 self.handlers.clear()
356 self.frames.get_static_pad("src").remove_probe(self.probe)
357 self.bus.remove_signal_watch()
358 self.pipeline.set_state(Gst.State.NULL)
359
360 def send(self, message):
361 if not self.writer.is_closing():
362 self.writer.write(json.dumps(message, separators=(",", ":")).encode() + b"\n")
363 if self.writer.transport.get_write_buffer_size() > 262144:
364 self.writer.close()
365
366 def fail(self, message):
367 self.send({"type": "error", "message": message, "fatal": True})
368 self.writer.close()
369
370 def pipeline_error(self, _, message):
371 error, debug = message.parse_error()
372 logging.error("VM screen stream: %s (%s)", error, debug)
373 self.loop.call_soon_threadsafe(self.fail, "The screen stream stopped. Reconnect to try again.")
374
375 def resize(self):
376 width, height = self.source.width, self.source.height
377 self.frames.set_property("caps", Gst.Caps.from_string(
378 f"video/x-raw,format=BGRx,width={width},height={height},framerate=30/1"))
379 self.pipeline.get_by_name("size").set_property("caps", Gst.Caps.from_string(
380 f"video/x-raw,format=NV12,width={width + width % 2},height={height + height % 2}"))
381 self.send({"type": "display", "width": width, "height": height,
382 "source": "guest" if self.source is self.guest else "qemu",
383 "encoder": "nvenc" if self.pipeline.get_by_name("encoder").get_factory().get_name() == "nvh264enc" else "x264"})
384
385 def refresh(self, _, info):
386 event = info.get_event().get_structure()
387 if event and event.get_name() == "GstForceKeyUnit":
388 self.loop.call_soon_threadsafe(self.present)
389 return Gst.PadProbeReturn.OK
390
391 def connection_state(self, rtc, _):
392 state = rtc.get_property("connection-state")
393 if state == GstWebRTC.WebRTCPeerConnectionState.CONNECTED:
394 self.loop.call_soon_threadsafe(self.present)
395 self.loop.call_soon_threadsafe(self.desktop.request, False)
396 elif state in (GstWebRTC.WebRTCPeerConnectionState.FAILED, GstWebRTC.WebRTCPeerConnectionState.CLOSED):
397 self.loop.call_soon_threadsafe(self.fail, "The screen disconnected. Reconnect to try again.")
398
399 def data_channel(self, _, channel):
400 if channel.get_property("label") != "input" or self.channel:
401 channel.emit("close")
402 return
403 self.channel = channel
404 self.handlers.append((channel, channel.connect("on-message-string", lambda _, message: self.loop.call_soon_threadsafe(self.input, message))))
405
406 def input(self, raw):
407 if self.writer.is_closing():
408 return
409 try:
410 if len(raw) > 131072:
411 raise ValueError("Shorten the clipboard text, then send it again.")
412 message = json.loads(raw)
413 if not isinstance(message, dict):
414 raise ValueError("Unsupported input event.")
415 self.desktop.input(message)
416 if self.desktop.writer.transport.get_write_buffer_size() > 262144:
417 self.fail("The VM stopped accepting input. Reconnect to try again.")
418 except (ValueError, TypeError) as error:
419 self.send({"type": "error", "message": str(error)})
420
421 @property
422 def source(self):
423 return self.guest if self.guest and self.guest.pixels else self.desktop
424
425 def present(self):
426 if self.rtc.get_property("connection-state") == GstWebRTC.WebRTCPeerConnectionState.CONNECTED:
427 buffer = Gst.Buffer.new_wrapped(bytes(self.source.pixels))
428 self.frames.emit("push-buffer", buffer)
429
430 def signal(self, message):
431 if message.get("type") == "offer" and isinstance(message.get("sdp"), str):
432 if self.rtc.get_property("remote-description"):
433 raise ValueError("Reconnect to start a new screen connection.")
434 result, sdp = GstSdp.SDPMessage.new_from_text(message["sdp"])
435 if result != GstSdp.SDPResult.OK:
436 raise ValueError("Unable to connect the screen. Reconnect to try again.")
437 caps = None
438 for index in range(sdp.medias_len()):
439 media = sdp.get_media(index)
440 if media.get_media() != "video":
441 continue
442 for format_index in range(media.formats_len()):
443 payload = int(media.get_format(format_index))
444 offered = media.get_caps_from_media(payload)
445 if not offered:
446 continue
447 codec = offered.get_structure(0)
448 profile = codec.get_string("profile-level-id")
449 if codec.get_string("encoding-name") == "H264" and codec.get_string("packetization-mode") == "1" and (
450 not profile or profile.startswith("42")):
451 caps = offered
452 break
453 if caps:
454 break
455 if not caps:
456 raise ValueError("This browser doesn't support the VM's H.264 stream. Open it in a current browser.")
457 payload = caps.get_structure(0).get_value("payload")
458 caps = Gst.Caps.from_string(
459 f"application/x-rtp,media=video,encoding-name=H264,clock-rate=90000,payload={payload},"
460 "packetization-mode=(string)1,rtcp-fb-nack=(boolean)true,rtcp-fb-nack-pli=(boolean)true")
461 self.pipeline.get_by_name("pay").set_property("pt", payload)
462 self.pipeline.get_by_name("codec").set_property("caps", caps)
463 transceiver = self.rtc.get_static_pad("sink_0").get_property("transceiver")
464 transceiver.set_property("codec-preferences", caps)
465 transceiver.set_property("direction", GstWebRTC.WebRTCRTPTransceiverDirection.SENDONLY)
466 transceiver.set_property("do-nack", True)
467 description = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sdp)
468 promise = Gst.Promise.new_with_change_func(self.remote_description, None, None)
469 self.rtc.emit("set-remote-description", description, promise)
470 elif message.get("type") == "candidate":
471 line, candidate = message.get("sdpMLineIndex"), message.get("candidate")
472 if type(line) is not int or not 0 <= line <= 16 or not isinstance(candidate, str):
473 raise ValueError("Unsupported screen connection candidate.")
474 self.rtc.emit("add-ice-candidate", line, candidate)
475 else:
476 raise ValueError("Unsupported screen connection message.")
477
478 def remote_description(self, promise, *_):
479 reply = promise.get_reply()
480 if reply and reply.has_field("error"):
481 logging.error("VM screen offer: %s", reply.to_string())
482 self.loop.call_soon_threadsafe(self.fail, "Unable to connect the screen. Reconnect to try again.")
483 return
484 self.rtc.emit("create-answer", None, Gst.Promise.new_with_change_func(self.answer, None, None))
485
486 def answer(self, promise, *_):
487 reply = promise.get_reply()
488 if not reply or not reply.has_field("answer"):
489 logging.error("VM screen answer: %s", reply.to_string() if reply else "missing reply")
490 self.loop.call_soon_threadsafe(self.fail, "Unable to connect the screen. Reconnect to try again.")
491 return
492 answer = reply.get_value("answer")
493 self.rtc.emit("set-local-description", answer, Gst.Promise.new())
494 self.loop.call_soon_threadsafe(self.send, {"type": "answer", "sdp": answer.sdp.as_text()})
495
496
497async def serve(reader, writer):
498 desktop = screen = task = None
499 if SLOTS.locked():
500 writer.write(b'{"type":"error","fatal":true,"message":"Four VM screens are already open. Close one, then reconnect."}\n')
501 writer.close()
502 return
503 async with SLOTS:
504 try:
505 peer = writer.get_extra_info("socket")
506 if struct.unpack("3i", peer.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))[1] not in (0, os.getuid()):
507 return
508 name = json.loads(await asyncio.wait_for(reader.readline(), 10))["name"]
509 if not isinstance(name, str) or not name or len(name) > 63:
510 raise ValueError("Choose a VM from the list.")
511 host_reader, host_writer = await asyncio.wait_for(asyncio.open_unix_connection(
512 os.environ.get("STUDIO_HOST_SOCKET", "/run/studio-host/host.sock")), 10)
513 desktop = Desktop(host_reader, host_writer, lambda message: screen.send(message) if message["type"] != "cursor" or screen.source is desktop else None)
514 host_peer = host_writer.get_extra_info("socket")
515 if struct.unpack("3i", host_peer.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))[1] != 0:
516 raise ValueError("The host identity couldn't be verified. Check its configuration.")
517 host_writer.write(json.dumps({"operation": "vm.console", "payload": {"name": name}}).encode() + b"\n")
518 size = struct.unpack("!I", await asyncio.wait_for(host_reader.readexactly(4), 70))[0]
519 if size > 65536:
520 raise ValueError("The host screen response was too large.")
521 result = json.loads(await host_reader.readexactly(size))
522 if "error" in result:
523 raise ValueError(result["error"])
524 await asyncio.wait_for(desktop.start(), 10)
525 screen = Screen(writer, desktop)
526 if name not in GUESTS:
527 GUESTS[name] = GuestDesktop(name)
528 GUESTS[name].task = asyncio.create_task(GUESTS[name].run())
529 screen.guest = GUESTS[name]
530 screen.guest.screens.add(screen)
531 if screen.guest.pixels:
532 screen.resize()
533 if screen.guest.pointer:
534 screen.send(screen.guest.pointer)
535 for cursor in screen.guest.cursor:
536 screen.send(cursor)
537 screen.send({"type": "ready", "iceServers": [{"urls": STUN.replace("stun://", "stun:")}] if STUN else []})
538 task = asyncio.create_task(desktop.read(lambda: screen.present() if screen.source is desktop else None,
539 lambda: screen.resize() if screen.source is desktop else None))
540 def captured(done):
541 if done.cancelled():
542 return
543 error = done.exception()
544 if error:
545 logging.error("VM screen capture: %s", error)
546 screen.fail(str(error) if isinstance(error, ValueError) else
547 "The VM screen stopped. Check that the VM is running, then reconnect.")
548 task.add_done_callback(captured)
549 async with asyncio.timeout(8 * 3600):
550 while raw := await reader.readline():
551 message = json.loads(raw)
552 if not isinstance(message, dict):
553 raise ValueError("Unsupported screen connection message.")
554 screen.signal(message)
555 if task.done():
556 task.result()
557 except ValueError as error:
558 if screen:
559 screen.send({"type": "error", "message": str(error), "fatal": True})
560 else:
561 writer.write(json.dumps({"type": "error", "message": str(error), "fatal": True}).encode() + b"\n")
562 except (OSError, EOFError, TimeoutError, GLib.Error, asyncio.IncompleteReadError):
563 logging.exception("VM screen connection stopped")
564 if screen:
565 screen.send({"type": "error", "fatal": True, "message": "Unable to connect the VM screen. Reconnect to try again."})
566 finally:
567 if task:
568 task.cancel()
569 with contextlib.suppress(asyncio.CancelledError, Exception):
570 await task
571 if screen:
572 screen.close()
573 if desktop:
574 with contextlib.suppress(OSError):
575 desktop.input({"type": "release"})
576 desktop.writer.close()
577 try:
578 await asyncio.wait_for(desktop.writer.wait_closed(), 2)
579 except (OSError, TimeoutError):
580 desktop.writer.transport.abort()
581 writer.close()
582 try:
583 await asyncio.wait_for(writer.wait_closed(), 2)
584 except (OSError, TimeoutError):
585 writer.transport.abort()
586
587
588async def main():
589 Gst.init(None)
590 threading.Thread(target=GLib.MainLoop().run, daemon=True).start()
591 path = Path(os.environ.get("STUDIO_VM_SCREEN_SOCKET", "/run/studio-vm-screen/screen.sock"))
592 path.unlink(missing_ok=True)
593 server = await asyncio.start_unix_server(serve, path=str(path), limit=131072)
594 path.chmod(0o600)
595 async with server:
596 await server.serve_forever()
597
598
599if __name__ == "__main__":
600 asyncio.run(main())