alan
1ed57291f9
An Omarchy shell plugin: a Python/asyncio daemon that discovers network displays and serves an encoded screen capture for them to pull, plus a Quickshell bar widget and panel to drive it. The screen comes from the xdg-desktop-portal ScreenCast interface (asked for every time, never remembered), goes through GStreamer, and is served from a local HTTP port. Desktop audio is mixed in from the default output's monitor. The container follows the receiver: WebM/VP8 for Cast, MPEG-TS/H.264 for DLNA, HLS for AirPlay video. The stream port has to be reachable from the LAN, and none of these protocols can carry a credential, so the capability is the URL: a fresh random path per session, refused to anything off the LAN, capped at four concurrent readers.
386 lines
15 KiB
Python
386 lines
15 KiB
Python
"""The local HTTP server the receivers pull the screen stream from.
|
|
|
|
Receivers (Cast, DLNA, AirPlay) do not accept a push: they are handed a URL and
|
|
fetch it themselves. So the encoder writes one live byte stream and this fans it
|
|
out to whoever connected, plus serves the HLS directory AirPlay needs.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import ipaddress
|
|
import secrets
|
|
import socket
|
|
import time
|
|
from collections import deque
|
|
from pathlib import Path
|
|
|
|
from aiohttp import web
|
|
|
|
from .util import log
|
|
|
|
# A late client cannot start mid-container, so every stream keeps the bytes the
|
|
# encoder wrote before the first media chunk (EBML head + tracks for WebM, the
|
|
# moov/init for fMP4, PAT/PMT for MPEG-TS) and replays them on connect.
|
|
CLUSTER_MARK = b"\x1f\x43\xb6\x75" # WebM Cluster id
|
|
MOOF_MARK = b"moof" # fragmented MP4 fragment box
|
|
HEADER_LIMIT = 1 << 20 # never hoard more than 1 MiB as "header"
|
|
MIN_BACKLOG = 192 << 10 # never hold less than this, whatever the bitrate
|
|
BACKLOG_SECONDS = 1.5 # ...and never much more than this much screen
|
|
# (a few keyframe-aligned clusters: the drop unit)
|
|
|
|
|
|
class Reader:
|
|
"""One receiver's view of the stream: a short, self-trimming backlog.
|
|
|
|
A receiver pulls at playback speed, so anything we let pile up here is
|
|
latency we hand it and never get back — and a live stream that falls behind
|
|
never catches up on its own. Keeping under a second of screen queued, and
|
|
skipping forward when it does not drain, is what keeps the picture close to
|
|
now rather than however far behind it drifted an hour ago.
|
|
"""
|
|
|
|
def __init__(self, max_bytes: int) -> None:
|
|
self.items: deque[tuple[bytes, bool]] = deque()
|
|
self.queued = 0
|
|
self.dropped = 0
|
|
self.max_bytes = max_bytes
|
|
self.closed = False
|
|
self._wake = asyncio.Event()
|
|
|
|
def push(self, data: bytes, at_boundary: bool) -> None:
|
|
self.items.append((data, at_boundary))
|
|
self.queued += len(data)
|
|
if self.queued > self.max_bytes:
|
|
self._skip_forward()
|
|
self._wake.set()
|
|
|
|
def _skip_forward(self) -> None:
|
|
while self.queued > self.max_bytes and len(self.items) > 1:
|
|
data, _ = self.items.popleft()
|
|
self.queued -= len(data)
|
|
self.dropped += len(data)
|
|
# Resume on a container boundary: dropping into the middle of a cluster
|
|
# hands the receiver's demuxer a torn one.
|
|
while len(self.items) > 1 and not self.items[0][1]:
|
|
data, _ = self.items.popleft()
|
|
self.queued -= len(data)
|
|
self.dropped += len(data)
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
self._wake.set()
|
|
|
|
async def get(self) -> bytes | None:
|
|
while not self.items:
|
|
if self.closed:
|
|
return None
|
|
self._wake.clear()
|
|
await self._wake.wait()
|
|
data, _ = self.items.popleft()
|
|
self.queued -= len(data)
|
|
return data
|
|
|
|
|
|
class Fanout:
|
|
"""One encoder byte stream, many HTTP readers."""
|
|
|
|
def __init__(self, container: str, bitrate_kbps: int = 5000) -> None:
|
|
self.container = container
|
|
self.header = b""
|
|
self.header_done = False
|
|
self.bytes_out = 0
|
|
self.started = time.time()
|
|
self.max_bytes = max(MIN_BACKLOG, int(bitrate_kbps * 1000 / 8 * BACKLOG_SECONDS))
|
|
# Where this container can be cut without tearing a unit in half. MPEG-TS
|
|
# has no such mark - it resynchronises on its own - so it is cut anywhere.
|
|
self._mark = (CLUSTER_MARK if container in ("webm", "matroska")
|
|
else MOOF_MARK if container == "mp4" else None)
|
|
self._carry = b"" # bytes held back in case a mark straddles a chunk
|
|
self._at_boundary = True
|
|
self._readers: set[Reader] = set()
|
|
self._closed = False
|
|
|
|
@property
|
|
def clients(self) -> int:
|
|
return len(self._readers)
|
|
|
|
@property
|
|
def backlog(self) -> int:
|
|
return max((r.queued for r in self._readers), default=0)
|
|
|
|
@property
|
|
def dropped(self) -> int:
|
|
return max((r.dropped for r in self._readers), default=0)
|
|
|
|
def feed(self, chunk: bytes) -> None:
|
|
if self._closed or not chunk:
|
|
return
|
|
self.bytes_out += len(chunk)
|
|
if not self.header_done:
|
|
chunk = self._grow_header(chunk)
|
|
if not self.header_done or not chunk:
|
|
return
|
|
for data, at_boundary in self._cut(chunk):
|
|
for reader in list(self._readers):
|
|
reader.push(data, at_boundary)
|
|
|
|
def _cut(self, chunk: bytes) -> list[tuple[bytes, bool]]:
|
|
"""Split so that every piece starting a new container unit says so."""
|
|
if self._mark is None:
|
|
return [(chunk, True)]
|
|
buf = self._carry + chunk
|
|
width = len(self._mark)
|
|
pieces: list[tuple[bytes, bool]] = []
|
|
start = i = 0
|
|
while True:
|
|
idx = buf.find(self._mark, i)
|
|
if idx < 0:
|
|
break
|
|
if idx > start:
|
|
pieces.append((buf[start:idx], self._at_boundary))
|
|
self._at_boundary = False
|
|
start, i = idx, idx + width
|
|
self._at_boundary = True
|
|
tail = buf[start:]
|
|
hold = min(len(tail), width - 1)
|
|
body, self._carry = tail[: len(tail) - hold], tail[len(tail) - hold :]
|
|
if body:
|
|
pieces.append((body, self._at_boundary))
|
|
self._at_boundary = False
|
|
return pieces
|
|
|
|
def _grow_header(self, chunk: bytes) -> bytes:
|
|
"""Collect the container header; return whatever followed it."""
|
|
buf = self.header + chunk
|
|
mark = CLUSTER_MARK if self.container in ("webm", "matroska") else MOOF_MARK
|
|
idx = buf.find(mark, max(0, len(self.header) - 4))
|
|
if self.container == "mpegts":
|
|
# TS repeats its tables; the first packets are enough to start on.
|
|
if len(buf) >= 32768:
|
|
self.header, self.header_done = buf[:32768], True
|
|
return buf[32768:]
|
|
elif idx > 0:
|
|
self.header, self.header_done = buf[:idx], True
|
|
return buf[idx:]
|
|
if len(buf) >= HEADER_LIMIT:
|
|
self.header, self.header_done = buf[:HEADER_LIMIT], True
|
|
return buf[HEADER_LIMIT:]
|
|
self.header = buf
|
|
return b""
|
|
|
|
def subscribe(self) -> Reader:
|
|
reader = Reader(self.max_bytes)
|
|
self._readers.add(reader)
|
|
return reader
|
|
|
|
def unsubscribe(self, reader: Reader) -> None:
|
|
self._readers.discard(reader)
|
|
reader.close()
|
|
|
|
def close(self) -> None:
|
|
self._closed = True
|
|
for reader in list(self._readers):
|
|
reader.close()
|
|
self._readers.clear()
|
|
|
|
|
|
CONTENT_TYPES = {
|
|
"webm": "video/webm",
|
|
"matroska": "video/x-matroska",
|
|
"mp4": "video/mp4",
|
|
"mpegts": "video/mp2t",
|
|
}
|
|
PATHS = {"webm": "/live.webm", "mp4": "/live.mp4", "mpegts": "/live.ts", "matroska": "/live.mkv"}
|
|
|
|
# One receiver pulls one stream; a couple spare for a retry that has not timed
|
|
# out yet. Anything past this is either a bug or someone else's curiosity, and
|
|
# each reader costs a socket plus its share of the backlog.
|
|
MAX_CLIENTS = 4
|
|
|
|
|
|
def _is_lan(peer: str | None) -> bool:
|
|
"""Receivers live on the LAN. Anything routed in from outside it is not one."""
|
|
if not peer:
|
|
return False
|
|
try:
|
|
ip = ipaddress.ip_address(peer.split("%", 1)[0])
|
|
except ValueError:
|
|
return False
|
|
if ip.version == 6 and ip.ipv4_mapped:
|
|
ip = ip.ipv4_mapped
|
|
return ip.is_private or ip.is_loopback or ip.is_link_local
|
|
|
|
|
|
class StreamServer:
|
|
"""aiohttp server exposing whatever the current session is encoding.
|
|
|
|
`on_client` lets the session know a receiver actually connected (or that the
|
|
last one went away) — that is the only reliable signal that casting started.
|
|
"""
|
|
|
|
def __init__(self, on_client=None) -> None:
|
|
self.port = 0
|
|
self.fanout: Fanout | None = None
|
|
self.hls_dir: Path | None = None
|
|
self.on_client = on_client
|
|
# The stream port has to be reachable from the LAN or no receiver can
|
|
# pull it, and none of these protocols can carry a credential. So the
|
|
# capability lives in the URL: a fresh unguessable path per session,
|
|
# which is the only thing standing between a shared screen and everyone
|
|
# else on the network.
|
|
self.token = secrets.token_urlsafe(16)
|
|
self._runner: web.AppRunner | None = None
|
|
self._app = web.Application()
|
|
self._app.add_routes(
|
|
[
|
|
web.get("/", self._index),
|
|
web.get("/status", self._status),
|
|
web.options("/{token}/live.{ext}", self._preflight),
|
|
web.route("*", "/{token}/live.{ext}", self._live),
|
|
web.route("*", "/{token}/hls/{name}", self._hls),
|
|
]
|
|
)
|
|
|
|
def new_token(self) -> str:
|
|
"""Rotate the stream path. Every session gets its own; the old one dies."""
|
|
self.token = secrets.token_urlsafe(16)
|
|
return self.token
|
|
|
|
def _authorized(self, request: web.Request) -> bool:
|
|
return (
|
|
_is_lan(request.remote)
|
|
and secrets.compare_digest(request.match_info.get("token", ""), self.token)
|
|
)
|
|
|
|
async def start(self, port: int) -> None:
|
|
if self._runner is not None and self.port == port:
|
|
return
|
|
await self.stop()
|
|
self._runner = web.AppRunner(self._app, access_log=None)
|
|
await self._runner.setup()
|
|
site = web.TCPSite(self._runner, "0.0.0.0", port, reuse_address=True)
|
|
await site.start()
|
|
self.port = port
|
|
log.info("http stream server on 0.0.0.0:%d", port)
|
|
|
|
async def stop(self) -> None:
|
|
if self._runner is not None:
|
|
await self._runner.cleanup()
|
|
self._runner = None
|
|
self.port = 0
|
|
|
|
def url_path(self, container: str) -> str:
|
|
return f"/{self.token}" + PATHS.get(container, "/live.webm")
|
|
|
|
def hls_path(self) -> str:
|
|
return f"/{self.token}/hls/index.m3u8"
|
|
|
|
async def _index(self, request: web.Request) -> web.Response:
|
|
return web.Response(text="screencast\n")
|
|
|
|
async def _status(self, request: web.Request) -> web.Response:
|
|
# Whether this desktop is casting is nobody else's business.
|
|
if request.remote not in ("127.0.0.1", "::1"):
|
|
raise web.HTTPNotFound()
|
|
f = self.fanout
|
|
return web.json_response(
|
|
{
|
|
"streaming": f is not None,
|
|
"container": f.container if f else "",
|
|
"clients": f.clients if f else 0,
|
|
"bytes": f.bytes_out if f else 0,
|
|
}
|
|
)
|
|
|
|
async def _preflight(self, request: web.Request) -> web.Response:
|
|
return web.Response(
|
|
status=204,
|
|
headers={
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Methods": "GET, HEAD, OPTIONS",
|
|
"Access-Control-Allow-Headers": "*",
|
|
},
|
|
)
|
|
|
|
async def _hls(self, request: web.Request) -> web.StreamResponse:
|
|
if not self._authorized(request):
|
|
raise web.HTTPNotFound()
|
|
name = request.match_info["name"]
|
|
if self.hls_dir is None or "/" in name or name.startswith("."):
|
|
raise web.HTTPNotFound()
|
|
path = self.hls_dir / name
|
|
if not path.is_file():
|
|
raise web.HTTPNotFound()
|
|
ctype = "application/vnd.apple.mpegurl" if name.endswith(".m3u8") else "video/mp2t"
|
|
return web.FileResponse(
|
|
path,
|
|
headers={"Content-Type": ctype, "Cache-Control": "no-cache",
|
|
"Access-Control-Allow-Origin": "*"},
|
|
)
|
|
|
|
async def _live(self, request: web.Request) -> web.StreamResponse:
|
|
if not self._authorized(request):
|
|
raise web.HTTPNotFound()
|
|
f = self.fanout
|
|
if f is None:
|
|
raise web.HTTPServiceUnavailable(text="nothing is being cast")
|
|
if request.method != "HEAD" and f.clients >= MAX_CLIENTS:
|
|
log.warning("stream: refusing %s, already %d clients", request.remote, f.clients)
|
|
raise web.HTTPServiceUnavailable(text="too many clients")
|
|
headers = {
|
|
"Content-Type": CONTENT_TYPES.get(f.container, "application/octet-stream"),
|
|
# The Cast receiver is a web page playing our URL through the media
|
|
# stack: without these it never gets past "loading". Wide-open CORS
|
|
# is only safe because the path itself is the secret - a page that
|
|
# cannot guess the token cannot read the stream.
|
|
"Access-Control-Allow-Origin": "*",
|
|
"Access-Control-Allow-Methods": "GET, HEAD, OPTIONS",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Expose-Headers": "*",
|
|
"Cache-Control": "no-cache, no-store",
|
|
"Pragma": "no-cache",
|
|
# DLNA renderers refuse a stream that does not say it is one.
|
|
"transferMode.dlna.org": "Streaming",
|
|
"contentFeatures.dlna.org": "DLNA.ORG_OP=00;DLNA.ORG_CI=0;"
|
|
"DLNA.ORG_FLAGS=8D500000000000000000000000000000",
|
|
"Server": "Linux/1.0 UPnP/1.0 screencast/1.0",
|
|
"Accept-Ranges": "none",
|
|
}
|
|
resp = web.StreamResponse(status=200, headers=headers)
|
|
resp.enable_chunked_encoding()
|
|
await resp.prepare(request)
|
|
if request.method == "HEAD":
|
|
return resp
|
|
|
|
reader = f.subscribe()
|
|
peer = request.remote or "?"
|
|
log.info("stream client %s connected (%s)", peer, f.container)
|
|
if self.on_client:
|
|
self.on_client(peer, True)
|
|
# Nagle would hold a small write back waiting for company; on a live
|
|
# stream that is latency for nothing.
|
|
with contextlib.suppress(Exception):
|
|
request.transport.get_extra_info("socket").setsockopt(
|
|
socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
with contextlib.suppress(Exception):
|
|
request.transport.set_write_buffer_limits(high=128 << 10)
|
|
try:
|
|
if f.header:
|
|
await resp.write(f.header)
|
|
while True:
|
|
chunk = await reader.get()
|
|
if chunk is None:
|
|
break
|
|
await resp.write(chunk)
|
|
except (ConnectionResetError, asyncio.CancelledError):
|
|
pass
|
|
except Exception as exc: # noqa: BLE001 - a dropped receiver must not kill the server
|
|
log.debug("stream client %s error: %s", peer, exc)
|
|
finally:
|
|
f.unsubscribe(reader)
|
|
log.info("stream client %s gone", peer)
|
|
if self.on_client:
|
|
self.on_client(peer, False)
|
|
return resp
|