NetDesk 1.5.0 — source release
This commit is contained in:
commit
40318aafdd
227 files changed
+80798
No files matched your search
@@ -0,0 +1,82 @@
|
||||
[package]
|
||||
name = "netdesk-core"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
description = "NetDesk native remote-desktop core (signaling, E2E crypto, transport, codec). Developed by Akronlogic."
|
||||
|
||||
[lib]
|
||||
name = "netdesk_core"
|
||||
crate-type = ["lib", "staticlib"] # frb static-links core (rlib); the vestigial netdesk_core cdylib overflowed MinGW export ordinals once webrtc was linked
|
||||
|
||||
[dependencies]
|
||||
tokio = { workspace = true }
|
||||
tokio-tungstenite = { workspace = true }
|
||||
rustls = { workspace = true }
|
||||
futures-util = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
hmac = { workspace = true }
|
||||
# TOTP (RFC 6238) 2FA on the native host — HMAC-SHA1 over the time counter.
|
||||
sha1 = "0.10"
|
||||
pbkdf2 = { workspace = true }
|
||||
p256 = { workspace = true }
|
||||
aes-gcm = { workspace = true }
|
||||
base64 = { workspace = true }
|
||||
hex = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
rand = { workspace = true }
|
||||
openh264 = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
# Host-engine permission gates + approve flow log through the platform's logger
|
||||
# (env_logger in the Windows host, logcat via android_logger on Android).
|
||||
log = "0.4"
|
||||
# Session recording: repackage the received H.264 into a playable .mp4 (no re-encode).
|
||||
mp4 = "0.14"
|
||||
webrtc = "0.17.2"
|
||||
async-trait = "0.1"
|
||||
bytes = "1"
|
||||
|
||||
# Windows-only: Media Foundation hardware H.264 decode (DXVA via the system
|
||||
# H264 decoder MFT). Guarded behind cfg(windows); the core still builds on the
|
||||
# other four targets with openh264 software decode only. The "windows" crate
|
||||
# links the GNU import libraries from windows-targets, so it builds on the
|
||||
# x86_64-pc-windows-gnu toolchain used for the portable build.
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
# VP9 software decode (libvpx) for constrained links. Windows-only for now: the
|
||||
# vendored static libvpx targets the x86_64-pc-windows-gnu ABI.
|
||||
netdesk-vpx = { path = "../vpx" }
|
||||
windows = { version = "0.58", features = [
|
||||
"Win32_Foundation",
|
||||
"Win32_System_Com",
|
||||
"Win32_Media_MediaFoundation",
|
||||
"Win32_Graphics_Direct3D",
|
||||
"Win32_Graphics_Direct3D11",
|
||||
"Win32_Graphics_Dxgi",
|
||||
"Win32_Graphics_Dxgi_Common",
|
||||
] }
|
||||
|
||||
# macOS / iOS: VideoToolbox hardware H.264 decode (decode_vt.rs). The module
|
||||
# links the system CoreMedia / CoreVideo / VideoToolbox frameworks directly and
|
||||
# uses core-foundation only to build the pixel-buffer attribute dictionary that
|
||||
# bridges into those calls. cfg-gated behind the Apple targets, so the Windows,
|
||||
# Linux and Android builds never reference these crates.
|
||||
[target.'cfg(any(target_os = "macos", target_os = "ios"))'.dependencies]
|
||||
core-foundation = "0.10"
|
||||
core-foundation-sys = "0.8"
|
||||
|
||||
# Android: MediaCodec hardware H.264 decode (decode_mediacodec.rs) via the NDK
|
||||
# Media C ABI in libmediandk — the same ABI the ndk-sys crate binds. cfg-gated
|
||||
# behind target_os = "android".
|
||||
[target.'cfg(target_os = "android")'.dependencies]
|
||||
ndk-sys = "0.6"
|
||||
|
||||
# Linux: VAAPI hardware H.264 decode (decode_vaapi.rs) through FFmpeg /
|
||||
# libavcodec. This pulls libav* in via ffmpeg-next and needs the system FFmpeg
|
||||
# (>= 4) development libraries present at build time — on Linux only; the crate
|
||||
# and its build script are never compiled for the other targets.
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
ffmpeg-next = "9.0"
|
||||
@@ -0,0 +1,214 @@
|
||||
//! White-label support — a renamed executable self-configures. A reseller /
|
||||
//! MSP can ship a pre-pointed client WITHOUT rebuilding, just by renaming the
|
||||
//! exe so its name encodes the relay + registration key (RustDesk's
|
||||
//! `rustdesk-host=...,key=....exe` trick). The viewer reads its OWN filename at
|
||||
//! startup and writes the encoded settings into config.json.
|
||||
//!
|
||||
//! Format (filename stem): `<BrandName>-<k=v,k=v,...>`, e.g.
|
||||
//! `NetDesk-server=relay.corp.example,key=CORP1.exe`
|
||||
//! `AcmeRemote-server=vpn.acme.io,port=8443,tls=1,name=Acme,id=100200300.exe`
|
||||
//! Keys: `server` (bare host — filenames can't hold `://` or `:`, so the scheme
|
||||
//! and `/ws` path are reconstructed), `port`, `tls` (0/1), `key` (registration
|
||||
//! key), `id`, `name`. The part before the first recognised `key=` is the brand
|
||||
//! name (empty / "NetDesk" means default). Pure + deterministic for unit tests.
|
||||
|
||||
/// Settings decoded from an executable's filename.
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct FilenameConfig {
|
||||
/// Product name from the filename prefix (None = default "NetDesk").
|
||||
pub brand: Option<String>,
|
||||
/// Full relay URL, reconstructed from the bare host (+ port/tls).
|
||||
pub server: Option<String>,
|
||||
pub id: Option<String>,
|
||||
pub key: Option<String>,
|
||||
pub name: Option<String>,
|
||||
}
|
||||
|
||||
impl FilenameConfig {
|
||||
/// True when the filename carried any self-config settings.
|
||||
pub fn has_config(&self) -> bool {
|
||||
self.server.is_some() || self.id.is_some() || self.key.is_some() || self.name.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a filename STEM (no extension) into its encoded settings.
|
||||
pub fn parse_filename_config(stem: &str) -> FilenameConfig {
|
||||
const KEYS: &[&str] = &["server=", "id=", "key=", "name=", "port=", "tls="];
|
||||
// Find where the k=v block starts (the earliest recognised key).
|
||||
let start = KEYS.iter().filter_map(|k| stem.find(k)).min();
|
||||
let Some(start) = start else {
|
||||
return FilenameConfig::default(); // no config encoded (e.g. "NetDesk-Portable")
|
||||
};
|
||||
|
||||
let brand_raw = stem[..start].trim_end_matches('-').trim();
|
||||
let brand = if brand_raw.is_empty() || brand_raw.eq_ignore_ascii_case("netdesk") {
|
||||
None
|
||||
} else {
|
||||
Some(brand_raw.to_string())
|
||||
};
|
||||
|
||||
let mut out = FilenameConfig {
|
||||
brand,
|
||||
..Default::default()
|
||||
};
|
||||
let mut host: Option<String> = None;
|
||||
let mut port: Option<String> = None;
|
||||
let mut tls = true;
|
||||
for part in stem[start..].split(',') {
|
||||
let Some((k, v)) = part.split_once('=') else {
|
||||
continue;
|
||||
};
|
||||
let v = v.trim().to_string();
|
||||
if v.is_empty() {
|
||||
continue;
|
||||
}
|
||||
match k.trim() {
|
||||
"server" => host = Some(v),
|
||||
"port" => port = Some(v),
|
||||
"tls" => tls = !(v == "0" || v.eq_ignore_ascii_case("false") || v.eq_ignore_ascii_case("no")),
|
||||
"id" => out.id = Some(v),
|
||||
"key" => out.key = Some(v),
|
||||
"name" => out.name = Some(v),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// Rebuild a full ws URL from the bare host (filenames can't contain :// or :).
|
||||
if let Some(h) = host {
|
||||
let scheme = if tls { "wss" } else { "ws" };
|
||||
let hostport = match port {
|
||||
Some(p) => format!("{h}:{p}"),
|
||||
None => h,
|
||||
};
|
||||
out.server = Some(format!("{scheme}://{hostport}/ws"));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn plain_name_has_no_config() {
|
||||
assert_eq!(parse_filename_config("NetDesk"), FilenameConfig::default());
|
||||
// "NetDesk-Portable" must NOT be mistaken for a config.
|
||||
assert_eq!(parse_filename_config("NetDesk-Portable"), FilenameConfig::default());
|
||||
assert!(!parse_filename_config("NetDesk").has_config());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decodes_server_and_key_default_brand() {
|
||||
let c = parse_filename_config("NetDesk-server=relay.corp.example,key=CORP1");
|
||||
assert_eq!(c.server.as_deref(), Some("wss://relay.corp.example/ws"));
|
||||
assert_eq!(c.key.as_deref(), Some("CORP1"));
|
||||
assert_eq!(c.brand, None);
|
||||
assert!(c.has_config());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decodes_brand_port_tls_name_id() {
|
||||
let c = parse_filename_config("AcmeRemote-server=vpn.acme.io,port=8443,tls=1,name=Acme HQ,id=100200300");
|
||||
assert_eq!(c.brand.as_deref(), Some("AcmeRemote"));
|
||||
assert_eq!(c.server.as_deref(), Some("wss://vpn.acme.io:8443/ws"));
|
||||
assert_eq!(c.name.as_deref(), Some("Acme HQ"));
|
||||
assert_eq!(c.id.as_deref(), Some("100200300"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tls_zero_gives_ws_scheme() {
|
||||
let c = parse_filename_config("NetDesk-server=192.168.1.10,tls=0");
|
||||
assert_eq!(c.server.as_deref(), Some("ws://192.168.1.10/ws"));
|
||||
}
|
||||
}
|
||||
|
||||
/// A runtime brand manifest (`brand.json` dropped next to the exe) for full
|
||||
/// white-labeling without a rebuild: product name, vendor, accent colors, a
|
||||
/// logo PNG (filename, loaded from next to the exe), and a default relay.
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct BrandManifest {
|
||||
pub name: String,
|
||||
pub vendor: String,
|
||||
/// Primary accent as `#RRGGBB` (empty = keep the default).
|
||||
pub accent: String,
|
||||
pub accent_hi: String,
|
||||
/// Logo PNG filename, resolved next to the executable.
|
||||
pub logo: String,
|
||||
/// Default relay URL (same role as the filename `server=`).
|
||||
pub server: String,
|
||||
}
|
||||
|
||||
impl BrandManifest {
|
||||
/// The accent parsed to RGB, if a valid `#RRGGBB` was given.
|
||||
pub fn accent_rgb(&self) -> Option<(u8, u8, u8)> {
|
||||
parse_hex_rgb(&self.accent)
|
||||
}
|
||||
pub fn accent_hi_rgb(&self) -> Option<(u8, u8, u8)> {
|
||||
parse_hex_rgb(&self.accent_hi)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse `brand.json`. Unknown/missing keys default to empty (a partial manifest
|
||||
/// is fine — only the fields present override). Invalid JSON -> all defaults.
|
||||
pub fn parse_brand_manifest(json: &str) -> BrandManifest {
|
||||
let v: serde_json::Value = serde_json::from_str(json).unwrap_or(serde_json::Value::Null);
|
||||
let s = |k: &str| {
|
||||
v.get(k)
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string()
|
||||
};
|
||||
BrandManifest {
|
||||
name: s("name"),
|
||||
vendor: s("vendor"),
|
||||
accent: s("accent"),
|
||||
accent_hi: s("accent_hi"),
|
||||
logo: s("logo"),
|
||||
server: s("server"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse `#RRGGBB` (or `RRGGBB`) into (r, g, b).
|
||||
pub fn parse_hex_rgb(s: &str) -> Option<(u8, u8, u8)> {
|
||||
let s = s.trim().trim_start_matches('#');
|
||||
if s.len() != 6 || !s.bytes().all(|b| b.is_ascii_hexdigit()) {
|
||||
return None;
|
||||
}
|
||||
let r = u8::from_str_radix(&s[0..2], 16).ok()?;
|
||||
let g = u8::from_str_radix(&s[2..4], 16).ok()?;
|
||||
let b = u8::from_str_radix(&s[4..6], 16).ok()?;
|
||||
Some((r, g, b))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod manifest_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn hex_parses_and_rejects() {
|
||||
assert_eq!(parse_hex_rgb("#22C55E"), Some((34, 197, 94)));
|
||||
assert_eq!(parse_hex_rgb("22c55e"), Some((34, 197, 94)));
|
||||
assert_eq!(parse_hex_rgb("#1E88E5"), Some((30, 136, 229)));
|
||||
assert_eq!(parse_hex_rgb("#12345"), None); // too short
|
||||
assert_eq!(parse_hex_rgb("#zzzzzz"), None); // not hex
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_parses_fields() {
|
||||
let json = r##"{"name":"AcmeRemote","vendor":"Acme Corp","accent":"#1E88E5",
|
||||
"logo":"acme.png","server":"wss://relay.acme.io/ws"}"##;
|
||||
let m = parse_brand_manifest(json);
|
||||
assert_eq!(m.name, "AcmeRemote");
|
||||
assert_eq!(m.vendor, "Acme Corp");
|
||||
assert_eq!(m.accent_rgb(), Some((30, 136, 229)));
|
||||
assert_eq!(m.logo, "acme.png");
|
||||
assert_eq!(m.server, "wss://relay.acme.io/ws");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_and_invalid_manifests() {
|
||||
assert_eq!(parse_brand_manifest(r#"{"name":"OnlyName"}"#).name, "OnlyName");
|
||||
assert_eq!(parse_brand_manifest("not json").name, "");
|
||||
assert_eq!(parse_brand_manifest(r#"{"name":"X"}"#).accent_rgb(), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
//! CIDR allow-list for encrypted direct-IP access. An empty list allows any
|
||||
//! peer (the admin opts into restriction by listing CIDRs); otherwise the peer
|
||||
//! IP must fall inside one of the comma/space/newline-separated CIDRs. Pure +
|
||||
//! deterministic, so it unit-tests without a socket. IPv4 and IPv6.
|
||||
|
||||
use std::net::IpAddr;
|
||||
|
||||
/// Is `ip` allowed by `list`? An empty/whitespace list allows everything (the
|
||||
/// permissive default); otherwise `ip` must match one entry. A bare address
|
||||
/// (no `/bits`) matches exactly (host route, /32 or /128).
|
||||
pub fn allowed(ip: &IpAddr, list: &str) -> bool {
|
||||
let mut any = false;
|
||||
for e in list.split(|c| c == ',' || c == ' ' || c == '\n' || c == '\t') {
|
||||
let e = e.trim();
|
||||
if e.is_empty() {
|
||||
continue;
|
||||
}
|
||||
any = true;
|
||||
if cidr_match(ip, e) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
!any // no entries -> allow all
|
||||
}
|
||||
|
||||
fn cidr_match(ip: &IpAddr, cidr: &str) -> bool {
|
||||
let (net_str, bits) = match cidr.split_once('/') {
|
||||
Some((n, b)) => (n, b.trim().parse::<u32>().ok()),
|
||||
None => (cidr, None),
|
||||
};
|
||||
let net: IpAddr = match net_str.trim().parse() {
|
||||
Ok(a) => a,
|
||||
Err(_) => return false,
|
||||
};
|
||||
match (ip, net) {
|
||||
(IpAddr::V4(ip), IpAddr::V4(net)) => {
|
||||
let bits = bits.unwrap_or(32).min(32);
|
||||
let mask: u32 = if bits == 0 { 0 } else { u32::MAX << (32 - bits) };
|
||||
(u32::from(*ip) & mask) == (u32::from(net) & mask)
|
||||
}
|
||||
(IpAddr::V6(ip), IpAddr::V6(net)) => {
|
||||
let bits = bits.unwrap_or(128).min(128);
|
||||
let mask: u128 = if bits == 0 { 0 } else { u128::MAX << (128 - bits) };
|
||||
(u128::from(*ip) & mask) == (u128::from(net) & mask)
|
||||
}
|
||||
_ => false, // v4 vs v6 never match
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::net::IpAddr;
|
||||
|
||||
fn ip(s: &str) -> IpAddr {
|
||||
s.parse().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_list_allows_all() {
|
||||
assert!(allowed(&ip("8.8.8.8"), ""));
|
||||
assert!(allowed(&ip("10.0.0.1"), " "));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_inside_cidr_blocks() {
|
||||
assert!(allowed(&ip("10.1.2.3"), "10.0.0.0/8"));
|
||||
assert!(!allowed(&ip("11.1.2.3"), "10.0.0.0/8"));
|
||||
assert!(allowed(&ip("192.168.1.50"), "192.168.1.0/24"));
|
||||
assert!(!allowed(&ip("192.168.2.50"), "192.168.1.0/24"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_entries_and_exact_host() {
|
||||
let list = "127.0.0.1, 10.0.0.0/8 192.168.0.0/16";
|
||||
assert!(allowed(&ip("127.0.0.1"), list)); // exact host
|
||||
assert!(allowed(&ip("10.9.9.9"), list));
|
||||
assert!(allowed(&ip("192.168.5.5"), list));
|
||||
assert!(!allowed(&ip("172.16.0.1"), list));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v4_and_v6_do_not_cross_match() {
|
||||
assert!(!allowed(&ip("10.0.0.1"), "::1/128"));
|
||||
assert!(allowed(&ip("::1"), "::1/128"));
|
||||
assert!(allowed(&ip("fe80::1"), "fe80::/16"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
//! End-to-end crypto matching the web viewer exactly, so the native client
|
||||
//! authenticates to and decrypts from the same hosts:
|
||||
//! verifier = hex( PBKDF2-HMAC-SHA256(password, salt, iters, 32) )
|
||||
//! proof = hex( SHA256( nonce_str || verifier ) )
|
||||
//! E2E key = SHA256( ECDH_P256(host_pub, our_priv).x || "e2e" ) -> AES-256-GCM
|
||||
//! frame = AES-GCM(iv = bytes[0..12], ct = bytes[12..])
|
||||
|
||||
use crate::error::{NdError, Result};
|
||||
use aes_gcm::aead::Aead;
|
||||
use aes_gcm::{Aes256Gcm, Key, KeyInit, Nonce};
|
||||
use base64::{engine::general_purpose::STANDARD as B64, Engine};
|
||||
use p256::ecdh::EphemeralSecret;
|
||||
use p256::elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
|
||||
use p256::{EncodedPoint, PublicKey};
|
||||
use rand::RngCore;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// Password proof the host verifies. `salt_b64` and `nonce` come from the challenge.
|
||||
pub fn password_proof(password: &str, salt_b64: &str, nonce: &str, iters: u32) -> Result<String> {
|
||||
let salt = B64
|
||||
.decode(salt_b64)
|
||||
.map_err(|e| NdError::Crypto(format!("salt b64: {e}")))?;
|
||||
let mut dk = [0u8; 32];
|
||||
pbkdf2::pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, iters, &mut dk);
|
||||
let verifier = hex::encode(dk);
|
||||
let mut h = Sha256::new();
|
||||
h.update(nonce.as_bytes());
|
||||
h.update(verifier.as_bytes());
|
||||
Ok(hex::encode(h.finalize()))
|
||||
}
|
||||
|
||||
/// Channel-bound password proof (MITM hardening). Binds the host's ephemeral
|
||||
/// ECDH public key into the proof: `hex(SHA256(nonce || verifier || host_pub))`.
|
||||
/// A malicious relay that swaps the host's `pub` in the challenge can't produce
|
||||
/// a matching proof (it lacks the verifier), and the host verifies against ITS
|
||||
/// OWN pub — so a key swap is detected and auth fails. `host_pub_b64` is the
|
||||
/// exact base64 string the host put in `session.challenge` `pub`.
|
||||
pub fn password_proof_bound(
|
||||
password: &str,
|
||||
salt_b64: &str,
|
||||
nonce: &str,
|
||||
iters: u32,
|
||||
host_pub_b64: &str,
|
||||
) -> Result<String> {
|
||||
let salt = B64
|
||||
.decode(salt_b64)
|
||||
.map_err(|e| NdError::Crypto(format!("salt b64: {e}")))?;
|
||||
let mut dk = [0u8; 32];
|
||||
pbkdf2::pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, iters, &mut dk);
|
||||
let verifier = hex::encode(dk);
|
||||
Ok(bound_proof_from_verifier(nonce, &verifier, host_pub_b64))
|
||||
}
|
||||
|
||||
/// Shared core of the bound proof, from an already-computed `verifier_hex`.
|
||||
/// Used by the viewer (via [`password_proof_bound`]) and the host verify side
|
||||
/// (which stores `verifier_hex`), so both compute the digest identically.
|
||||
pub fn bound_proof_from_verifier(nonce: &str, verifier_hex: &str, host_pub_b64: &str) -> String {
|
||||
let mut h = Sha256::new();
|
||||
h.update(nonce.as_bytes());
|
||||
h.update(verifier_hex.as_bytes());
|
||||
h.update(host_pub_b64.as_bytes());
|
||||
hex::encode(h.finalize())
|
||||
}
|
||||
|
||||
/// HOST side: verify a viewer's channel-bound proof against the host's OWN
|
||||
/// ephemeral pubkey. `host_pub_b64` is what the host sent in the challenge.
|
||||
/// A relay that swapped the pub can't satisfy this.
|
||||
pub fn verify_proof_bound(
|
||||
nonce: &str,
|
||||
verifier_hex: &str,
|
||||
host_pub_b64: &str,
|
||||
proof: &str,
|
||||
) -> bool {
|
||||
let want = bound_proof_from_verifier(nonce, verifier_hex, host_pub_b64);
|
||||
ct_eq(want.as_bytes(), proof.to_lowercase().as_bytes())
|
||||
}
|
||||
|
||||
/// Our ephemeral ECDH keypair + the derived AES-256-GCM session key.
|
||||
pub struct E2e {
|
||||
pub our_pub_b64: String,
|
||||
pub cipher: Aes256Gcm,
|
||||
}
|
||||
|
||||
/// Derive the end-to-end key from the host's public key (base64 raw, 65-byte point).
|
||||
pub fn derive_e2e(host_pub_b64: &str) -> Result<E2e> {
|
||||
let host_raw = B64
|
||||
.decode(host_pub_b64)
|
||||
.map_err(|e| NdError::Crypto(format!("host pub b64: {e}")))?;
|
||||
let host_pt = EncodedPoint::from_bytes(&host_raw)
|
||||
.map_err(|e| NdError::Crypto(format!("host pub point: {e}")))?;
|
||||
let host_pub = Option::from(PublicKey::from_encoded_point(&host_pt))
|
||||
.ok_or_else(|| NdError::Crypto("invalid host public key".into()))?;
|
||||
|
||||
let secret = EphemeralSecret::random(&mut rand::rngs::OsRng);
|
||||
let our_pub = secret.public_key().to_encoded_point(false); // uncompressed 0x04||X||Y
|
||||
let our_pub_b64 = B64.encode(our_pub.as_bytes());
|
||||
|
||||
let shared = secret.diffie_hellman(&host_pub);
|
||||
let mut h = Sha256::new();
|
||||
h.update(shared.raw_secret_bytes()); // the x-coordinate (32 bytes), as WebCrypto deriveBits
|
||||
h.update(b"e2e");
|
||||
let key_bytes = h.finalize();
|
||||
let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&key_bytes));
|
||||
|
||||
Ok(E2e {
|
||||
our_pub_b64,
|
||||
cipher,
|
||||
})
|
||||
}
|
||||
|
||||
/// Decrypt one binary frame (iv is the first 12 bytes; the rest is ciphertext+tag).
|
||||
pub fn decrypt_frame(cipher: &Aes256Gcm, data: &[u8]) -> Result<Vec<u8>> {
|
||||
if data.len() < 12 + 16 {
|
||||
return Err(NdError::Crypto("frame too short".into()));
|
||||
}
|
||||
let nonce = Nonce::from_slice(&data[..12]);
|
||||
cipher
|
||||
.decrypt(nonce, &data[12..])
|
||||
.map_err(|_| NdError::Crypto("GCM decrypt failed".into()))
|
||||
}
|
||||
|
||||
/// Encrypt one binary frame for the wire: `iv(12) || AES-GCM(plaintext)`.
|
||||
/// The HOST side of [`decrypt_frame`] — the Android host uses this to seal its
|
||||
/// H.264 access units exactly the way the Python engine does (12-byte random
|
||||
/// nonce prefix), so the same web/desktop/native viewers decrypt them.
|
||||
pub fn encrypt_frame(cipher: &Aes256Gcm, plaintext: &[u8]) -> Result<Vec<u8>> {
|
||||
let mut iv = [0u8; 12];
|
||||
rand::rngs::OsRng.fill_bytes(&mut iv);
|
||||
let nonce = Nonce::from_slice(&iv);
|
||||
let ct = cipher
|
||||
.encrypt(nonce, plaintext)
|
||||
.map_err(|_| NdError::Crypto("GCM encrypt failed".into()))?;
|
||||
let mut out = Vec::with_capacity(12 + ct.len());
|
||||
out.extend_from_slice(&iv);
|
||||
out.extend_from_slice(&ct);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Constant-time byte compare (mirrors Python's `secrets.compare_digest`).
|
||||
fn ct_eq(a: &[u8], b: &[u8]) -> bool {
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
let mut d = 0u8;
|
||||
for (x, y) in a.iter().zip(b.iter()) {
|
||||
d |= x ^ y;
|
||||
}
|
||||
d == 0
|
||||
}
|
||||
|
||||
/// HOST side of the auth handshake: verify a viewer's password proof.
|
||||
/// `proof` is what the viewer sent; it must equal `hex(SHA256(nonce || verifier))`
|
||||
/// where `verifier` is the host's stored `verifier_hex`. Matches the Python
|
||||
/// host's `Config.check_proof`.
|
||||
pub fn verify_proof(nonce: &str, verifier_hex: &str, proof: &str) -> bool {
|
||||
let mut h = Sha256::new();
|
||||
h.update(nonce.as_bytes());
|
||||
h.update(verifier_hex.as_bytes());
|
||||
let want = hex::encode(h.finalize());
|
||||
ct_eq(want.as_bytes(), proof.to_lowercase().as_bytes())
|
||||
}
|
||||
|
||||
/// HOST side of the E2E ECDH. Created when the host sends `session.challenge`
|
||||
/// (exposing [`our_pub_b64`] in the challenge), then consumed at `session.auth`
|
||||
/// via [`derive`] with the viewer's public key to yield the shared AES-256-GCM
|
||||
/// cipher. Mirror of [`derive_e2e`] (the viewer side): both sides compute
|
||||
/// `SHA256( ECDH_P256(x) || "e2e" )`, so they end up with the SAME key.
|
||||
pub struct HostE2e {
|
||||
pub our_pub_b64: String,
|
||||
secret: EphemeralSecret,
|
||||
}
|
||||
|
||||
impl HostE2e {
|
||||
/// Generate a fresh ephemeral keypair for one challenge.
|
||||
pub fn new() -> Self {
|
||||
let secret = EphemeralSecret::random(&mut rand::rngs::OsRng);
|
||||
let our_pub = secret.public_key().to_encoded_point(false); // 0x04||X||Y
|
||||
let our_pub_b64 = B64.encode(our_pub.as_bytes());
|
||||
HostE2e {
|
||||
our_pub_b64,
|
||||
secret,
|
||||
}
|
||||
}
|
||||
|
||||
/// Derive the AES-256-GCM session cipher from the viewer's public key
|
||||
/// (base64 raw 65-byte point). Consumes the keypair (used once).
|
||||
pub fn derive(self, viewer_pub_b64: &str) -> Result<Aes256Gcm> {
|
||||
let raw = B64
|
||||
.decode(viewer_pub_b64)
|
||||
.map_err(|e| NdError::Crypto(format!("viewer pub b64: {e}")))?;
|
||||
let pt = EncodedPoint::from_bytes(&raw)
|
||||
.map_err(|e| NdError::Crypto(format!("viewer pub point: {e}")))?;
|
||||
let viewer_pub = Option::from(PublicKey::from_encoded_point(&pt))
|
||||
.ok_or_else(|| NdError::Crypto("invalid viewer public key".into()))?;
|
||||
let shared = self.secret.diffie_hellman(&viewer_pub);
|
||||
let mut h = Sha256::new();
|
||||
h.update(shared.raw_secret_bytes());
|
||||
h.update(b"e2e");
|
||||
let key_bytes = h.finalize();
|
||||
Ok(Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&key_bytes)))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for HostE2e {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn host_viewer_e2e_agree_and_frames_roundtrip() {
|
||||
// Host sends its pub in the challenge; viewer derives against it.
|
||||
let host = HostE2e::new();
|
||||
let host_pub = host.our_pub_b64.clone();
|
||||
let viewer = derive_e2e(&host_pub).expect("viewer derive");
|
||||
// Host derives against the viewer's pub -> both keys must match.
|
||||
let host_cipher = host.derive(&viewer.our_pub_b64).expect("host derive");
|
||||
|
||||
// Host-sealed frame decrypts on the viewer (the real media direction).
|
||||
let au = b"\x00\x00\x00\x01fake annex-b access unit";
|
||||
let sealed = encrypt_frame(&host_cipher, au).expect("seal");
|
||||
let opened = decrypt_frame(&viewer.cipher, &sealed).expect("open");
|
||||
assert_eq!(&opened, au);
|
||||
|
||||
// And the reverse (viewer-sealed -> host opens), proving identical keys.
|
||||
let sealed2 = encrypt_frame(&viewer.cipher, au).expect("seal2");
|
||||
let opened2 = decrypt_frame(&host_cipher, &sealed2).expect("open2");
|
||||
assert_eq!(&opened2, au);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn totp_matches_python_and_rfc6238() {
|
||||
// RFC 6238 SHA-1 test secret ("12345678901234567890" base32). These are
|
||||
// the exact codes the shipping Python host's totp_now() produced — so
|
||||
// Rust ≡ Python ≡ any authenticator app.
|
||||
let s = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ";
|
||||
assert_eq!(totp_at(s, 59, 30, 6).unwrap(), "287082");
|
||||
assert_eq!(totp_at(s, 1_111_111_109, 30, 6).unwrap(), "081804");
|
||||
assert_eq!(totp_at(s, 1_234_567_890, 30, 6).unwrap(), "005924");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn totp_verify_window_and_reject() {
|
||||
let s = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ";
|
||||
// A non-numeric / empty code is always rejected.
|
||||
assert!(!totp_verify(s, "", 1));
|
||||
assert!(!totp_verify(s, "abcdef", 1));
|
||||
// The helper that drives verify accepts the matching code for a step and
|
||||
// rejects a wrong one (tested time-independently via totp_at parity).
|
||||
let now_code = totp_at(s, 1_000_000_020, 30, 6).unwrap();
|
||||
let wrong = if now_code == "000000" { "111111" } else { "000000" };
|
||||
assert_ne!(now_code, wrong);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bound_proof_known_answer_cross_language() {
|
||||
// Shared known-answer vector: Rust, the web viewer (JS sha256Hex), and
|
||||
// the Python host all compute hex(SHA256(nonce||verifier||host_pub)) over
|
||||
// the same UTF-8 string, so this one hex proves all three agree.
|
||||
// (Python produced the identical digest for these inputs.)
|
||||
let got = bound_proof_from_verifier("abc123", "deadbeef", "AAAA_host_pub_b64");
|
||||
assert_eq!(
|
||||
got,
|
||||
"ecc84429d061a1580268f829aa3da0d046f1d6aa8a1c6e86ad46c8deff5d835f"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bound_proof_detects_pub_swap() {
|
||||
let nonce = "0123456789abcdef0123456789abcdef";
|
||||
let password = "correct horse battery staple";
|
||||
// Salt + verifier MUST come from the same credentials (new_host_credentials
|
||||
// picks a random salt), so the viewer's bound proof and the host's verify
|
||||
// use the same verifier — exactly as a real host stores them.
|
||||
let (salt_b64, iters, verifier) = new_host_credentials(password);
|
||||
// The host's real ephemeral pub vs. a relay-swapped one.
|
||||
let real = HostE2e::new();
|
||||
let swapped = HostE2e::new();
|
||||
// A viewer binding the REAL pub verifies; one binding the swapped pub does not.
|
||||
let good = password_proof_bound(password, &salt_b64, nonce, iters, &real.our_pub_b64)
|
||||
.expect("bound proof");
|
||||
assert!(verify_proof_bound(nonce, &verifier, &real.our_pub_b64, &good));
|
||||
let swapped_proof =
|
||||
password_proof_bound(password, &salt_b64, nonce, iters, &swapped.our_pub_b64)
|
||||
.expect("bound proof");
|
||||
assert!(
|
||||
!verify_proof_bound(nonce, &verifier, &real.our_pub_b64, &swapped_proof),
|
||||
"a proof bound to a swapped pub must NOT verify against the real pub"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proof_verifies_against_viewer_proof() {
|
||||
let salt = gen_salt();
|
||||
let salt_b64 = B64.encode(salt);
|
||||
let iters = 200_000u32;
|
||||
let nonce = "0123456789abcdef0123456789abcdef";
|
||||
let verifier = verifier_hex("hunter2", &salt, iters);
|
||||
// Viewer computes the proof the way password_proof does.
|
||||
let proof = password_proof("hunter2", &salt_b64, nonce, iters).expect("proof");
|
||||
assert!(verify_proof(nonce, &verifier, &proof), "valid proof accepted");
|
||||
assert!(
|
||||
!verify_proof(nonce, &verifier, &"00".repeat(32)),
|
||||
"wrong proof rejected"
|
||||
);
|
||||
// Wrong password must not verify.
|
||||
let bad = password_proof("wrong", &salt_b64, nonce, iters).expect("bad proof");
|
||||
assert!(!verify_proof(nonce, &verifier, &bad), "wrong password rejected");
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Host unattended-password credentials. These mirror the Python host engine's
|
||||
// `make_verifier` / `gen_password` (netdesk_host.py) EXACTLY so a verifier the
|
||||
// native panel writes to config.json is accepted by the same engine:
|
||||
// verifier = hex( PBKDF2-HMAC-SHA256(password, salt, iters=200000, dklen=32) )
|
||||
// In config.json `salt` and `verifier` are stored hex-encoded, `iters` = 200000.
|
||||
// (The viewer's host-control module reuses these instead of pulling pbkdf2/rand
|
||||
// into the viewer crate.)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// PBKDF2-HMAC-SHA256 verifier (hex) over RAW salt bytes. Matches the host
|
||||
/// engine's `make_verifier` and the `verifier` field in config.json.
|
||||
pub fn verifier_hex(password: &str, salt: &[u8], iters: u32) -> String {
|
||||
let mut dk = [0u8; 32];
|
||||
pbkdf2::pbkdf2_hmac::<Sha256>(password.as_bytes(), salt, iters, &mut dk);
|
||||
hex::encode(dk)
|
||||
}
|
||||
|
||||
/// A fresh random 16-byte salt (matches the host `make_verifier` salt length).
|
||||
pub fn gen_salt() -> [u8; 16] {
|
||||
let mut s = [0u8; 16];
|
||||
rand::rngs::OsRng.fill_bytes(&mut s);
|
||||
s
|
||||
}
|
||||
|
||||
/// Fresh credentials for `password`: returns `(salt_hex, iters, verifier_hex)`
|
||||
/// ready to drop into config.json's `salt` / `iters` / `verifier` fields.
|
||||
pub fn new_password_record(password: &str) -> (String, u32, String) {
|
||||
let salt = gen_salt();
|
||||
let iters = 200_000u32;
|
||||
let verifier = verifier_hex(password, &salt, iters);
|
||||
(hex::encode(salt), iters, verifier)
|
||||
}
|
||||
|
||||
/// Fresh host credentials ready to drop into [`crate::host::HostConfig`]:
|
||||
/// `(salt_b64, iters, verifier_hex)` where `salt_b64` is base64 of the RAW salt
|
||||
/// (what `session.challenge` carries) and `verifier_hex` checks viewer proofs.
|
||||
pub fn new_host_credentials(password: &str) -> (String, u32, String) {
|
||||
let salt = gen_salt();
|
||||
let iters = 200_000u32;
|
||||
let verifier = verifier_hex(password, &salt, iters);
|
||||
(B64.encode(salt), iters, verifier)
|
||||
}
|
||||
|
||||
/// One-click unattended password, no look-alike characters (l/1, o/0) —
|
||||
/// identical alphabet to the host engine's `gen_password`. The alphabet is 32
|
||||
/// chars, so `u32 % 32` is perfectly uniform (2^32 is divisible by 32).
|
||||
pub fn gen_password(n: usize) -> String {
|
||||
const ALPHABET: &[u8] = b"abcdefghijkmnpqrstuvwxyz23456789";
|
||||
let mut rng = rand::rngs::OsRng;
|
||||
(0..n)
|
||||
.map(|_| ALPHABET[(rng.next_u32() as usize) % ALPHABET.len()] as char)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A fresh base32 TOTP secret (32 chars, A–Z + 2–7) — matches the host engine's
|
||||
/// 2FA secret generation (`string.ascii_uppercase + "234567"`).
|
||||
pub fn gen_totp_secret() -> String {
|
||||
const B32: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
|
||||
let mut rng = rand::rngs::OsRng;
|
||||
(0..32)
|
||||
.map(|_| B32[(rng.next_u32() as usize) % B32.len()] as char)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Decode an RFC 4648 base32 secret (A–Z, 2–7; padding + case tolerated), as the
|
||||
/// TOTP secret is stored. Returns None on an invalid character.
|
||||
fn b32_decode(s: &str) -> Option<Vec<u8>> {
|
||||
const A: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
|
||||
let mut bits: u32 = 0;
|
||||
let mut nbits: u32 = 0;
|
||||
let mut out = Vec::new();
|
||||
for c in s.trim().trim_end_matches('=').bytes() {
|
||||
let c = c.to_ascii_uppercase();
|
||||
let v = A.iter().position(|&x| x == c)? as u32;
|
||||
bits = (bits << 5) | v;
|
||||
nbits += 5;
|
||||
if nbits >= 8 {
|
||||
nbits -= 8;
|
||||
out.push((bits >> nbits) as u8);
|
||||
}
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// RFC 6238 TOTP code for a specific unix time (SHA-1, `step`s, `digits`).
|
||||
/// Exposed for tests; `totp_verify` is the runtime entry. Matches the Python
|
||||
/// host's `totp_now` exactly so the same authenticator app works everywhere.
|
||||
pub fn totp_at(secret: &str, unix_secs: u64, step: u64, digits: u32) -> Option<String> {
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha1::Sha1;
|
||||
let key = b32_decode(secret)?;
|
||||
let counter = (unix_secs / step).to_be_bytes();
|
||||
let mut mac = <Hmac<Sha1> as Mac>::new_from_slice(&key).ok()?;
|
||||
mac.update(&counter);
|
||||
let digest = mac.finalize().into_bytes();
|
||||
let off = (digest[19] & 0x0f) as usize;
|
||||
let code = ((u32::from(digest[off]) & 0x7f) << 24
|
||||
| u32::from(digest[off + 1]) << 16
|
||||
| u32::from(digest[off + 2]) << 8
|
||||
| u32::from(digest[off + 3]))
|
||||
% 10u32.pow(digits);
|
||||
Some(format!("{:0width$}", code, width = digits as usize))
|
||||
}
|
||||
|
||||
/// Verify a 6-digit TOTP code against `secret`, allowing ±`window` 30s steps for
|
||||
/// clock skew (window=1 ⇒ ±30s). Constant-time compare. Mirrors the Python
|
||||
/// host's `totp_verify`.
|
||||
pub fn totp_verify(secret: &str, code: &str, window: i64) -> bool {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let code = code.trim();
|
||||
if code.is_empty() || !code.bytes().all(|b| b.is_ascii_digit()) {
|
||||
return false;
|
||||
}
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0) as i64;
|
||||
for i in -window..=window {
|
||||
let t = (now + i * 30).max(0) as u64;
|
||||
if let Some(c) = totp_at(secret, t, 30, 6) {
|
||||
if ct_eq(c.as_bytes(), code.as_bytes()) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// A fresh 9-digit NetDesk host id, matching the engine's
|
||||
/// `str(secrets.randbelow(900000000) + 100000000)`.
|
||||
pub fn gen_host_id() -> String {
|
||||
let n = (rand::rngs::OsRng.next_u64() % 900_000_000) + 100_000_000;
|
||||
n.to_string()
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
//! H.264 decode. Phase 2 uses openh264 software decode (portable across all five
|
||||
//! targets) to turn the Annex-B access units from the host into RGB frames the UI
|
||||
//! can paint. Hardware decoders (Media Foundation / VideoToolbox / MediaCodec) are
|
||||
//! layered in later behind this same `Frame` output.
|
||||
|
||||
use crate::error::{NdError, Result};
|
||||
use openh264::decoder::Decoder;
|
||||
use openh264::formats::YUVSource;
|
||||
|
||||
/// One decoded frame, packed RGB8 (w*h*3 bytes), ready to upload to a texture.
|
||||
pub struct Frame {
|
||||
pub w: usize,
|
||||
pub h: usize,
|
||||
pub rgb: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Object-safe decode interface shared by every backend.
|
||||
///
|
||||
/// Both the portable openh264 software decoder ([`H264Decoder`]) and the Windows
|
||||
/// Media Foundation hardware decoder implement this, so the viewer can own a
|
||||
/// `Box<dyn VideoDecoder + Send>` and not care which one is running. The method
|
||||
/// contract is identical to [`H264Decoder::decode`]: feed one Annex-B access
|
||||
/// unit, get back a frame when the decoder emits a picture.
|
||||
pub trait VideoDecoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>>;
|
||||
}
|
||||
|
||||
pub struct H264Decoder {
|
||||
inner: Decoder,
|
||||
}
|
||||
|
||||
impl H264Decoder {
|
||||
pub fn new() -> Result<Self> {
|
||||
let inner =
|
||||
Decoder::new().map_err(|e| NdError::Other(anyhow::anyhow!("openh264 init: {e}")))?;
|
||||
Ok(Self { inner })
|
||||
}
|
||||
|
||||
/// Feed one Annex-B access unit (keyframes carry SPS+PPS+IDR; deltas carry a
|
||||
/// slice). Returns the decoded RGB frame when the decoder emits a picture.
|
||||
pub fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
match self.inner.decode(au) {
|
||||
Ok(Some(yuv)) => {
|
||||
let (w, h) = yuv.dimensions();
|
||||
let mut rgb = vec![0u8; w * h * 3];
|
||||
yuv.write_rgb8(&mut rgb);
|
||||
Ok(Some(Frame { w, h, rgb }))
|
||||
}
|
||||
Ok(None) => Ok(None),
|
||||
Err(e) => Err(NdError::Other(anyhow::anyhow!("decode: {e}"))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl VideoDecoder for H264Decoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
// Reuse the existing inherent method so both call sites stay in lockstep.
|
||||
H264Decoder::decode(self, au)
|
||||
}
|
||||
}
|
||||
|
||||
// Per-platform hardware decoders. Each lives in a sibling file and is compiled
|
||||
// only on its own OS; everywhere else it is cfg'd out and the openh264 software
|
||||
// decoder below is used. All implement the same `VideoDecoder` trait, so
|
||||
// `best_decoder` can hand back a `Box<dyn VideoDecoder + Send>` without the
|
||||
// viewer knowing or caring which backend is running.
|
||||
|
||||
// Windows: DXVA H.264 via Media Foundation.
|
||||
#[cfg(windows)]
|
||||
#[path = "decode_mf.rs"]
|
||||
mod mf;
|
||||
|
||||
// macOS / iOS: VideoToolbox (VTDecompressionSession).
|
||||
#[cfg(any(target_os = "macos", target_os = "ios"))]
|
||||
#[path = "decode_vt.rs"]
|
||||
mod vt;
|
||||
|
||||
// Android: MediaCodec (NDK AMediaCodec).
|
||||
#[cfg(target_os = "android")]
|
||||
#[path = "decode_mediacodec.rs"]
|
||||
mod mediacodec;
|
||||
|
||||
// Linux: VAAPI via FFmpeg/libavcodec.
|
||||
#[cfg(target_os = "linux")]
|
||||
#[path = "decode_vaapi.rs"]
|
||||
mod vaapi;
|
||||
|
||||
/// Last-resort decoder that drops every access unit. Only returned by
|
||||
/// [`best_decoder`] in the (practically impossible) case that even openh264
|
||||
/// fails to initialize, so callers never have to handle a `None` decoder.
|
||||
struct NullDecoder;
|
||||
|
||||
impl VideoDecoder for NullDecoder {
|
||||
fn decode(&mut self, _au: &[u8]) -> Result<Option<Frame>> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Pick the fastest decoder available on this platform.
|
||||
///
|
||||
/// Each OS gets its native hardware H.264 decoder, which is dramatically
|
||||
/// lower-CPU (and lower-power on laptops/phones) than the portable openh264
|
||||
/// software decoder:
|
||||
///
|
||||
/// * **Windows** — Media Foundation: DXVA2 GPU decode when the adapter supports
|
||||
/// it, otherwise Microsoft's optimized software MFT.
|
||||
/// * **macOS / iOS** — VideoToolbox (`VTDecompressionSession`) on the platform
|
||||
/// video engine.
|
||||
/// * **Android** — `MediaCodec` (`AMediaCodec`), i.e. the SoC video decoder.
|
||||
/// * **Linux** — VAAPI through FFmpeg/libavcodec (Intel/AMD GPUs via `/dev/dri`).
|
||||
///
|
||||
/// If the platform backend cannot initialize at runtime — no GPU support, a
|
||||
/// missing runtime, a headless box — this transparently falls back to the
|
||||
/// openh264 software decoder, which is always available on every target.
|
||||
///
|
||||
/// The returned decoder is `Send`, so the viewer's background network thread can
|
||||
/// own it directly.
|
||||
pub fn best_decoder() -> Box<dyn VideoDecoder + Send> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
match mf::MediaFoundationDecoder::new() {
|
||||
Ok(d) => {
|
||||
tracing::info!(
|
||||
"decode: Media Foundation H.264 ({})",
|
||||
if d.is_hardware() {
|
||||
"DXVA hardware"
|
||||
} else {
|
||||
"MS software MFT"
|
||||
}
|
||||
);
|
||||
return Box::new(d);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"decode: Media Foundation init failed ({e}); using openh264 software decode"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "macos", target_os = "ios"))]
|
||||
{
|
||||
match vt::VideoToolboxDecoder::new() {
|
||||
Ok(d) => {
|
||||
tracing::info!("decode: VideoToolbox H.264 (hardware)");
|
||||
return Box::new(d);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"decode: VideoToolbox init failed ({e}); using openh264 software decode"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_os = "android")]
|
||||
{
|
||||
match mediacodec::MediaCodecDecoder::new() {
|
||||
Ok(d) => {
|
||||
tracing::info!("decode: MediaCodec H.264 (hardware)");
|
||||
return Box::new(d);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"decode: MediaCodec init failed ({e}); using openh264 software decode"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
match vaapi::VaapiDecoder::new() {
|
||||
Ok(d) => {
|
||||
tracing::info!("decode: VAAPI H.264 (FFmpeg hardware)");
|
||||
return Box::new(d);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("decode: VAAPI init failed ({e}); using openh264 software decode");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match H264Decoder::new() {
|
||||
Ok(d) => Box::new(d),
|
||||
Err(e) => {
|
||||
tracing::error!("decode: openh264 init failed ({e}); video will not render");
|
||||
Box::new(NullDecoder)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// VP9 software decoder (libvpx via the netdesk-vpx crate). Windows-only for now
|
||||
/// — the vendored static libvpx is built for the x86_64-pc-windows-gnu ABI; the
|
||||
/// other platforms keep hardware H.264 until their libvpx is vendored too.
|
||||
#[cfg(windows)]
|
||||
pub struct Vp9Decoder {
|
||||
inner: netdesk_vpx::Decoder,
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
impl Vp9Decoder {
|
||||
pub fn new() -> Result<Self> {
|
||||
let inner = netdesk_vpx::Decoder::new(netdesk_vpx::VpxCodec::Vp9)
|
||||
.map_err(|e| NdError::Other(anyhow::anyhow!("vp9 decoder init: {e}")))?;
|
||||
Ok(Self { inner })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
impl VideoDecoder for Vp9Decoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
let frames = self
|
||||
.inner
|
||||
.decode_to_bgra(au)
|
||||
.map_err(|e| NdError::Other(anyhow::anyhow!("vp9 decode: {e}")))?;
|
||||
// One compressed packet yields 0 or 1 pictures; take it.
|
||||
let Some(f) = frames.into_iter().next() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let (w, h) = (f.width as usize, f.height as usize);
|
||||
let mut rgb = vec![0u8; w * h * 3];
|
||||
for i in 0..(w * h) {
|
||||
// BGRA -> RGB (drop alpha, swap B/R)
|
||||
rgb[i * 3] = f.bgra[i * 4 + 2];
|
||||
rgb[i * 3 + 1] = f.bgra[i * 4 + 1];
|
||||
rgb[i * 3 + 2] = f.bgra[i * 4];
|
||||
}
|
||||
Ok(Some(Frame { w, h, rgb }))
|
||||
}
|
||||
}
|
||||
|
||||
/// Pick a decoder for the given codec id (`proto::VideoCodec`): VP9 (1) uses the
|
||||
/// libvpx software decoder; anything else uses the platform H.264 decoder from
|
||||
/// [`best_decoder`]. The viewer calls this whenever a frame's codec tag changes
|
||||
/// (always at a keyframe), so it can follow the host switching codecs live.
|
||||
pub fn best_decoder_for(codec: u8) -> Box<dyn VideoDecoder + Send> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
if codec == crate::proto::VideoCodec::Vp9 as u8 {
|
||||
match Vp9Decoder::new() {
|
||||
Ok(d) => {
|
||||
tracing::info!("decode: VP9 (libvpx software)");
|
||||
return Box::new(d);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("decode: VP9 init failed ({e}); falling back to H.264");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = codec;
|
||||
best_decoder()
|
||||
}
|
||||
@@ -0,0 +1,512 @@
|
||||
//! Android hardware H.264 decode via `MediaCodec` (the NDK `AMediaCodec` API).
|
||||
//!
|
||||
//! This drives a byte-buffer-mode `AMediaCodec` H.264 decoder: the host's
|
||||
//! Annex-B access units go in as input buffers, and the decoder hands back
|
||||
//! YUV 4:2:0 output buffers (NV12 semi-planar or I420 planar depending on the
|
||||
//! device) which we convert to the same packed-RGB8 [`Frame`] the openh264 path
|
||||
//! produces, so the viewer is unchanged. On essentially all Android devices this
|
||||
//! binds to the SoC's hardware video decoder (OMX / Codec2).
|
||||
//!
|
||||
//! The module is only compiled on Android (see the `#[cfg(target_os = "android")]`
|
||||
//! gate in `decode.rs`). It talks to the stable NDK Media C ABI in `libmediandk`
|
||||
//! directly — the same ABI the `ndk`/`ndk-sys` crates wrap — so the FFI surface
|
||||
//! is self-contained and version-independent. If init fails at runtime,
|
||||
//! [`super::best_decoder`] falls back to the portable openh264 decoder.
|
||||
//!
|
||||
//! ## Flow
|
||||
//! * The codec must be configured with the stream's SPS/PPS as codec-specific
|
||||
//! data (`csd-0`/`csd-1`) before it can start, so it is created lazily on the
|
||||
//! first access unit carrying both (a keyframe). AUs seen before that are
|
||||
//! dropped (`Ok(None)`).
|
||||
//! * Each `decode` call queues the AU as one input buffer and dequeues at most
|
||||
//! one decoded output buffer; the host stream is low-latency with no
|
||||
//! B-frames, so the steady state is one frame in, one frame out.
|
||||
//! * The decoder reports the true picture rectangle through the crop keys on the
|
||||
//! output format, so the emitted `Frame` is the real display size (e.g.
|
||||
//! 1920x1080) even when the YUV planes are stride/slice-height padded.
|
||||
|
||||
use super::{Frame, VideoDecoder};
|
||||
use crate::error::{NdError, Result};
|
||||
|
||||
use core::ffi::{c_char, c_void};
|
||||
|
||||
use self::sys::*;
|
||||
|
||||
/// `MediaCodec` hardware H.264 decoder.
|
||||
pub struct MediaCodecDecoder {
|
||||
/// The codec; `null` until the first keyframe primes and starts it.
|
||||
codec: *mut AMediaCodec,
|
||||
sps: Vec<u8>,
|
||||
pps: Vec<u8>,
|
||||
pts_us: i64,
|
||||
}
|
||||
|
||||
impl MediaCodecDecoder {
|
||||
/// Probe that an H.264 decoder can be created on this device, then discard
|
||||
/// it (the real codec is configured lazily once the SPS/PPS arrive). Returns
|
||||
/// an error — so the caller falls back to openh264 — if no AVC decoder
|
||||
/// exists at all.
|
||||
pub fn new() -> Result<Self> {
|
||||
unsafe {
|
||||
let probe = AMediaCodec_createDecoderByType(MIME.as_ptr() as *const c_char);
|
||||
if probe.is_null() {
|
||||
return Err(other("no AMediaCodec AVC decoder available"));
|
||||
}
|
||||
AMediaCodec_delete(probe);
|
||||
}
|
||||
Ok(Self {
|
||||
codec: core::ptr::null_mut(),
|
||||
sps: Vec::new(),
|
||||
pps: Vec::new(),
|
||||
pts_us: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create, configure and start the codec from the cached SPS/PPS.
|
||||
fn prime(&mut self) -> Result<()> {
|
||||
self.teardown();
|
||||
unsafe {
|
||||
let codec = AMediaCodec_createDecoderByType(MIME.as_ptr() as *const c_char);
|
||||
if codec.is_null() {
|
||||
return Err(other("AMediaCodec_createDecoderByType returned null"));
|
||||
}
|
||||
|
||||
let fmt = AMediaFormat_new();
|
||||
if fmt.is_null() {
|
||||
AMediaCodec_delete(codec);
|
||||
return Err(other("AMediaFormat_new returned null"));
|
||||
}
|
||||
AMediaFormat_setString(fmt, KEY_MIME.as_ptr() as _, MIME.as_ptr() as _);
|
||||
// A nominal size is required; the decoder learns the real one from
|
||||
// the codec-specific data and reports it back via the output format.
|
||||
AMediaFormat_setInt32(fmt, KEY_WIDTH.as_ptr() as _, 1920);
|
||||
AMediaFormat_setInt32(fmt, KEY_HEIGHT.as_ptr() as _, 1080);
|
||||
|
||||
// csd-0 = SPS, csd-1 = PPS, each as a start-code-prefixed NAL.
|
||||
let csd0 = with_start_code(&self.sps);
|
||||
let csd1 = with_start_code(&self.pps);
|
||||
AMediaFormat_setBuffer(
|
||||
fmt,
|
||||
KEY_CSD0.as_ptr() as _,
|
||||
csd0.as_ptr() as *const c_void,
|
||||
csd0.len(),
|
||||
);
|
||||
AMediaFormat_setBuffer(
|
||||
fmt,
|
||||
KEY_CSD1.as_ptr() as _,
|
||||
csd1.as_ptr() as *const c_void,
|
||||
csd1.len(),
|
||||
);
|
||||
|
||||
let st = AMediaCodec_configure(
|
||||
codec,
|
||||
fmt,
|
||||
core::ptr::null_mut(), // no output Surface: byte-buffer mode
|
||||
core::ptr::null_mut(), // no crypto
|
||||
0,
|
||||
);
|
||||
AMediaFormat_delete(fmt);
|
||||
if st != AMEDIA_OK {
|
||||
AMediaCodec_delete(codec);
|
||||
return Err(other(&format!("AMediaCodec_configure failed ({st})")));
|
||||
}
|
||||
let st = AMediaCodec_start(codec);
|
||||
if st != AMEDIA_OK {
|
||||
AMediaCodec_delete(codec);
|
||||
return Err(other(&format!("AMediaCodec_start failed ({st})")));
|
||||
}
|
||||
self.codec = codec;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Queue one Annex-B access unit as an input buffer.
|
||||
unsafe fn queue_input(&mut self, au: &[u8]) {
|
||||
let idx = AMediaCodec_dequeueInputBuffer(self.codec, DEQUEUE_TIMEOUT_US);
|
||||
if idx < 0 {
|
||||
return;
|
||||
}
|
||||
let mut cap: usize = 0;
|
||||
let buf = AMediaCodec_getInputBuffer(self.codec, idx as usize, &mut cap);
|
||||
if buf.is_null() || au.len() > cap {
|
||||
// Still submit an empty buffer so the index is not leaked.
|
||||
let _ = AMediaCodec_queueInputBuffer(self.codec, idx as usize, 0, 0, self.pts_us, 0);
|
||||
return;
|
||||
}
|
||||
core::ptr::copy_nonoverlapping(au.as_ptr(), buf, au.len());
|
||||
let _ = AMediaCodec_queueInputBuffer(self.codec, idx as usize, 0, au.len(), self.pts_us, 0);
|
||||
self.pts_us += PTS_STEP_US;
|
||||
}
|
||||
|
||||
/// Dequeue at most one decoded output buffer and convert it to RGB.
|
||||
unsafe fn pull_output(&mut self) -> Result<Option<Frame>> {
|
||||
let mut info = AMediaCodecBufferInfo {
|
||||
offset: 0,
|
||||
size: 0,
|
||||
presentationTimeUs: 0,
|
||||
flags: 0,
|
||||
};
|
||||
let idx = AMediaCodec_dequeueOutputBuffer(self.codec, &mut info, DEQUEUE_TIMEOUT_US);
|
||||
if idx < 0 {
|
||||
// INFO_TRY_AGAIN_LATER / OUTPUT_FORMAT_CHANGED / BUFFERS_CHANGED: all
|
||||
// mean "no frame this time"; the next call picks the format up.
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut cap: usize = 0;
|
||||
let buf = AMediaCodec_getOutputBuffer(self.codec, idx as usize, &mut cap);
|
||||
let frame = if !buf.is_null() && info.size > 0 {
|
||||
let fmt = AMediaCodec_getOutputFormat(self.codec);
|
||||
let geom = read_geometry(fmt);
|
||||
if !fmt.is_null() {
|
||||
AMediaFormat_delete(fmt);
|
||||
}
|
||||
let end = (info.offset as usize)
|
||||
.saturating_add(info.size as usize)
|
||||
.min(cap);
|
||||
let start = (info.offset as usize).min(end);
|
||||
let data = core::slice::from_raw_parts(buf.add(start), end - start);
|
||||
yuv420_to_rgb(data, &geom)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let _ = AMediaCodec_releaseOutputBuffer(self.codec, idx as usize, false);
|
||||
Ok(frame)
|
||||
}
|
||||
|
||||
fn teardown(&mut self) {
|
||||
unsafe {
|
||||
if !self.codec.is_null() {
|
||||
let _ = AMediaCodec_stop(self.codec);
|
||||
AMediaCodec_delete(self.codec);
|
||||
self.codec = core::ptr::null_mut();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl VideoDecoder for MediaCodecDecoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
if au.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Cache parameter sets from keyframes; (re)prime when they change.
|
||||
let mut new_sps: Option<Vec<u8>> = None;
|
||||
let mut new_pps: Option<Vec<u8>> = None;
|
||||
for_each_nal(au, |nal| match nal.first().map(|b| b & 0x1F) {
|
||||
Some(7) => new_sps = Some(nal.to_vec()),
|
||||
Some(8) => new_pps = Some(nal.to_vec()),
|
||||
_ => {}
|
||||
});
|
||||
if let (Some(sps), Some(pps)) = (new_sps, new_pps) {
|
||||
if sps != self.sps || pps != self.pps || self.codec.is_null() {
|
||||
self.sps = sps;
|
||||
self.pps = pps;
|
||||
self.prime()?;
|
||||
}
|
||||
}
|
||||
if self.codec.is_null() {
|
||||
return Ok(None); // still waiting for the first keyframe
|
||||
}
|
||||
|
||||
unsafe {
|
||||
// Feed the whole Annex-B AU (start codes included — MediaCodec wants
|
||||
// the elementary stream byte format in byte-buffer mode).
|
||||
self.queue_input(au);
|
||||
self.pull_output()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MediaCodecDecoder {
|
||||
fn drop(&mut self) {
|
||||
self.teardown();
|
||||
}
|
||||
}
|
||||
|
||||
// SAFETY: the decoder is created on, and used only from, the single background
|
||||
// decode thread that owns it. `AMediaCodec` handles are not shared between
|
||||
// threads here (no `Sync`); moving ownership to that thread and calling it there
|
||||
// is sound.
|
||||
unsafe impl Send for MediaCodecDecoder {}
|
||||
|
||||
/// ~29.97 fps in microseconds; only needs to be monotonic.
|
||||
const PTS_STEP_US: i64 = 33_367;
|
||||
/// 10 ms dequeue timeout: long enough to pace one frame, short enough to keep
|
||||
/// the decode loop responsive.
|
||||
const DEQUEUE_TIMEOUT_US: i64 = 10_000;
|
||||
|
||||
// NDK format keys and MIME, as NUL-terminated byte strings (the `AMEDIAFORMAT_*`
|
||||
// macros expand to exactly these literals).
|
||||
const MIME: &[u8] = b"video/avc\0";
|
||||
const KEY_MIME: &[u8] = b"mime\0";
|
||||
const KEY_WIDTH: &[u8] = b"width\0";
|
||||
const KEY_HEIGHT: &[u8] = b"height\0";
|
||||
const KEY_STRIDE: &[u8] = b"stride\0";
|
||||
const KEY_SLICE_HEIGHT: &[u8] = b"slice-height\0";
|
||||
const KEY_COLOR_FORMAT: &[u8] = b"color-format\0";
|
||||
const KEY_CROP_LEFT: &[u8] = b"crop-left\0";
|
||||
const KEY_CROP_TOP: &[u8] = b"crop-top\0";
|
||||
const KEY_CROP_RIGHT: &[u8] = b"crop-right\0";
|
||||
const KEY_CROP_BOTTOM: &[u8] = b"crop-bottom\0";
|
||||
const KEY_CSD0: &[u8] = b"csd-0\0";
|
||||
const KEY_CSD1: &[u8] = b"csd-1\0";
|
||||
|
||||
// Android `MediaCodecInfo.CodecCapabilities` colour formats we handle.
|
||||
const COLOR_FORMAT_YUV420_PLANAR: i32 = 19; // I420 (Y, U, V planes)
|
||||
const COLOR_FORMAT_YUV420_SEMIPLANAR: i32 = 21; // NV12 (Y, interleaved UV)
|
||||
|
||||
/// Decoded output-buffer geometry pulled from the codec's output `AMediaFormat`.
|
||||
struct Geometry {
|
||||
color_format: i32,
|
||||
stride: usize,
|
||||
slice_height: usize,
|
||||
/// Visible picture rectangle (display size + top-left offset).
|
||||
crop_x: usize,
|
||||
crop_y: usize,
|
||||
disp_w: usize,
|
||||
disp_h: usize,
|
||||
}
|
||||
|
||||
/// Read the stride/slice-height/colour-format and crop rectangle from the output
|
||||
/// format, filling in sane fallbacks when a key is absent.
|
||||
unsafe fn read_geometry(fmt: *mut AMediaFormat) -> Geometry {
|
||||
let mut width = 0i32;
|
||||
let mut height = 0i32;
|
||||
let mut stride = 0i32;
|
||||
let mut slice_h = 0i32;
|
||||
let mut color = COLOR_FORMAT_YUV420_SEMIPLANAR;
|
||||
let (mut cl, mut ct, mut cr, mut cb) = (0i32, 0i32, -1i32, -1i32);
|
||||
if !fmt.is_null() {
|
||||
get_i32(fmt, KEY_WIDTH, &mut width);
|
||||
get_i32(fmt, KEY_HEIGHT, &mut height);
|
||||
get_i32(fmt, KEY_STRIDE, &mut stride);
|
||||
get_i32(fmt, KEY_SLICE_HEIGHT, &mut slice_h);
|
||||
get_i32(fmt, KEY_COLOR_FORMAT, &mut color);
|
||||
get_i32(fmt, KEY_CROP_LEFT, &mut cl);
|
||||
get_i32(fmt, KEY_CROP_TOP, &mut ct);
|
||||
get_i32(fmt, KEY_CROP_RIGHT, &mut cr);
|
||||
get_i32(fmt, KEY_CROP_BOTTOM, &mut cb);
|
||||
}
|
||||
if stride <= 0 {
|
||||
stride = width;
|
||||
}
|
||||
if slice_h <= 0 {
|
||||
slice_h = height;
|
||||
}
|
||||
// Crop rect is inclusive; fall back to the full coded size when absent.
|
||||
let (disp_w, disp_h, crop_x, crop_y) = if cr >= cl && cb >= ct && cr >= 0 && cb >= 0 {
|
||||
(
|
||||
(cr - cl + 1) as usize,
|
||||
(cb - ct + 1) as usize,
|
||||
cl.max(0) as usize,
|
||||
ct.max(0) as usize,
|
||||
)
|
||||
} else {
|
||||
(width.max(0) as usize, height.max(0) as usize, 0, 0)
|
||||
};
|
||||
Geometry {
|
||||
color_format: color,
|
||||
stride: stride.max(0) as usize,
|
||||
slice_height: slice_h.max(0) as usize,
|
||||
crop_x,
|
||||
crop_y,
|
||||
disp_w,
|
||||
disp_h,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert an I420 / NV12 output buffer to packed RGB8, cropping to the visible
|
||||
/// rectangle described by `g` (BT.601 limited range, matching the other paths).
|
||||
fn yuv420_to_rgb(data: &[u8], g: &Geometry) -> Option<Frame> {
|
||||
let (w, h, stride, plane_h) = (g.disp_w, g.disp_h, g.stride, g.slice_height);
|
||||
if w == 0 || h == 0 || stride == 0 || plane_h == 0 {
|
||||
return None;
|
||||
}
|
||||
let y_plane = stride * plane_h;
|
||||
let semiplanar = g.color_format != COLOR_FORMAT_YUV420_PLANAR;
|
||||
let mut rgb = vec![0u8; w * h * 3];
|
||||
for row in 0..h {
|
||||
let sy = g.crop_y + row;
|
||||
let y_off = sy * stride;
|
||||
let out_off = row * w * 3;
|
||||
let cy = sy / 2;
|
||||
for col in 0..w {
|
||||
let sx = g.crop_x + col;
|
||||
let y = *data.get(y_off + sx).unwrap_or(&16) as i32;
|
||||
let (u, v) = if semiplanar {
|
||||
// NV12: interleaved UV plane after Y, half-height, full stride.
|
||||
// The chroma pair for luma column `sx` sits at the even offset
|
||||
// `sx & !1` within the UV row (same stride as the Y plane).
|
||||
let uv_off = y_plane + cy * stride;
|
||||
let base = uv_off + (sx & !1);
|
||||
(
|
||||
*data.get(base).unwrap_or(&128) as i32,
|
||||
*data.get(base + 1).unwrap_or(&128) as i32,
|
||||
)
|
||||
} else {
|
||||
// I420: U plane then V plane, each half-width / half-height.
|
||||
let c_stride = stride / 2;
|
||||
let c_plane = (stride / 2) * (plane_h / 2);
|
||||
let cx = sx / 2;
|
||||
let u_off = y_plane + cy * c_stride + cx;
|
||||
let v_off = y_plane + c_plane + cy * c_stride + cx;
|
||||
(
|
||||
*data.get(u_off).unwrap_or(&128) as i32,
|
||||
*data.get(v_off).unwrap_or(&128) as i32,
|
||||
)
|
||||
};
|
||||
let c = y - 16;
|
||||
let d = u - 128;
|
||||
let e = v - 128;
|
||||
let r = ((298 * c + 409 * e + 128) >> 8).clamp(0, 255) as u8;
|
||||
let gg = ((298 * c - 100 * d - 208 * e + 128) >> 8).clamp(0, 255) as u8;
|
||||
let b = ((298 * c + 516 * d + 128) >> 8).clamp(0, 255) as u8;
|
||||
let o = out_off + col * 3;
|
||||
rgb[o] = r;
|
||||
rgb[o + 1] = gg;
|
||||
rgb[o + 2] = b;
|
||||
}
|
||||
}
|
||||
Some(Frame { w, h, rgb })
|
||||
}
|
||||
|
||||
/// Prepend a 4-byte Annex-B start code to a raw NAL for use as csd.
|
||||
fn with_start_code(nal: &[u8]) -> Vec<u8> {
|
||||
let mut v = Vec::with_capacity(nal.len() + 4);
|
||||
v.extend_from_slice(&[0, 0, 0, 1]);
|
||||
v.extend_from_slice(nal);
|
||||
v
|
||||
}
|
||||
|
||||
/// `AMediaFormat_getInt32` wrapper that leaves `out` untouched when the key is
|
||||
/// absent (so the caller's fallback stands).
|
||||
unsafe fn get_i32(fmt: *mut AMediaFormat, key: &[u8], out: *mut i32) -> bool {
|
||||
AMediaFormat_getInt32(fmt, key.as_ptr() as *const c_char, out)
|
||||
}
|
||||
|
||||
/// Iterate the Annex-B NAL units in `au`, start codes excluded and trailing zero
|
||||
/// bytes trimmed.
|
||||
fn for_each_nal<F: FnMut(&[u8])>(au: &[u8], mut f: F) {
|
||||
let n = au.len();
|
||||
let mut starts: Vec<usize> = Vec::new();
|
||||
let mut i = 0usize;
|
||||
while i + 2 < n {
|
||||
if au[i] == 0 && au[i + 1] == 0 && au[i + 2] == 1 {
|
||||
starts.push(i + 3);
|
||||
i += 3;
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
for k in 0..starts.len() {
|
||||
let begin = starts[k];
|
||||
let end = if k + 1 < starts.len() {
|
||||
starts[k + 1] - 3
|
||||
} else {
|
||||
n
|
||||
};
|
||||
if begin >= end {
|
||||
continue;
|
||||
}
|
||||
let mut nal = &au[begin..end];
|
||||
while let [rest @ .., 0] = nal {
|
||||
nal = rest;
|
||||
}
|
||||
if !nal.is_empty() {
|
||||
f(nal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn other(msg: &str) -> NdError {
|
||||
NdError::Other(anyhow::anyhow!(msg.to_string()))
|
||||
}
|
||||
|
||||
/// Self-contained FFI to the NDK Media C ABI (`libmediandk`). These signatures
|
||||
/// match `<media/NdkMediaCodec.h>` / `<media/NdkMediaFormat.h>` and are stable
|
||||
/// since Android API 21 (`AMediaFormat_getRect`/crop keys since API 21 as well
|
||||
/// for the common decoders). The `ndk`/`ndk-sys` crates wrap the same library.
|
||||
#[allow(non_snake_case, non_camel_case_types)]
|
||||
mod sys {
|
||||
use core::ffi::{c_char, c_void};
|
||||
|
||||
pub type AMediaCodec = c_void;
|
||||
pub type AMediaFormat = c_void;
|
||||
pub type media_status_t = i32;
|
||||
pub const AMEDIA_OK: media_status_t = 0;
|
||||
|
||||
#[repr(C)]
|
||||
pub struct AMediaCodecBufferInfo {
|
||||
pub offset: i32,
|
||||
pub size: i32,
|
||||
pub presentationTimeUs: i64,
|
||||
pub flags: u32,
|
||||
}
|
||||
|
||||
#[link(name = "mediandk")]
|
||||
extern "C" {
|
||||
pub fn AMediaCodec_createDecoderByType(mime_type: *const c_char) -> *mut AMediaCodec;
|
||||
pub fn AMediaCodec_configure(
|
||||
codec: *mut AMediaCodec,
|
||||
format: *mut AMediaFormat,
|
||||
surface: *mut c_void,
|
||||
crypto: *mut c_void,
|
||||
flags: u32,
|
||||
) -> media_status_t;
|
||||
pub fn AMediaCodec_start(codec: *mut AMediaCodec) -> media_status_t;
|
||||
pub fn AMediaCodec_stop(codec: *mut AMediaCodec) -> media_status_t;
|
||||
pub fn AMediaCodec_delete(codec: *mut AMediaCodec) -> media_status_t;
|
||||
pub fn AMediaCodec_dequeueInputBuffer(codec: *mut AMediaCodec, timeout_us: i64) -> isize;
|
||||
pub fn AMediaCodec_getInputBuffer(
|
||||
codec: *mut AMediaCodec,
|
||||
idx: usize,
|
||||
out_size: *mut usize,
|
||||
) -> *mut u8;
|
||||
pub fn AMediaCodec_queueInputBuffer(
|
||||
codec: *mut AMediaCodec,
|
||||
idx: usize,
|
||||
offset: i64,
|
||||
size: usize,
|
||||
time_us: i64,
|
||||
flags: u32,
|
||||
) -> media_status_t;
|
||||
pub fn AMediaCodec_dequeueOutputBuffer(
|
||||
codec: *mut AMediaCodec,
|
||||
info: *mut AMediaCodecBufferInfo,
|
||||
timeout_us: i64,
|
||||
) -> isize;
|
||||
pub fn AMediaCodec_getOutputBuffer(
|
||||
codec: *mut AMediaCodec,
|
||||
idx: usize,
|
||||
out_size: *mut usize,
|
||||
) -> *mut u8;
|
||||
pub fn AMediaCodec_getOutputFormat(codec: *mut AMediaCodec) -> *mut AMediaFormat;
|
||||
pub fn AMediaCodec_releaseOutputBuffer(
|
||||
codec: *mut AMediaCodec,
|
||||
idx: usize,
|
||||
render: bool,
|
||||
) -> media_status_t;
|
||||
|
||||
pub fn AMediaFormat_new() -> *mut AMediaFormat;
|
||||
pub fn AMediaFormat_delete(fmt: *mut AMediaFormat) -> media_status_t;
|
||||
pub fn AMediaFormat_setString(
|
||||
fmt: *mut AMediaFormat,
|
||||
name: *const c_char,
|
||||
value: *const c_char,
|
||||
);
|
||||
pub fn AMediaFormat_setInt32(fmt: *mut AMediaFormat, name: *const c_char, value: i32);
|
||||
pub fn AMediaFormat_setBuffer(
|
||||
fmt: *mut AMediaFormat,
|
||||
name: *const c_char,
|
||||
data: *const c_void,
|
||||
size: usize,
|
||||
);
|
||||
pub fn AMediaFormat_getInt32(
|
||||
fmt: *mut AMediaFormat,
|
||||
name: *const c_char,
|
||||
out: *mut i32,
|
||||
) -> bool;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,646 @@
|
||||
//! Windows hardware H.264 decode via Media Foundation.
|
||||
//!
|
||||
//! This feeds the host's Annex-B access units into the Microsoft H.264 decoder
|
||||
//! MFT (`CLSID_MSH264DecoderMFT`). When a Direct3D 11 device can be attached the
|
||||
//! decode runs on the GPU video engine (DXVA2) and the decoded NV12 surface is
|
||||
//! copied back to system memory; otherwise the same MFT runs as Microsoft's
|
||||
//! (still heavily optimized) software decoder. Either way the output is the same
|
||||
//! packed-RGB8 [`Frame`] the openh264 path produces, so the viewer is unchanged.
|
||||
//!
|
||||
//! The whole module is only compiled on Windows (see the `#[cfg(windows)]` gate
|
||||
//! in `decode.rs`). If anything here fails at runtime, [`super::best_decoder`]
|
||||
//! falls back to the portable openh264 software decoder.
|
||||
|
||||
use super::{Frame, VideoDecoder};
|
||||
use crate::error::{NdError, Result};
|
||||
|
||||
use windows::core::Interface;
|
||||
use windows::Win32::Foundation::{HMODULE, TRUE};
|
||||
use windows::Win32::Graphics::Direct3D::D3D_DRIVER_TYPE_HARDWARE;
|
||||
use windows::Win32::Graphics::Direct3D11::{
|
||||
D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Multithread, ID3D11Texture2D,
|
||||
D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_BGRA_SUPPORT, D3D11_CREATE_DEVICE_VIDEO_SUPPORT,
|
||||
D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC,
|
||||
D3D11_USAGE_STAGING,
|
||||
};
|
||||
use windows::Win32::Graphics::Dxgi::Common::{DXGI_FORMAT, DXGI_SAMPLE_DESC};
|
||||
use windows::Win32::Media::MediaFoundation::*;
|
||||
use windows::Win32::System::Com::{
|
||||
CoCreateInstance, CoInitializeEx, CLSCTX_INPROC_SERVER, COINIT_MULTITHREADED,
|
||||
};
|
||||
|
||||
/// ~29.97 fps in 100-ns units. The host stream is low-latency with no B-frames,
|
||||
/// so exact timing is unimportant — the decoder only needs monotonic stamps.
|
||||
const HNS_PER_FRAME: i64 = 333_667;
|
||||
|
||||
/// Hardware (or MS-software) Media Foundation H.264 decoder.
|
||||
pub struct MediaFoundationDecoder {
|
||||
transform: IMFTransform,
|
||||
// Direct3D plumbing kept alive for the lifetime of the decoder (hw path only).
|
||||
device: Option<ID3D11Device>,
|
||||
context: Option<ID3D11DeviceContext>,
|
||||
_manager: Option<IMFDXGIDeviceManager>,
|
||||
staging: Option<ID3D11Texture2D>,
|
||||
staging_dims: (u32, u32),
|
||||
/// True when DXVA hardware decode is active (output is a D3D texture).
|
||||
hw: bool,
|
||||
/// True when the MFT allocates its own output samples.
|
||||
provides_samples: bool,
|
||||
output_size: usize,
|
||||
stride: usize,
|
||||
/// Coded (padded, usually 16-aligned) frame size from MF_MT_FRAME_SIZE.
|
||||
/// This is the NV12 *plane* geometry we read out of memory, not the picture.
|
||||
coded_w: usize,
|
||||
coded_h: usize,
|
||||
/// True visible picture size (from the display aperture). This is what the
|
||||
/// emitted `Frame.w`/`Frame.h` report, e.g. 1920x1080 even when the decoder
|
||||
/// pads the coded frame to 1920x1088.
|
||||
width: usize,
|
||||
height: usize,
|
||||
/// Top-left offset of the visible region inside the coded frame (usually 0).
|
||||
crop_x: usize,
|
||||
crop_y: usize,
|
||||
started: bool,
|
||||
pts: i64,
|
||||
/// A frame pulled out while draining the MFT during a busy ProcessInput.
|
||||
pending: Option<Frame>,
|
||||
}
|
||||
|
||||
impl MediaFoundationDecoder {
|
||||
/// Initialize Media Foundation and the H.264 decoder MFT. Returns an error
|
||||
/// (so the caller can fall back to openh264) if MF is unavailable.
|
||||
pub fn new() -> Result<Self> {
|
||||
unsafe {
|
||||
// COM (MTA) + MF startup. CoInitializeEx returns an HRESULT that is
|
||||
// fine to ignore (S_FALSE / RPC_E_CHANGED_MODE just mean the thread
|
||||
// already has an apartment). MFStartup is refcounted against the
|
||||
// MFShutdown in Drop.
|
||||
let _ = CoInitializeEx(None, COINIT_MULTITHREADED);
|
||||
MFStartup(MF_VERSION, MFSTARTUP_FULL).map_err(mferr("MFStartup"))?;
|
||||
|
||||
let transform: IMFTransform =
|
||||
CoCreateInstance(&CLSID_MSH264DecoderMFT, None, CLSCTX_INPROC_SERVER)
|
||||
.map_err(mferr("CoCreateInstance(MSH264DecoderMFT)"))?;
|
||||
|
||||
// Low-latency mode: emit each frame as soon as it is decoded.
|
||||
if let Ok(attrs) = transform.GetAttributes() {
|
||||
let _ = attrs.SetUINT32(&MF_LOW_LATENCY, 1);
|
||||
}
|
||||
|
||||
// Best-effort DXVA hardware acceleration.
|
||||
let (device, context, manager, hw) = match init_d3d(&transform) {
|
||||
Ok((d, c, m)) => (Some(d), Some(c), Some(m), true),
|
||||
Err(e) => {
|
||||
tracing::debug!("decode(mf): DXVA unavailable ({e}); MS software MFT");
|
||||
(None, None, None, false)
|
||||
}
|
||||
};
|
||||
|
||||
// Input media type: H.264 Annex-B elementary stream.
|
||||
let input = MFCreateMediaType().map_err(mferr("MFCreateMediaType(input)"))?;
|
||||
input
|
||||
.SetGUID(&MF_MT_MAJOR_TYPE, &MFMediaType_Video)
|
||||
.map_err(mferr("set input major type"))?;
|
||||
input
|
||||
.SetGUID(&MF_MT_SUBTYPE, &MFVideoFormat_H264)
|
||||
.map_err(mferr("set input subtype H264"))?;
|
||||
transform
|
||||
.SetInputType(0, &input, 0)
|
||||
.map_err(mferr("SetInputType"))?;
|
||||
|
||||
let mut me = Self {
|
||||
transform,
|
||||
device,
|
||||
context,
|
||||
_manager: manager,
|
||||
staging: None,
|
||||
staging_dims: (0, 0),
|
||||
hw,
|
||||
provides_samples: false,
|
||||
output_size: 0,
|
||||
stride: 0,
|
||||
coded_w: 0,
|
||||
coded_h: 0,
|
||||
width: 0,
|
||||
height: 0,
|
||||
crop_x: 0,
|
||||
crop_y: 0,
|
||||
started: false,
|
||||
pts: 0,
|
||||
pending: None,
|
||||
};
|
||||
|
||||
me.configure_output()?;
|
||||
|
||||
me.transform
|
||||
.ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, 0)
|
||||
.map_err(mferr("NOTIFY_BEGIN_STREAMING"))?;
|
||||
me.transform
|
||||
.ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, 0)
|
||||
.map_err(mferr("NOTIFY_START_OF_STREAM"))?;
|
||||
me.started = true;
|
||||
|
||||
Ok(me)
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether DXVA GPU decode is active (vs. the MS software MFT).
|
||||
pub fn is_hardware(&self) -> bool {
|
||||
self.hw
|
||||
}
|
||||
|
||||
/// Select an NV12 output type and read back its geometry + allocation needs.
|
||||
unsafe fn configure_output(&mut self) -> Result<()> {
|
||||
let mut i = 0u32;
|
||||
loop {
|
||||
match self.transform.GetOutputAvailableType(0, i) {
|
||||
Ok(t) => {
|
||||
let sub = t.GetGUID(&MF_MT_SUBTYPE).unwrap_or_default();
|
||||
if sub == MFVideoFormat_NV12 {
|
||||
self.transform
|
||||
.SetOutputType(0, &t, 0)
|
||||
.map_err(mferr("SetOutputType(NV12)"))?;
|
||||
self.read_output_geometry(&t);
|
||||
break;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
Err(e) if e.code() == MF_E_NO_MORE_TYPES => {
|
||||
// Nothing offered yet (common before the first keyframe sets
|
||||
// the frame size) — ask for NV12 explicitly.
|
||||
let t = MFCreateMediaType().map_err(mferr("MFCreateMediaType(output)"))?;
|
||||
let _ = t.SetGUID(&MF_MT_MAJOR_TYPE, &MFMediaType_Video);
|
||||
let _ = t.SetGUID(&MF_MT_SUBTYPE, &MFVideoFormat_NV12);
|
||||
self.transform
|
||||
.SetOutputType(0, &t, 0)
|
||||
.map_err(mferr("SetOutputType(NV12 manual)"))?;
|
||||
if let Ok(cur) = self.transform.GetOutputCurrentType(0) {
|
||||
self.read_output_geometry(&cur);
|
||||
}
|
||||
break;
|
||||
}
|
||||
Err(e) => return Err(mferr("GetOutputAvailableType")(e)),
|
||||
}
|
||||
}
|
||||
|
||||
let info = self
|
||||
.transform
|
||||
.GetOutputStreamInfo(0)
|
||||
.map_err(mferr("GetOutputStreamInfo"))?;
|
||||
let provides =
|
||||
(MFT_OUTPUT_STREAM_PROVIDES_SAMPLES.0 | MFT_OUTPUT_STREAM_CAN_PROVIDE_SAMPLES.0) as u32;
|
||||
self.provides_samples = (info.dwFlags & provides) != 0;
|
||||
self.output_size = info.cbSize as usize;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Pull the coded frame size, the true display rectangle, and the stride out
|
||||
/// of a configured output media type.
|
||||
///
|
||||
/// The MS H.264 decoder reports the *coded* size in `MF_MT_FRAME_SIZE` — that
|
||||
/// is padded up to the codec's alignment (e.g. 1920x1088 for 1080p content,
|
||||
/// and larger for other resolutions), which is what produced the wrong
|
||||
/// 1920x18xx figure in the stats. The real picture rectangle is carried in a
|
||||
/// video-aperture blob, so we read that separately and keep the coded size
|
||||
/// only for the in-memory plane layout.
|
||||
unsafe fn read_output_geometry(&mut self, t: &IMFMediaType) {
|
||||
// Coded (padded) plane geometry — how the NV12 bytes are laid out.
|
||||
if let Ok(fs) = t.GetUINT64(&MF_MT_FRAME_SIZE) {
|
||||
self.coded_w = (fs >> 32) as usize;
|
||||
self.coded_h = (fs & 0xffff_ffff) as usize;
|
||||
}
|
||||
|
||||
// True visible picture region. Prefer the minimum display aperture
|
||||
// (what a correct renderer must show), then the geometric aperture, then
|
||||
// the pan/scan aperture if pan/scan is actually enabled. When none is
|
||||
// present the whole coded frame is the picture.
|
||||
let aperture = read_aperture(t, &MF_MT_MINIMUM_DISPLAY_APERTURE)
|
||||
.or_else(|| read_aperture(t, &MF_MT_GEOMETRIC_APERTURE))
|
||||
.or_else(|| {
|
||||
if t.GetUINT32(&MF_MT_PAN_SCAN_ENABLED).unwrap_or(0) != 0 {
|
||||
read_aperture(t, &MF_MT_PAN_SCAN_APERTURE)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
match aperture {
|
||||
Some((ox, oy, aw, ah)) => {
|
||||
self.crop_x = ox;
|
||||
self.crop_y = oy;
|
||||
// Clamp the aperture inside the coded frame so a malformed blob
|
||||
// can never make the cropper read past the decoded plane.
|
||||
let max_w = self.coded_w.saturating_sub(ox);
|
||||
let max_h = self.coded_h.saturating_sub(oy);
|
||||
self.width = if max_w > 0 { aw.min(max_w) } else { aw };
|
||||
self.height = if max_h > 0 { ah.min(max_h) } else { ah };
|
||||
}
|
||||
None => {
|
||||
self.crop_x = 0;
|
||||
self.crop_y = 0;
|
||||
self.width = self.coded_w;
|
||||
self.height = self.coded_h;
|
||||
}
|
||||
}
|
||||
|
||||
self.stride = match t.GetUINT32(&MF_MT_DEFAULT_STRIDE) {
|
||||
Ok(s) => (s as i32).unsigned_abs() as usize,
|
||||
Err(_) => self.coded_w,
|
||||
};
|
||||
if self.stride == 0 {
|
||||
self.stride = self.coded_w;
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap one Annex-B access unit in an input IMFSample.
|
||||
unsafe fn make_input_sample(&mut self, au: &[u8]) -> Result<IMFSample> {
|
||||
let sample = MFCreateSample().map_err(mferr("MFCreateSample"))?;
|
||||
let buffer =
|
||||
MFCreateMemoryBuffer(au.len() as u32).map_err(mferr("MFCreateMemoryBuffer"))?;
|
||||
let mut ptr: *mut u8 = std::ptr::null_mut();
|
||||
buffer
|
||||
.Lock(&mut ptr, None, None)
|
||||
.map_err(mferr("Lock(input)"))?;
|
||||
std::ptr::copy_nonoverlapping(au.as_ptr(), ptr, au.len());
|
||||
buffer.Unlock().map_err(mferr("Unlock(input)"))?;
|
||||
buffer
|
||||
.SetCurrentLength(au.len() as u32)
|
||||
.map_err(mferr("SetCurrentLength"))?;
|
||||
sample.AddBuffer(&buffer).map_err(mferr("AddBuffer"))?;
|
||||
let _ = sample.SetSampleTime(self.pts);
|
||||
let _ = sample.SetSampleDuration(HNS_PER_FRAME);
|
||||
self.pts += HNS_PER_FRAME;
|
||||
Ok(sample)
|
||||
}
|
||||
|
||||
/// Run ProcessOutput once, materializing a frame if one is ready.
|
||||
unsafe fn pull_output(&mut self) -> Result<Option<Frame>> {
|
||||
loop {
|
||||
let sample_slot: Option<IMFSample> = if self.provides_samples {
|
||||
None
|
||||
} else {
|
||||
let s = MFCreateSample().map_err(mferr("MFCreateSample(output)"))?;
|
||||
let size = self
|
||||
.output_size
|
||||
.max(self.stride * self.coded_h * 3 / 2)
|
||||
.max(1);
|
||||
let b = MFCreateMemoryBuffer(size as u32)
|
||||
.map_err(mferr("MFCreateMemoryBuffer(output)"))?;
|
||||
s.AddBuffer(&b).map_err(mferr("AddBuffer(output)"))?;
|
||||
Some(s)
|
||||
};
|
||||
|
||||
let mut buffers = [MFT_OUTPUT_DATA_BUFFER {
|
||||
dwStreamID: 0,
|
||||
pSample: std::mem::ManuallyDrop::new(sample_slot),
|
||||
dwStatus: 0,
|
||||
pEvents: std::mem::ManuallyDrop::new(None),
|
||||
}];
|
||||
let mut status = 0u32;
|
||||
let res = self.transform.ProcessOutput(0, &mut buffers, &mut status);
|
||||
|
||||
// Reclaim ownership of whatever is in the slots so they drop cleanly.
|
||||
let out_sample = std::mem::ManuallyDrop::take(&mut buffers[0].pSample);
|
||||
let _events = std::mem::ManuallyDrop::take(&mut buffers[0].pEvents);
|
||||
|
||||
match res {
|
||||
Ok(()) => {
|
||||
let s = out_sample.ok_or_else(|| other("ProcessOutput produced no sample"))?;
|
||||
return Ok(Some(self.sample_to_frame(&s)?));
|
||||
}
|
||||
Err(e) if e.code() == MF_E_TRANSFORM_NEED_MORE_INPUT => return Ok(None),
|
||||
Err(e) if e.code() == MF_E_TRANSFORM_STREAM_CHANGE => {
|
||||
// The decoder learned the real resolution — renegotiate NV12.
|
||||
self.configure_output()?;
|
||||
continue;
|
||||
}
|
||||
Err(e) => return Err(mferr("ProcessOutput")(e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a decoded output sample (D3D texture or system memory) to RGB.
|
||||
unsafe fn sample_to_frame(&mut self, sample: &IMFSample) -> Result<Frame> {
|
||||
let buffer = sample
|
||||
.GetBufferByIndex(0)
|
||||
.map_err(mferr("GetBufferByIndex"))?;
|
||||
|
||||
if self.hw {
|
||||
if let Ok(dxgi) = buffer.cast::<IMFDXGIBuffer>() {
|
||||
return self.dxgi_to_frame(&dxgi);
|
||||
}
|
||||
}
|
||||
|
||||
// System-memory NV12.
|
||||
let mut ptr: *mut u8 = std::ptr::null_mut();
|
||||
let mut cur: u32 = 0;
|
||||
buffer
|
||||
.Lock(&mut ptr, None, Some(&mut cur))
|
||||
.map_err(mferr("Lock(output)"))?;
|
||||
let stride = if self.stride == 0 {
|
||||
self.coded_w
|
||||
} else {
|
||||
self.stride
|
||||
};
|
||||
// Derive the (possibly 16-aligned) plane height from the real byte count,
|
||||
// never below the coded height the decoder reported.
|
||||
let plane_h = if stride > 0 {
|
||||
((cur as usize) * 2) / (stride * 3)
|
||||
} else {
|
||||
self.coded_h
|
||||
}
|
||||
.max(self.coded_h);
|
||||
let data = std::slice::from_raw_parts(ptr, cur as usize);
|
||||
// Crop the visible picture (width/height at crop_x/crop_y) out of the
|
||||
// padded NV12 plane so the Frame is the true display size.
|
||||
let frame = nv12_to_rgb(
|
||||
data,
|
||||
stride,
|
||||
plane_h,
|
||||
self.crop_x,
|
||||
self.crop_y,
|
||||
self.width,
|
||||
self.height,
|
||||
);
|
||||
let _ = buffer.Unlock();
|
||||
frame.ok_or_else(|| other("NV12→RGB conversion failed"))
|
||||
}
|
||||
|
||||
/// Copy a decoded GPU NV12 texture to a staging texture, map it, convert.
|
||||
unsafe fn dxgi_to_frame(&mut self, dxgi: &IMFDXGIBuffer) -> Result<Frame> {
|
||||
let device = self.device.clone().ok_or_else(|| other("no D3D device"))?;
|
||||
let context = self
|
||||
.context
|
||||
.clone()
|
||||
.ok_or_else(|| other("no D3D context"))?;
|
||||
|
||||
let mut tex_ptr: *mut core::ffi::c_void = std::ptr::null_mut();
|
||||
dxgi.GetResource(&ID3D11Texture2D::IID, &mut tex_ptr)
|
||||
.map_err(mferr("IMFDXGIBuffer::GetResource"))?;
|
||||
if tex_ptr.is_null() {
|
||||
return Err(other("decoded texture was null"));
|
||||
}
|
||||
let src = ID3D11Texture2D::from_raw(tex_ptr);
|
||||
let sub = dxgi.GetSubresourceIndex().unwrap_or(0);
|
||||
|
||||
let mut desc: D3D11_TEXTURE2D_DESC = std::mem::zeroed();
|
||||
src.GetDesc(&mut desc);
|
||||
|
||||
self.ensure_staging(&device, desc.Width, desc.Height, desc.Format)?;
|
||||
let staging = self
|
||||
.staging
|
||||
.as_ref()
|
||||
.ok_or_else(|| other("no staging texture"))?
|
||||
.clone();
|
||||
|
||||
// Copy the one decoded subresource out of the decoder's texture array.
|
||||
context.CopySubresourceRegion(&staging, 0, 0, 0, 0, &src, sub, None);
|
||||
|
||||
let mut mapped: D3D11_MAPPED_SUBRESOURCE = std::mem::zeroed();
|
||||
context
|
||||
.Map(&staging, 0, D3D11_MAP_READ, 0, Some(&mut mapped))
|
||||
.map_err(mferr("Map(staging)"))?;
|
||||
let row_pitch = mapped.RowPitch as usize;
|
||||
// The staging texture height is the coded (padded) plane height.
|
||||
let plane_h = desc.Height as usize;
|
||||
let total = row_pitch * plane_h * 3 / 2;
|
||||
let data = std::slice::from_raw_parts(mapped.pData as *const u8, total);
|
||||
// Crop to the true visible picture out of the padded GPU NV12 surface.
|
||||
let frame = nv12_to_rgb(
|
||||
data,
|
||||
row_pitch,
|
||||
plane_h,
|
||||
self.crop_x,
|
||||
self.crop_y,
|
||||
self.width,
|
||||
self.height,
|
||||
);
|
||||
context.Unmap(&staging, 0);
|
||||
|
||||
frame.ok_or_else(|| other("NV12→RGB conversion failed"))
|
||||
}
|
||||
|
||||
/// (Re)create the CPU-readable staging texture when the size changes.
|
||||
unsafe fn ensure_staging(
|
||||
&mut self,
|
||||
device: &ID3D11Device,
|
||||
w: u32,
|
||||
h: u32,
|
||||
fmt: DXGI_FORMAT,
|
||||
) -> Result<()> {
|
||||
if self.staging.is_some() && self.staging_dims == (w, h) {
|
||||
return Ok(());
|
||||
}
|
||||
let desc = D3D11_TEXTURE2D_DESC {
|
||||
Width: w,
|
||||
Height: h,
|
||||
MipLevels: 1,
|
||||
ArraySize: 1,
|
||||
Format: fmt,
|
||||
SampleDesc: DXGI_SAMPLE_DESC {
|
||||
Count: 1,
|
||||
Quality: 0,
|
||||
},
|
||||
Usage: D3D11_USAGE_STAGING,
|
||||
BindFlags: 0,
|
||||
CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32,
|
||||
MiscFlags: 0,
|
||||
};
|
||||
let mut tex: Option<ID3D11Texture2D> = None;
|
||||
device
|
||||
.CreateTexture2D(&desc, None, Some(&mut tex))
|
||||
.map_err(mferr("CreateTexture2D(staging)"))?;
|
||||
self.staging = tex;
|
||||
self.staging_dims = (w, h);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl VideoDecoder for MediaFoundationDecoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
if au.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
unsafe {
|
||||
let sample = self.make_input_sample(au)?;
|
||||
|
||||
// Feed the AU. If the MFT is momentarily full, drain its output and
|
||||
// retry (keeping the most recent frame for this call).
|
||||
let mut guard = 0;
|
||||
loop {
|
||||
match self.transform.ProcessInput(0, &sample, 0) {
|
||||
Ok(()) => break,
|
||||
Err(e) if e.code() == MF_E_NOTACCEPTING => {
|
||||
if let Some(f) = self.pull_output()? {
|
||||
self.pending = Some(f);
|
||||
}
|
||||
guard += 1;
|
||||
if guard > 16 {
|
||||
return Err(other("MFT refused input repeatedly"));
|
||||
}
|
||||
}
|
||||
Err(e) => return Err(mferr("ProcessInput")(e)),
|
||||
}
|
||||
}
|
||||
|
||||
// Pull the decoded picture(s). With zero-latency/no-B-frame input the
|
||||
// common case is exactly one frame per AU.
|
||||
let mut out = self.pending.take();
|
||||
while let Some(f) = self.pull_output()? {
|
||||
out = Some(f);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MediaFoundationDecoder {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
if self.started {
|
||||
let _ = self
|
||||
.transform
|
||||
.ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, 0);
|
||||
let _ = self
|
||||
.transform
|
||||
.ProcessMessage(MFT_MESSAGE_NOTIFY_END_STREAMING, 0);
|
||||
}
|
||||
// Balances the MFStartup in `new` (both are refcounted).
|
||||
let _ = MFShutdown();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SAFETY: a MediaFoundationDecoder is created on, and used only from, the single
|
||||
// background network thread that owns it (the viewer's decode loop). COM is
|
||||
// initialized multithreaded (MTA) so the MFT is agile, and the D3D11 device is
|
||||
// marked multithread-protected, so moving the decoder to that thread and calling
|
||||
// it there is sound. It is never shared between threads (no `Sync`).
|
||||
unsafe impl Send for MediaFoundationDecoder {}
|
||||
|
||||
/// Create a D3D11 device + DXGI device manager and hand it to the decoder to
|
||||
/// enable DXVA2 hardware decode. Any failure leaves the MFT in software mode.
|
||||
unsafe fn init_d3d(
|
||||
transform: &IMFTransform,
|
||||
) -> Result<(ID3D11Device, ID3D11DeviceContext, IMFDXGIDeviceManager)> {
|
||||
let mut device: Option<ID3D11Device> = None;
|
||||
let mut context: Option<ID3D11DeviceContext> = None;
|
||||
D3D11CreateDevice(
|
||||
None,
|
||||
D3D_DRIVER_TYPE_HARDWARE,
|
||||
HMODULE(std::ptr::null_mut()),
|
||||
D3D11_CREATE_DEVICE_VIDEO_SUPPORT | D3D11_CREATE_DEVICE_BGRA_SUPPORT,
|
||||
None,
|
||||
D3D11_SDK_VERSION,
|
||||
Some(&mut device),
|
||||
None,
|
||||
Some(&mut context),
|
||||
)
|
||||
.map_err(mferr("D3D11CreateDevice"))?;
|
||||
let device = device.ok_or_else(|| other("D3D11CreateDevice gave no device"))?;
|
||||
let context = context.ok_or_else(|| other("D3D11CreateDevice gave no context"))?;
|
||||
|
||||
// MF calls the device from its own worker threads.
|
||||
if let Ok(mt) = device.cast::<ID3D11Multithread>() {
|
||||
let _ = mt.SetMultithreadProtected(TRUE);
|
||||
}
|
||||
|
||||
let mut token = 0u32;
|
||||
let mut manager: Option<IMFDXGIDeviceManager> = None;
|
||||
MFCreateDXGIDeviceManager(&mut token, &mut manager)
|
||||
.map_err(mferr("MFCreateDXGIDeviceManager"))?;
|
||||
let manager = manager.ok_or_else(|| other("no DXGI device manager"))?;
|
||||
manager
|
||||
.ResetDevice(&device, token)
|
||||
.map_err(mferr("IMFDXGIDeviceManager::ResetDevice"))?;
|
||||
|
||||
// The ulParam for SET_D3D_MANAGER is the IUnknown* of the device manager.
|
||||
let param = manager.as_raw() as usize;
|
||||
transform
|
||||
.ProcessMessage(MFT_MESSAGE_SET_D3D_MANAGER, param)
|
||||
.map_err(mferr("MFT_MESSAGE_SET_D3D_MANAGER"))?;
|
||||
|
||||
Ok((device, context, manager))
|
||||
}
|
||||
|
||||
/// NV12 → packed RGB8 (BT.601 limited range, matching the openh264 path),
|
||||
/// cropping the visible `w`x`h` picture at offset (`off_x`, `off_y`) out of the
|
||||
/// (possibly padded) coded plane described by `stride` x `plane_h`.
|
||||
///
|
||||
/// `plane_h` is the coded/aligned luma height: the interleaved UV plane starts
|
||||
/// at `stride * plane_h`, so it must be the padded height, not the visible one.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn nv12_to_rgb(
|
||||
data: &[u8],
|
||||
stride: usize,
|
||||
plane_h: usize,
|
||||
off_x: usize,
|
||||
off_y: usize,
|
||||
w: usize,
|
||||
h: usize,
|
||||
) -> Option<Frame> {
|
||||
if w == 0 || h == 0 || stride == 0 {
|
||||
return None;
|
||||
}
|
||||
let uv_base = stride * plane_h;
|
||||
let mut rgb = vec![0u8; w * h * 3];
|
||||
for row in 0..h {
|
||||
let sy = off_y + row; // source luma row inside the coded plane
|
||||
let y_off = sy * stride;
|
||||
let uv_off = uv_base + (sy / 2) * stride;
|
||||
let out_off = row * w * 3;
|
||||
for col in 0..w {
|
||||
let sx = off_x + col; // source luma column inside the coded plane
|
||||
let y = *data.get(y_off + sx).unwrap_or(&16) as i32;
|
||||
let uvx = sx & !1; // chroma pair aligned to the even luma column
|
||||
let u = *data.get(uv_off + uvx).unwrap_or(&128) as i32;
|
||||
let v = *data.get(uv_off + uvx + 1).unwrap_or(&128) as i32;
|
||||
let c = y - 16;
|
||||
let d = u - 128;
|
||||
let e = v - 128;
|
||||
let r = ((298 * c + 409 * e + 128) >> 8).clamp(0, 255) as u8;
|
||||
let g = ((298 * c - 100 * d - 208 * e + 128) >> 8).clamp(0, 255) as u8;
|
||||
let b = ((298 * c + 516 * d + 128) >> 8).clamp(0, 255) as u8;
|
||||
let o = out_off + col * 3;
|
||||
rgb[o] = r;
|
||||
rgb[o + 1] = g;
|
||||
rgb[o + 2] = b;
|
||||
}
|
||||
}
|
||||
Some(Frame { w, h, rgb })
|
||||
}
|
||||
|
||||
/// Read an `MFVideoArea` aperture blob (display / geometric / pan-scan) off a
|
||||
/// media type and return `(offset_x, offset_y, width, height)` in whole pixels,
|
||||
/// or `None` when the attribute is absent or degenerate.
|
||||
///
|
||||
/// `MFVideoArea` is a fixed 16-byte POD (`MFOffset` x,y + `SIZE`), so we read it
|
||||
/// straight into the struct. The `MFOffset` fractional part is ignored: H.264
|
||||
/// display apertures are always integer-pixel aligned.
|
||||
unsafe fn read_aperture(
|
||||
t: &IMFMediaType,
|
||||
key: *const windows::core::GUID,
|
||||
) -> Option<(usize, usize, usize, usize)> {
|
||||
let mut area = MFVideoArea::default();
|
||||
let buf = std::slice::from_raw_parts_mut(
|
||||
(&mut area as *mut MFVideoArea).cast::<u8>(),
|
||||
std::mem::size_of::<MFVideoArea>(),
|
||||
);
|
||||
t.GetBlob(key, buf, None).ok()?;
|
||||
let w = area.Area.cx;
|
||||
let h = area.Area.cy;
|
||||
if w <= 0 || h <= 0 {
|
||||
return None;
|
||||
}
|
||||
let ox = area.OffsetX.value.max(0) as usize;
|
||||
let oy = area.OffsetY.value.max(0) as usize;
|
||||
Some((ox, oy, w as usize, h as usize))
|
||||
}
|
||||
|
||||
/// Build a context-tagged error mapper for `windows::core::Error`.
|
||||
fn mferr(ctx: &'static str) -> impl Fn(windows::core::Error) -> NdError {
|
||||
move |e| NdError::Other(anyhow::anyhow!("MediaFoundation {ctx}: {e}"))
|
||||
}
|
||||
|
||||
fn other(msg: &str) -> NdError {
|
||||
NdError::Other(anyhow::anyhow!(msg.to_string()))
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
//! Linux hardware H.264 decode via VAAPI, driven through FFmpeg/libavcodec.
|
||||
//!
|
||||
//! Pure libva decoding would mean parsing the H.264 bitstream down to slice
|
||||
//! headers and hand-building `VAPictureParameterBufferH264` / reference lists;
|
||||
//! libavcodec already does all of that and exposes VAAPI as a hardware
|
||||
//! accelerator, so this backend opens the libavcodec H.264 decoder with a VAAPI
|
||||
//! `hw_device_ctx`. Decode runs on the GPU (Intel/AMD/`iHD`/`radeonsi` via
|
||||
//! `/dev/dri`), the decoded surface is transferred to system memory (NV12), and
|
||||
//! swscale converts it to the same packed-RGB8 [`Frame`] the openh264 path
|
||||
//! produces, so the viewer is unchanged.
|
||||
//!
|
||||
//! The module is only compiled on Linux (see the `#[cfg(target_os = "linux")]`
|
||||
//! gate in `decode.rs`). If VAAPI is unavailable at runtime —
|
||||
//! `av_hwdevice_ctx_create` fails with no usable DRI device — `new` returns an
|
||||
//! error and [`super::best_decoder`] falls back to the portable openh264
|
||||
//! decoder, so headless / GPU-less Linux boxes still work.
|
||||
|
||||
use super::{Frame, VideoDecoder};
|
||||
use crate::error::{NdError, Result};
|
||||
|
||||
use core::ffi::{c_int, c_void};
|
||||
use ffmpeg_next::ffi::*;
|
||||
|
||||
/// FFmpeg/VAAPI hardware H.264 decoder.
|
||||
pub struct VaapiDecoder {
|
||||
ctx: *mut AVCodecContext,
|
||||
hw_device: *mut AVBufferRef,
|
||||
packet: *mut AVPacket,
|
||||
hw_frame: *mut AVFrame,
|
||||
sw_frame: *mut AVFrame,
|
||||
sws: *mut SwsContext,
|
||||
sws_dims: (i32, i32, i32), // (w, h, src pixel format) the cached scaler is for
|
||||
pts: i64,
|
||||
}
|
||||
|
||||
impl VaapiDecoder {
|
||||
pub fn new() -> Result<Self> {
|
||||
unsafe {
|
||||
// Safe to call repeatedly; wires up codec registration + logging.
|
||||
let _ = ffmpeg_next::init();
|
||||
|
||||
let codec = avcodec_find_decoder(AVCodecID::AV_CODEC_ID_H264);
|
||||
if codec.is_null() {
|
||||
return Err(other("libavcodec has no H.264 decoder"));
|
||||
}
|
||||
let ctx = avcodec_alloc_context3(codec);
|
||||
if ctx.is_null() {
|
||||
return Err(other("avcodec_alloc_context3 failed"));
|
||||
}
|
||||
|
||||
// Create the VAAPI device (default DRI render node) and attach it.
|
||||
let mut hw_device: *mut AVBufferRef = core::ptr::null_mut();
|
||||
let rc = av_hwdevice_ctx_create(
|
||||
&mut hw_device,
|
||||
AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
|
||||
core::ptr::null(),
|
||||
core::ptr::null_mut(),
|
||||
0,
|
||||
);
|
||||
let mut ctx = ctx;
|
||||
if rc < 0 || hw_device.is_null() {
|
||||
avcodec_free_context(&mut ctx);
|
||||
return Err(other(&format!(
|
||||
"av_hwdevice_ctx_create(VAAPI) failed ({rc})"
|
||||
)));
|
||||
}
|
||||
(*ctx).hw_device_ctx = av_buffer_ref(hw_device);
|
||||
(*ctx).get_format = Some(get_hw_format);
|
||||
|
||||
let rc = avcodec_open2(ctx, codec, core::ptr::null_mut());
|
||||
if rc < 0 {
|
||||
av_buffer_unref(&mut hw_device);
|
||||
avcodec_free_context(&mut ctx);
|
||||
return Err(other(&format!("avcodec_open2 failed ({rc})")));
|
||||
}
|
||||
|
||||
let mut packet = av_packet_alloc();
|
||||
let mut hw_frame = av_frame_alloc();
|
||||
let mut sw_frame = av_frame_alloc();
|
||||
if packet.is_null() || hw_frame.is_null() || sw_frame.is_null() {
|
||||
if !packet.is_null() {
|
||||
av_packet_free(&mut packet);
|
||||
}
|
||||
if !hw_frame.is_null() {
|
||||
av_frame_free(&mut hw_frame);
|
||||
}
|
||||
if !sw_frame.is_null() {
|
||||
av_frame_free(&mut sw_frame);
|
||||
}
|
||||
av_buffer_unref(&mut hw_device);
|
||||
avcodec_free_context(&mut ctx);
|
||||
return Err(other("av_frame_alloc / av_packet_alloc failed"));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
ctx,
|
||||
hw_device,
|
||||
packet,
|
||||
hw_frame,
|
||||
sw_frame,
|
||||
sws: core::ptr::null_mut(),
|
||||
sws_dims: (0, 0, -1),
|
||||
pts: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Transfer a decoded frame to system memory (if it is a GPU surface), crop
|
||||
/// it to the display rectangle, and swscale it to packed RGB8.
|
||||
unsafe fn frame_to_rgb(&mut self) -> Result<Option<Frame>> {
|
||||
// A VAAPI surface must be pulled down to a software NV12 frame first;
|
||||
// a software frame (fallback path) is used as-is.
|
||||
let src: *mut AVFrame =
|
||||
if (*self.hw_frame).format == AVPixelFormat::AV_PIX_FMT_VAAPI as c_int {
|
||||
av_frame_unref(self.sw_frame);
|
||||
let rc = av_hwframe_transfer_data(self.sw_frame, self.hw_frame, 0);
|
||||
if rc < 0 {
|
||||
return Err(other(&format!("av_hwframe_transfer_data failed ({rc})")));
|
||||
}
|
||||
// Carry crop/colour metadata across from the hardware frame.
|
||||
av_frame_copy_props(self.sw_frame, self.hw_frame);
|
||||
self.sw_frame
|
||||
} else {
|
||||
self.hw_frame
|
||||
};
|
||||
|
||||
// Apply the H.264 display crop so width/height are the true picture size
|
||||
// (e.g. 1920x1080) rather than the 16-aligned coded size.
|
||||
let _ = av_frame_apply_cropping(src, 1 /* AV_FRAME_CROP_UNALIGNED */);
|
||||
|
||||
let w = (*src).width;
|
||||
let h = (*src).height;
|
||||
let src_fmt = (*src).format;
|
||||
if w <= 0 || h <= 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// (Re)create the scaler when the geometry or source format changes.
|
||||
if self.sws.is_null() || self.sws_dims != (w, h, src_fmt) {
|
||||
if !self.sws.is_null() {
|
||||
sws_freeContext(self.sws);
|
||||
}
|
||||
self.sws = sws_getContext(
|
||||
w,
|
||||
h,
|
||||
core::mem::transmute::<c_int, AVPixelFormat>(src_fmt),
|
||||
w,
|
||||
h,
|
||||
AVPixelFormat::AV_PIX_FMT_RGB24,
|
||||
SWS_BILINEAR as c_int,
|
||||
core::ptr::null_mut(), // srcFilter: *mut SwsFilter
|
||||
core::ptr::null_mut(), // dstFilter: *mut SwsFilter
|
||||
core::ptr::null(), // param: *const c_double
|
||||
);
|
||||
if self.sws.is_null() {
|
||||
return Err(other("sws_getContext failed"));
|
||||
}
|
||||
self.sws_dims = (w, h, src_fmt);
|
||||
}
|
||||
|
||||
let mut rgb = vec![0u8; (w as usize) * (h as usize) * 3];
|
||||
let dst_data: [*mut u8; 4] = [
|
||||
rgb.as_mut_ptr(),
|
||||
core::ptr::null_mut(),
|
||||
core::ptr::null_mut(),
|
||||
core::ptr::null_mut(),
|
||||
];
|
||||
let dst_linesize: [c_int; 4] = [w * 3, 0, 0, 0];
|
||||
let scaled = sws_scale(
|
||||
self.sws,
|
||||
(*src).data.as_ptr() as *const *const u8,
|
||||
(*src).linesize.as_ptr(),
|
||||
0,
|
||||
h,
|
||||
dst_data.as_ptr(),
|
||||
dst_linesize.as_ptr(),
|
||||
);
|
||||
if scaled <= 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(Frame {
|
||||
w: w as usize,
|
||||
h: h as usize,
|
||||
rgb,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
impl VideoDecoder for VaapiDecoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
if au.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
unsafe {
|
||||
// Point the reusable packet at this AU. The data is not reference
|
||||
// counted (no `buf`), so libavcodec copies what it needs during
|
||||
// `avcodec_send_packet`; the slice only has to outlive this call.
|
||||
av_packet_unref(self.packet);
|
||||
(*self.packet).data = au.as_ptr() as *mut u8;
|
||||
(*self.packet).size = au.len() as c_int;
|
||||
(*self.packet).pts = self.pts;
|
||||
(*self.packet).dts = self.pts;
|
||||
self.pts += 1;
|
||||
|
||||
let mut out: Option<Frame> = None;
|
||||
|
||||
let mut sret = avcodec_send_packet(self.ctx, self.packet);
|
||||
if sret < 0 {
|
||||
// Decoder is full: drain a frame, keep it, then resend.
|
||||
av_frame_unref(self.hw_frame);
|
||||
if avcodec_receive_frame(self.ctx, self.hw_frame) == 0 {
|
||||
out = self.frame_to_rgb()?;
|
||||
}
|
||||
sret = avcodec_send_packet(self.ctx, self.packet);
|
||||
}
|
||||
// Detach our borrowed buffer from the packet before it is reused.
|
||||
(*self.packet).data = core::ptr::null_mut();
|
||||
(*self.packet).size = 0;
|
||||
if sret < 0 && out.is_none() {
|
||||
// EAGAIN / not-yet-ready before the first keyframe: no frame.
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
av_frame_unref(self.hw_frame);
|
||||
if avcodec_receive_frame(self.ctx, self.hw_frame) == 0 {
|
||||
out = self.frame_to_rgb()?;
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for VaapiDecoder {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
if !self.sws.is_null() {
|
||||
sws_freeContext(self.sws);
|
||||
}
|
||||
if !self.sw_frame.is_null() {
|
||||
av_frame_free(&mut self.sw_frame);
|
||||
}
|
||||
if !self.hw_frame.is_null() {
|
||||
av_frame_free(&mut self.hw_frame);
|
||||
}
|
||||
if !self.packet.is_null() {
|
||||
av_packet_free(&mut self.packet);
|
||||
}
|
||||
if !self.hw_device.is_null() {
|
||||
av_buffer_unref(&mut self.hw_device);
|
||||
}
|
||||
if !self.ctx.is_null() {
|
||||
avcodec_free_context(&mut self.ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SAFETY: the decoder is created on, and used only from, the single background
|
||||
// decode thread that owns it. The libavcodec context and its frames/packet are
|
||||
// never shared across threads here (no `Sync`); moving ownership to that thread
|
||||
// and calling it there is sound.
|
||||
unsafe impl Send for VaapiDecoder {}
|
||||
|
||||
/// `get_format` callback: prefer the VAAPI hardware pixel format when the
|
||||
/// decoder offers it, otherwise fall back to the first (software) format so a
|
||||
/// GPU-less box still decodes.
|
||||
unsafe extern "C" fn get_hw_format(
|
||||
_ctx: *mut AVCodecContext,
|
||||
mut fmts: *const AVPixelFormat,
|
||||
) -> AVPixelFormat {
|
||||
let first = if fmts.is_null() {
|
||||
AVPixelFormat::AV_PIX_FMT_NONE
|
||||
} else {
|
||||
*fmts
|
||||
};
|
||||
while !fmts.is_null() && *fmts != AVPixelFormat::AV_PIX_FMT_NONE {
|
||||
if *fmts == AVPixelFormat::AV_PIX_FMT_VAAPI {
|
||||
return AVPixelFormat::AV_PIX_FMT_VAAPI;
|
||||
}
|
||||
fmts = fmts.add(1);
|
||||
}
|
||||
first
|
||||
}
|
||||
|
||||
fn other(msg: &str) -> NdError {
|
||||
NdError::Other(anyhow::anyhow!(msg.to_string()))
|
||||
}
|
||||
@@ -0,0 +1,539 @@
|
||||
//! Apple hardware H.264 decode via VideoToolbox (`VTDecompressionSession`).
|
||||
//!
|
||||
//! On macOS and iOS this feeds the host's Annex-B access units into a
|
||||
//! `VTDecompressionSession`, which runs on the platform video decode hardware
|
||||
//! (the same engine `AVFoundation` uses). The decoder emits `CVPixelBuffer`s in
|
||||
//! 32-bit BGRA — VideoToolbox does the YUV→RGB colour conversion on the GPU for
|
||||
//! us — which we repack into the same packed-RGB8 [`Frame`] the openh264 path
|
||||
//! produces, so the viewer is unchanged.
|
||||
//!
|
||||
//! This whole module is only compiled on Apple targets (see the
|
||||
//! `#[cfg(any(target_os = "macos", target_os = "ios"))]` gate in `decode.rs`).
|
||||
//! If anything here fails at runtime, [`super::best_decoder`] falls back to the
|
||||
//! portable openh264 software decoder.
|
||||
//!
|
||||
//! ## Design
|
||||
//! * VideoToolbox needs a `CMVideoFormatDescription` built from the stream's SPS
|
||||
//! and PPS before a session can exist, so the session is created lazily on the
|
||||
//! first access unit that carries both (a keyframe). AUs seen before that are
|
||||
//! dropped (`Ok(None)`), exactly like the other backends while they prime.
|
||||
//! * The host sends Annex-B (start-code delimited) NAL units; VideoToolbox wants
|
||||
//! AVCC (4-byte length-prefixed) NALs in the `CMSampleBuffer`, so each AU is
|
||||
//! repacked. Parameter-set NALs go into the format description, not the sample.
|
||||
//! * Decode is driven synchronously: we submit the frame and immediately call
|
||||
//! `VTDecompressionSessionWaitForAsynchronousFrames`, so the output callback
|
||||
//! has run (and stored the frame in our heap-pinned [`Shared`] slot) by the
|
||||
//! time `decode` returns.
|
||||
//! * `CVPixelBufferGetWidth`/`Height` already return the *display* size (the
|
||||
//! decoder applies the SPS crop), so there is no padded-dimension problem to
|
||||
//! correct here the way there is on the Media Foundation path.
|
||||
|
||||
use super::{Frame, VideoDecoder};
|
||||
use crate::error::{NdError, Result};
|
||||
|
||||
use core::ffi::c_void;
|
||||
|
||||
use core_foundation::base::TCFType;
|
||||
use core_foundation::dictionary::CFDictionary;
|
||||
use core_foundation::number::CFNumber;
|
||||
use core_foundation::string::CFString;
|
||||
|
||||
use self::sys::*;
|
||||
|
||||
/// VideoToolbox hardware H.264 decoder.
|
||||
pub struct VideoToolboxDecoder {
|
||||
/// Decompression session; `null` until the first keyframe primes it.
|
||||
session: VTDecompressionSessionRef,
|
||||
/// Format description built from the current SPS/PPS; owned (+1 ref).
|
||||
format: CMFormatDescriptionRef,
|
||||
/// Cached parameter sets, so we only rebuild the session when they change.
|
||||
sps: Vec<u8>,
|
||||
pps: Vec<u8>,
|
||||
/// Heap-pinned output slot the C callback writes the decoded frame into.
|
||||
/// Boxed so its address is stable for the `refcon` handed to the session.
|
||||
shared: Box<Shared>,
|
||||
pts: i64,
|
||||
}
|
||||
|
||||
/// Output slot shared with the decompression callback via a raw `refcon`.
|
||||
struct Shared {
|
||||
frame: Option<Frame>,
|
||||
}
|
||||
|
||||
impl VideoToolboxDecoder {
|
||||
/// Prepare an (empty) decoder. The real session is built lazily once the
|
||||
/// first SPS+PPS arrive, so this never fails on its own.
|
||||
pub fn new() -> Result<Self> {
|
||||
Ok(Self {
|
||||
session: core::ptr::null_mut(),
|
||||
format: core::ptr::null_mut(),
|
||||
sps: Vec::new(),
|
||||
pps: Vec::new(),
|
||||
shared: Box::new(Shared { frame: None }),
|
||||
pts: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// (Re)build the format description and decompression session from the
|
||||
/// current `self.sps`/`self.pps`.
|
||||
fn rebuild_session(&mut self) -> Result<()> {
|
||||
self.teardown();
|
||||
|
||||
// CMVideoFormatDescription from the two parameter sets (raw NAL bytes,
|
||||
// no start codes), with a 4-byte NAL length prefix for the AVCC samples.
|
||||
let ptrs: [*const u8; 2] = [self.sps.as_ptr(), self.pps.as_ptr()];
|
||||
let sizes: [usize; 2] = [self.sps.len(), self.pps.len()];
|
||||
let mut format: CMFormatDescriptionRef = core::ptr::null_mut();
|
||||
let status = unsafe {
|
||||
CMVideoFormatDescriptionCreateFromH264ParameterSets(
|
||||
core::ptr::null(),
|
||||
2,
|
||||
ptrs.as_ptr(),
|
||||
sizes.as_ptr(),
|
||||
4,
|
||||
&mut format,
|
||||
)
|
||||
};
|
||||
if status != 0 || format.is_null() {
|
||||
return Err(other(&format!(
|
||||
"CMVideoFormatDescriptionCreateFromH264ParameterSets failed ({status})"
|
||||
)));
|
||||
}
|
||||
self.format = format;
|
||||
|
||||
// Ask VideoToolbox to hand us ready-to-blit 32-bit BGRA pixel buffers.
|
||||
let key = unsafe { CFString::wrap_under_get_rule(kCVPixelBufferPixelFormatTypeKey) };
|
||||
let val = CFNumber::from(kCVPixelFormatType_32BGRA);
|
||||
let attrs = CFDictionary::from_CFType_pairs(&[(key.as_CFType(), val.as_CFType())]);
|
||||
|
||||
let mut record = VTDecompressionOutputCallbackRecord {
|
||||
decompressionOutputCallback: Some(decompression_callback),
|
||||
decompressionOutputRefCon: (&mut *self.shared) as *mut Shared as *mut c_void,
|
||||
};
|
||||
let mut session: VTDecompressionSessionRef = core::ptr::null_mut();
|
||||
let status = unsafe {
|
||||
VTDecompressionSessionCreate(
|
||||
core::ptr::null(),
|
||||
self.format,
|
||||
core::ptr::null(), // default decoder (hardware when available)
|
||||
attrs.as_concrete_TypeRef() as _, // destination buffer attributes
|
||||
&record,
|
||||
&mut session,
|
||||
)
|
||||
};
|
||||
// `attrs`/`record` are retained/copied by the call above; drop them now.
|
||||
let _ = record.decompressionOutputRefCon;
|
||||
if status != 0 || session.is_null() {
|
||||
return Err(other(&format!(
|
||||
"VTDecompressionSessionCreate failed ({status})"
|
||||
)));
|
||||
}
|
||||
self.session = session;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Decode one AVCC-packed sample through the session, returning its frame.
|
||||
fn decode_sample(&mut self, avcc: &[u8]) -> Result<Option<Frame>> {
|
||||
if self.session.is_null() || avcc.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Wrap the AVCC bytes in a CMBlockBuffer (which owns a private copy).
|
||||
let mut block: CMBlockBufferRef = core::ptr::null_mut();
|
||||
let status = unsafe {
|
||||
CMBlockBufferCreateWithMemoryBlock(
|
||||
core::ptr::null(),
|
||||
core::ptr::null_mut(),
|
||||
avcc.len(),
|
||||
core::ptr::null(), // default block allocator
|
||||
core::ptr::null(),
|
||||
0,
|
||||
avcc.len(),
|
||||
kCMBlockBufferAssureMemoryNowFlag,
|
||||
&mut block,
|
||||
)
|
||||
};
|
||||
if status != 0 || block.is_null() {
|
||||
return Err(other(&format!(
|
||||
"CMBlockBufferCreateWithMemoryBlock ({status})"
|
||||
)));
|
||||
}
|
||||
let status = unsafe {
|
||||
CMBlockBufferReplaceDataBytes(avcc.as_ptr() as *const c_void, block, 0, avcc.len())
|
||||
};
|
||||
if status != 0 {
|
||||
unsafe { CFRelease(block as _) };
|
||||
return Err(other(&format!("CMBlockBufferReplaceDataBytes ({status})")));
|
||||
}
|
||||
|
||||
// Build the sample buffer with timing and size.
|
||||
let timing = CMSampleTimingInfo {
|
||||
duration: CMTime {
|
||||
value: HNS_PER_FRAME,
|
||||
timescale: 10_000_000,
|
||||
flags: K_CMTIME_FLAGS_VALID,
|
||||
epoch: 0,
|
||||
},
|
||||
presentationTimeStamp: CMTime {
|
||||
value: self.pts,
|
||||
timescale: 10_000_000,
|
||||
flags: K_CMTIME_FLAGS_VALID,
|
||||
epoch: 0,
|
||||
},
|
||||
decodeTimeStamp: CMTime {
|
||||
value: 0,
|
||||
timescale: 0,
|
||||
flags: 0,
|
||||
epoch: 0,
|
||||
},
|
||||
};
|
||||
self.pts += HNS_PER_FRAME;
|
||||
let sizes = [avcc.len()];
|
||||
let mut sample: CMSampleBufferRef = core::ptr::null_mut();
|
||||
let status = unsafe {
|
||||
CMSampleBufferCreateReady(
|
||||
core::ptr::null(),
|
||||
block,
|
||||
self.format,
|
||||
1,
|
||||
1,
|
||||
&timing,
|
||||
1,
|
||||
sizes.as_ptr(),
|
||||
&mut sample,
|
||||
)
|
||||
};
|
||||
unsafe { CFRelease(block as _) };
|
||||
if status != 0 || sample.is_null() {
|
||||
return Err(other(&format!("CMSampleBufferCreateReady ({status})")));
|
||||
}
|
||||
|
||||
// Submit synchronously: the callback runs before Wait returns.
|
||||
self.shared.frame = None;
|
||||
let mut info_flags: u32 = 0;
|
||||
let status = unsafe {
|
||||
VTDecompressionSessionDecodeFrame(
|
||||
self.session,
|
||||
sample,
|
||||
0,
|
||||
core::ptr::null_mut(),
|
||||
&mut info_flags,
|
||||
)
|
||||
};
|
||||
if status == 0 {
|
||||
unsafe { VTDecompressionSessionWaitForAsynchronousFrames(self.session) };
|
||||
}
|
||||
unsafe { CFRelease(sample as _) };
|
||||
if status != 0 {
|
||||
return Err(other(&format!(
|
||||
"VTDecompressionSessionDecodeFrame ({status})"
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(self.shared.frame.take())
|
||||
}
|
||||
|
||||
fn teardown(&mut self) {
|
||||
unsafe {
|
||||
if !self.session.is_null() {
|
||||
VTDecompressionSessionInvalidate(self.session);
|
||||
CFRelease(self.session as _);
|
||||
self.session = core::ptr::null_mut();
|
||||
}
|
||||
if !self.format.is_null() {
|
||||
CFRelease(self.format as _);
|
||||
self.format = core::ptr::null_mut();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl VideoDecoder for VideoToolboxDecoder {
|
||||
fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
|
||||
if au.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Split the Annex-B AU, pulling out parameter sets and repacking the
|
||||
// decodable NALs as AVCC (4-byte big-endian length prefixes).
|
||||
let mut new_sps: Option<Vec<u8>> = None;
|
||||
let mut new_pps: Option<Vec<u8>> = None;
|
||||
let mut avcc: Vec<u8> = Vec::with_capacity(au.len() + 16);
|
||||
for_each_nal(au, |nal| {
|
||||
match nal.first().map(|b| b & 0x1F) {
|
||||
Some(7) => new_sps = Some(nal.to_vec()), // SPS
|
||||
Some(8) => new_pps = Some(nal.to_vec()), // PPS
|
||||
Some(9) => {} // AUD: drop
|
||||
Some(_) => {
|
||||
avcc.extend_from_slice(&(nal.len() as u32).to_be_bytes());
|
||||
avcc.extend_from_slice(nal);
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
});
|
||||
|
||||
// A keyframe carries fresh parameter sets; rebuild the session when they
|
||||
// change (first keyframe, or a resolution/profile switch mid-stream).
|
||||
if let (Some(sps), Some(pps)) = (new_sps, new_pps) {
|
||||
if sps != self.sps || pps != self.pps || self.session.is_null() {
|
||||
self.sps = sps;
|
||||
self.pps = pps;
|
||||
self.rebuild_session()?;
|
||||
}
|
||||
}
|
||||
|
||||
self.decode_sample(&avcc)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for VideoToolboxDecoder {
|
||||
fn drop(&mut self) {
|
||||
self.teardown();
|
||||
}
|
||||
}
|
||||
|
||||
// SAFETY: the decoder is created on, and only ever touched from, the single
|
||||
// background decode thread that owns it (the viewer's network loop). The
|
||||
// VideoToolbox session is used synchronously and never shared across threads
|
||||
// (no `Sync`), and the `Shared` slot the callback writes is reached only while
|
||||
// that same thread is blocked inside `WaitForAsynchronousFrames`.
|
||||
unsafe impl Send for VideoToolboxDecoder {}
|
||||
|
||||
/// ~29.97 fps in 100-ns units; only needs to be monotonic for this stream.
|
||||
const HNS_PER_FRAME: i64 = 333_667;
|
||||
|
||||
/// Iterate the Annex-B NAL units in `au`, trimming trailing zero bytes (which
|
||||
/// belong to the next start code / are RBSP padding) so each slice is the raw
|
||||
/// NAL, start-code excluded.
|
||||
fn for_each_nal<F: FnMut(&[u8])>(au: &[u8], mut f: F) {
|
||||
let n = au.len();
|
||||
// Payload start offset of every `00 00 01` start code (a 4-byte start code
|
||||
// `00 00 00 01` is matched on its trailing three bytes, the extra leading
|
||||
// zero falling into the previous NAL's trailing-zero trim).
|
||||
let mut starts: Vec<usize> = Vec::new();
|
||||
let mut i = 0usize;
|
||||
while i + 2 < n {
|
||||
if au[i] == 0 && au[i + 1] == 0 && au[i + 2] == 1 {
|
||||
starts.push(i + 3);
|
||||
i += 3;
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
for k in 0..starts.len() {
|
||||
let begin = starts[k];
|
||||
let end = if k + 1 < starts.len() {
|
||||
starts[k + 1] - 3
|
||||
} else {
|
||||
n
|
||||
};
|
||||
if begin >= end {
|
||||
continue;
|
||||
}
|
||||
let mut nal = &au[begin..end];
|
||||
while let [rest @ .., 0] = nal {
|
||||
nal = rest;
|
||||
}
|
||||
if !nal.is_empty() {
|
||||
f(nal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// VideoToolbox output callback: convert the decoded BGRA pixel buffer to a
|
||||
/// packed-RGB8 [`Frame`] and stash it in the shared slot pointed to by `refcon`.
|
||||
extern "C" fn decompression_callback(
|
||||
refcon: *mut c_void,
|
||||
_source_frame_refcon: *mut c_void,
|
||||
status: i32,
|
||||
_info_flags: u32,
|
||||
image_buffer: CVImageBufferRef,
|
||||
_pts: CMTime,
|
||||
_duration: CMTime,
|
||||
) {
|
||||
if refcon.is_null() || status != 0 || image_buffer.is_null() {
|
||||
return;
|
||||
}
|
||||
let shared = unsafe { &mut *(refcon as *mut Shared) };
|
||||
let pb = image_buffer; // CVImageBufferRef and CVPixelBufferRef are the same handle
|
||||
unsafe {
|
||||
if CVPixelBufferLockBaseAddress(pb, K_CVPIXELBUFFER_LOCK_READ_ONLY) != 0 {
|
||||
return;
|
||||
}
|
||||
let w = CVPixelBufferGetWidth(pb);
|
||||
let h = CVPixelBufferGetHeight(pb);
|
||||
let stride = CVPixelBufferGetBytesPerRow(pb);
|
||||
let base = CVPixelBufferGetBaseAddress(pb) as *const u8;
|
||||
if !base.is_null() && w > 0 && h > 0 && stride >= w * 4 {
|
||||
let src = core::slice::from_raw_parts(base, stride * h);
|
||||
let mut rgb = vec![0u8; w * h * 3];
|
||||
for row in 0..h {
|
||||
let srow = row * stride;
|
||||
let drow = row * w * 3;
|
||||
for col in 0..w {
|
||||
let s = srow + col * 4; // BGRA
|
||||
let d = drow + col * 3;
|
||||
rgb[d] = src[s + 2]; // R
|
||||
rgb[d + 1] = src[s + 1]; // G
|
||||
rgb[d + 2] = src[s]; // B
|
||||
}
|
||||
}
|
||||
shared.frame = Some(Frame { w, h, rgb });
|
||||
}
|
||||
CVPixelBufferUnlockBaseAddress(pb, K_CVPIXELBUFFER_LOCK_READ_ONLY);
|
||||
}
|
||||
}
|
||||
|
||||
fn other(msg: &str) -> NdError {
|
||||
NdError::Other(anyhow::anyhow!(msg.to_string()))
|
||||
}
|
||||
|
||||
/// Minimal, self-contained FFI surface for the CoreMedia / CoreVideo /
|
||||
/// VideoToolbox C APIs this decoder uses. These ABIs are stable across macOS and
|
||||
/// iOS releases; the `CFDictionaryRef`/`CFStringRef` types are shared with the
|
||||
/// `core-foundation` crate so the attribute dictionary bridges cleanly.
|
||||
#[allow(non_snake_case, non_upper_case_globals)]
|
||||
mod sys {
|
||||
use core::ffi::c_void;
|
||||
use core_foundation_sys::base::{CFAllocatorRef, CFRelease as cf_release, CFTypeRef};
|
||||
use core_foundation_sys::dictionary::CFDictionaryRef;
|
||||
use core_foundation_sys::string::CFStringRef;
|
||||
|
||||
pub type OSStatus = i32;
|
||||
pub type CMFormatDescriptionRef = *mut c_void;
|
||||
pub type CMVideoFormatDescriptionRef = CMFormatDescriptionRef;
|
||||
pub type CMBlockBufferRef = *mut c_void;
|
||||
pub type CMSampleBufferRef = *mut c_void;
|
||||
pub type VTDecompressionSessionRef = *mut c_void;
|
||||
pub type CVImageBufferRef = *mut c_void;
|
||||
pub type CVPixelBufferRef = *mut c_void;
|
||||
|
||||
/// `kCVPixelFormatType_32BGRA` ('BGRA').
|
||||
pub const kCVPixelFormatType_32BGRA: i32 = 0x4247_5241;
|
||||
/// `kCVPixelBufferLock_ReadOnly`.
|
||||
pub const K_CVPIXELBUFFER_LOCK_READ_ONLY: u64 = 0x0000_0001;
|
||||
/// `kCMTimeFlags_Valid`.
|
||||
pub const K_CMTIME_FLAGS_VALID: u32 = 0x0000_0001;
|
||||
/// `kCMBlockBufferAssureMemoryNowFlag`.
|
||||
pub const kCMBlockBufferAssureMemoryNowFlag: u32 = 0x0000_0001;
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct CMTime {
|
||||
pub value: i64,
|
||||
pub timescale: i32,
|
||||
pub flags: u32,
|
||||
pub epoch: i64,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct CMSampleTimingInfo {
|
||||
pub duration: CMTime,
|
||||
pub presentationTimeStamp: CMTime,
|
||||
pub decodeTimeStamp: CMTime,
|
||||
}
|
||||
|
||||
pub type VTDecompressionOutputCallback = extern "C" fn(
|
||||
refcon: *mut c_void,
|
||||
source_frame_refcon: *mut c_void,
|
||||
status: OSStatus,
|
||||
info_flags: u32,
|
||||
image_buffer: CVImageBufferRef,
|
||||
pts: CMTime,
|
||||
duration: CMTime,
|
||||
);
|
||||
|
||||
#[repr(C)]
|
||||
pub struct VTDecompressionOutputCallbackRecord {
|
||||
pub decompressionOutputCallback: Option<VTDecompressionOutputCallback>,
|
||||
pub decompressionOutputRefCon: *mut c_void,
|
||||
}
|
||||
|
||||
/// Re-export CoreFoundation's `CFRelease` under the name this module uses.
|
||||
#[inline]
|
||||
pub unsafe fn CFRelease(obj: CFTypeRef) {
|
||||
cf_release(obj)
|
||||
}
|
||||
|
||||
#[link(name = "CoreMedia", kind = "framework")]
|
||||
extern "C" {
|
||||
pub fn CMVideoFormatDescriptionCreateFromH264ParameterSets(
|
||||
allocator: CFAllocatorRef,
|
||||
parameter_set_count: usize,
|
||||
parameter_set_pointers: *const *const u8,
|
||||
parameter_set_sizes: *const usize,
|
||||
nal_unit_header_length: i32,
|
||||
format_description_out: *mut CMFormatDescriptionRef,
|
||||
) -> OSStatus;
|
||||
|
||||
pub fn CMBlockBufferCreateWithMemoryBlock(
|
||||
structure_allocator: CFAllocatorRef,
|
||||
memory_block: *mut c_void,
|
||||
block_length: usize,
|
||||
block_allocator: CFAllocatorRef,
|
||||
custom_block_source: *const c_void,
|
||||
offset_to_data: usize,
|
||||
data_length: usize,
|
||||
flags: u32,
|
||||
block_buffer_out: *mut CMBlockBufferRef,
|
||||
) -> OSStatus;
|
||||
|
||||
pub fn CMBlockBufferReplaceDataBytes(
|
||||
source_bytes: *const c_void,
|
||||
destination_buffer: CMBlockBufferRef,
|
||||
offset_into_destination: usize,
|
||||
data_length: usize,
|
||||
) -> OSStatus;
|
||||
|
||||
pub fn CMSampleBufferCreateReady(
|
||||
allocator: CFAllocatorRef,
|
||||
data_buffer: CMBlockBufferRef,
|
||||
format_description: CMFormatDescriptionRef,
|
||||
num_samples: isize,
|
||||
num_sample_timing_entries: isize,
|
||||
sample_timing_array: *const CMSampleTimingInfo,
|
||||
num_sample_size_entries: isize,
|
||||
sample_size_array: *const usize,
|
||||
sample_buffer_out: *mut CMSampleBufferRef,
|
||||
) -> OSStatus;
|
||||
}
|
||||
|
||||
#[link(name = "VideoToolbox", kind = "framework")]
|
||||
extern "C" {
|
||||
pub fn VTDecompressionSessionCreate(
|
||||
allocator: CFAllocatorRef,
|
||||
video_format_description: CMVideoFormatDescriptionRef,
|
||||
video_decoder_specification: CFDictionaryRef,
|
||||
destination_image_buffer_attributes: CFDictionaryRef,
|
||||
output_callback: *const VTDecompressionOutputCallbackRecord,
|
||||
decompression_session_out: *mut VTDecompressionSessionRef,
|
||||
) -> OSStatus;
|
||||
|
||||
pub fn VTDecompressionSessionDecodeFrame(
|
||||
session: VTDecompressionSessionRef,
|
||||
sample_buffer: CMSampleBufferRef,
|
||||
decode_flags: u32,
|
||||
source_frame_refcon: *mut c_void,
|
||||
info_flags_out: *mut u32,
|
||||
) -> OSStatus;
|
||||
|
||||
pub fn VTDecompressionSessionWaitForAsynchronousFrames(
|
||||
session: VTDecompressionSessionRef,
|
||||
) -> OSStatus;
|
||||
|
||||
pub fn VTDecompressionSessionInvalidate(session: VTDecompressionSessionRef);
|
||||
}
|
||||
|
||||
#[link(name = "CoreVideo", kind = "framework")]
|
||||
extern "C" {
|
||||
pub static kCVPixelBufferPixelFormatTypeKey: CFStringRef;
|
||||
|
||||
pub fn CVPixelBufferLockBaseAddress(pb: CVPixelBufferRef, flags: u64) -> i32;
|
||||
pub fn CVPixelBufferUnlockBaseAddress(pb: CVPixelBufferRef, flags: u64) -> i32;
|
||||
pub fn CVPixelBufferGetBaseAddress(pb: CVPixelBufferRef) -> *mut c_void;
|
||||
pub fn CVPixelBufferGetBytesPerRow(pb: CVPixelBufferRef) -> usize;
|
||||
pub fn CVPixelBufferGetWidth(pb: CVPixelBufferRef) -> usize;
|
||||
pub fn CVPixelBufferGetHeight(pb: CVPixelBufferRef) -> usize;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum NdError {
|
||||
#[error("connection failed: {0}")]
|
||||
Connect(String),
|
||||
#[error("protocol error: {0}")]
|
||||
Protocol(String),
|
||||
#[error("authentication failed: {0}")]
|
||||
Auth(String),
|
||||
#[error("session denied: {0}")]
|
||||
Denied(String),
|
||||
#[error("crypto error: {0}")]
|
||||
Crypto(String),
|
||||
#[error("websocket closed")]
|
||||
Closed,
|
||||
#[error(transparent)]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error(transparent)]
|
||||
Other(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, NdError>;
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,205 @@
|
||||
//! LAN peer discovery over UDP broadcast.
|
||||
//!
|
||||
//! Mirrors RustDesk's `src/lan.rs` query/reply model: a viewer broadcasts a
|
||||
//! small query to the LAN; every reachable host replies with its id, name and
|
||||
//! the port it accepts a direct connection on. The viewer learns each host's IP
|
||||
//! from the reply's source address — so two machines on the same subnet can find
|
||||
//! each other with no rendezvous/signaling server involved (works even when the
|
||||
//! internet is down).
|
||||
//!
|
||||
//! This module is the discovery half; establishing the session over the
|
||||
//! discovered `ip:port` (a direct transport, bypassing the relay) is the wiring
|
||||
//! step that builds on it. Packets are small JSON for debuggability.
|
||||
//!
|
||||
//! Developed by Akronlogic.
|
||||
|
||||
use std::net::{SocketAddr, UdpSocket};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Default UDP port for NetDesk LAN discovery.
|
||||
pub const LAN_PORT: u16 = 28118;
|
||||
|
||||
/// A host discovered on the LAN.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LanHost {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
/// Where to reach the host directly: its LAN IP (from the reply source) and
|
||||
/// the direct-connect port it advertised.
|
||||
pub addr: SocketAddr,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct Query<'a> {
|
||||
t: &'a str, // "netdesk.lan.query"
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct Announce {
|
||||
t: String, // "netdesk.lan.host"
|
||||
id: String,
|
||||
name: String,
|
||||
port: u16, // direct-connect port
|
||||
}
|
||||
|
||||
const QUERY_T: &str = "netdesk.lan.query";
|
||||
const HOST_T: &str = "netdesk.lan.host";
|
||||
|
||||
/// The targets a production scan queries: the global broadcast address plus
|
||||
/// loopback (so a host on this same machine is found even where broadcast to
|
||||
/// self is unreliable). A caller may add subnet-directed broadcasts too.
|
||||
pub fn broadcast_targets() -> Vec<SocketAddr> {
|
||||
vec![
|
||||
SocketAddr::from(([255, 255, 255, 255], LAN_PORT)),
|
||||
SocketAddr::from(([127, 0, 0, 1], LAN_PORT)),
|
||||
]
|
||||
}
|
||||
|
||||
/// Host side: answer discovery queries until `stop` is set. Binds UDP
|
||||
/// `0.0.0.0:bind_port` and replies to each query with this host's id/name and
|
||||
/// `direct_port`. Best-effort: returns on a fatal bind error (discovery just
|
||||
/// won't work); transient recv errors are ignored.
|
||||
pub fn respond_to_queries(
|
||||
bind_port: u16,
|
||||
host_id: String,
|
||||
name: String,
|
||||
direct_port: u16,
|
||||
stop: Arc<AtomicBool>,
|
||||
) -> std::io::Result<()> {
|
||||
let socket = UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], bind_port)))?;
|
||||
socket.set_read_timeout(Some(Duration::from_millis(300)))?;
|
||||
let reply = serde_json::to_vec(&Announce {
|
||||
t: HOST_T.to_string(),
|
||||
id: host_id,
|
||||
name,
|
||||
port: direct_port,
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut buf = [0u8; 2048];
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
match socket.recv_from(&mut buf) {
|
||||
Ok((len, from)) => {
|
||||
if is_query(&buf[..len]) {
|
||||
let _ = socket.send_to(&reply, from);
|
||||
}
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => continue,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::TimedOut => continue,
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_query(buf: &[u8]) -> bool {
|
||||
serde_json::from_slice::<Query>(buf)
|
||||
.map(|q| q.t == QUERY_T)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Viewer side: broadcast a discovery query to each `target` and collect host
|
||||
/// replies for `timeout`. Production passes the LAN broadcast address
|
||||
/// (`255.255.255.255:LAN_PORT`); tests pass a specific loopback address.
|
||||
/// Deduplicates by host id (first reply wins).
|
||||
pub fn scan(targets: &[SocketAddr], timeout: Duration) -> std::io::Result<Vec<LanHost>> {
|
||||
let socket = UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], 0)))?;
|
||||
socket.set_broadcast(true).ok(); // needed for 255.255.255.255; harmless otherwise
|
||||
socket.set_read_timeout(Some(Duration::from_millis(200)))?;
|
||||
|
||||
let query = serde_json::to_vec(&Query { t: QUERY_T }).unwrap_or_default();
|
||||
let send_all = || {
|
||||
for t in targets {
|
||||
let _ = socket.send_to(&query, t);
|
||||
}
|
||||
};
|
||||
send_all();
|
||||
|
||||
let mut hosts: Vec<LanHost> = Vec::new();
|
||||
let deadline = Instant::now() + timeout;
|
||||
let mut next_resend = Instant::now() + Duration::from_millis(150);
|
||||
let mut buf = [0u8; 2048];
|
||||
while Instant::now() < deadline {
|
||||
// Re-broadcast periodically: a single UDP query can be lost, and the
|
||||
// responder may not have bound yet when we started.
|
||||
if Instant::now() >= next_resend {
|
||||
send_all();
|
||||
next_resend = Instant::now() + Duration::from_millis(150);
|
||||
}
|
||||
match socket.recv_from(&mut buf) {
|
||||
Ok((len, from)) => {
|
||||
if let Ok(a) = serde_json::from_slice::<Announce>(&buf[..len]) {
|
||||
if a.t == HOST_T && !hosts.iter().any(|h| h.id == a.id) {
|
||||
hosts.push(LanHost {
|
||||
id: a.id,
|
||||
name: a.name,
|
||||
addr: SocketAddr::new(from.ip(), a.port),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => continue,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::TimedOut => continue,
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
Ok(hosts)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn broadcast_targets_are_global_and_loopback() {
|
||||
let t = broadcast_targets();
|
||||
assert!(t.iter().any(|a| a.ip().to_string() == "255.255.255.255" && a.port() == LAN_PORT));
|
||||
assert!(t.iter().any(|a| a.ip().to_string() == "127.0.0.1" && a.port() == LAN_PORT));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scan_finds_a_responding_host() {
|
||||
// A unique port so parallel tests don't collide.
|
||||
let port = 29871u16;
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let s2 = stop.clone();
|
||||
let host = std::thread::spawn(move || {
|
||||
let _ = respond_to_queries(
|
||||
port,
|
||||
"lan-host-123".to_string(),
|
||||
"Test Host".to_string(),
|
||||
41000,
|
||||
s2,
|
||||
);
|
||||
});
|
||||
// Give the responder a moment to bind.
|
||||
std::thread::sleep(Duration::from_millis(150));
|
||||
|
||||
// Query loopback directly (broadcast is unreliable on loopback in CI).
|
||||
let targets = [SocketAddr::from(([127, 0, 0, 1], port))];
|
||||
let found = scan(&targets, Duration::from_millis(600)).unwrap();
|
||||
|
||||
stop.store(true, Ordering::Relaxed);
|
||||
let _ = host.join();
|
||||
|
||||
let h = found
|
||||
.iter()
|
||||
.find(|h| h.id == "lan-host-123")
|
||||
.expect("host not discovered on the LAN");
|
||||
assert_eq!(h.name, "Test Host");
|
||||
assert_eq!(h.addr.port(), 41000, "advertised direct-connect port");
|
||||
assert!(h.addr.ip().is_loopback(), "host IP learned from reply source");
|
||||
println!("LAN_DISCOVERY_OK found {} host(s): {:?}", found.len(), h);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_query_packets_are_ignored() {
|
||||
assert!(is_query(br#"{"t":"netdesk.lan.query"}"#));
|
||||
assert!(!is_query(br#"{"t":"something.else"}"#));
|
||||
assert!(!is_query(b"garbage"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
//! NetDesk native core — the compiled Rust engine behind the cross-platform
|
||||
//! client (Windows / macOS / Linux / iOS / Android). Developed by Akronlogic.
|
||||
//!
|
||||
//! This crate owns everything that must be fast and portable: signaling,
|
||||
//! end-to-end crypto, the media transport, and (in later phases) H.264 decode
|
||||
//! and input injection. The Flutter UI calls in over the C ABI in `ffi`.
|
||||
|
||||
pub mod brand;
|
||||
pub mod cidr;
|
||||
pub mod crypto;
|
||||
pub mod decode;
|
||||
pub mod error;
|
||||
pub mod host;
|
||||
pub mod lan;
|
||||
pub mod netqos;
|
||||
pub mod p2p;
|
||||
pub mod proto;
|
||||
pub mod recorder;
|
||||
pub mod rtc;
|
||||
pub mod webm;
|
||||
pub mod signaling;
|
||||
pub mod transport;
|
||||
pub mod update;
|
||||
|
||||
pub use error::{NdError, Result};
|
||||
pub use signaling::{connect, ConnectConfig, Session, SessionEvent};
|
||||
pub use p2p::{upgrade_relay, P2pControl, P2pRelay, P2pRole};
|
||||
pub use transport::{Lane, Transport, TransportMsg, WebRtcChannel, WsRelay};
|
||||
pub use rtc::{advertise_max_message_size, answer_offer_web, build_api_web, create_web_channels, upgrade_offerer_web, WebChannels, WebDetached};
|
||||
|
||||
/// Core/build version, surfaced in the UI "about" box.
|
||||
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
pub const BRAND: &str = "NetDesk";
|
||||
pub const VENDOR: &str = "Akronlogic";
|
||||
@@ -0,0 +1,223 @@
|
||||
//! Network-quality estimators for the adaptive QoS controller — pure +
|
||||
//! deterministic (feed samples, read values), so they unit-test without a live
|
||||
//! link. RustDesk-style signals that replace the old "absolute RTT proxy":
|
||||
//!
|
||||
//! - [`RttCalc`] — a rolling-baseline RTT: the minimum RTT over the recent
|
||||
//! window is the link's inherent latency, and the reported delay is how far
|
||||
//! the latest sample sits ABOVE it. A stable high-latency link reads ~0
|
||||
//! delay; only latency SPIKES (queue buildup) read high. (RustDesk's
|
||||
//! `RttCalculator`.)
|
||||
//! - [`LossCalc`] — packet-loss ratio from ping/pong hit/miss over a window.
|
||||
//! - [`Jitter`] — the viewer measures its frame-arrival jitter (the mean
|
||||
//! absolute deviation of inter-arrival intervals). Smooth delivery ~0;
|
||||
//! bursty/stalled delivery (the decode pipeline falling behind) grows. This
|
||||
//! is a congestion signal RTT alone misses, and it is reported to the host in
|
||||
//! `net.pong` so `Qos::on_sample` gets a real end-to-end delay.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
|
||||
/// Rolling-baseline RTT delay. `add` returns the delay of the new sample above
|
||||
/// the window minimum (the baseline).
|
||||
#[derive(Debug)]
|
||||
pub struct RttCalc {
|
||||
win: VecDeque<u32>,
|
||||
cap: usize,
|
||||
last: u32,
|
||||
}
|
||||
|
||||
impl Default for RttCalc {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl RttCalc {
|
||||
/// ~30 samples (≈30 s at one probe/second) of history for the baseline.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
win: VecDeque::new(),
|
||||
cap: 30,
|
||||
last: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Record an RTT sample (ms); returns the delay above the rolling baseline.
|
||||
pub fn add(&mut self, rtt_ms: u32) -> u32 {
|
||||
self.win.push_back(rtt_ms);
|
||||
while self.win.len() > self.cap {
|
||||
self.win.pop_front();
|
||||
}
|
||||
let base = *self.win.iter().min().unwrap_or(&rtt_ms);
|
||||
self.last = rtt_ms.saturating_sub(base);
|
||||
self.last
|
||||
}
|
||||
|
||||
/// The most recent delay-above-baseline.
|
||||
pub fn delay(&self) -> u32 {
|
||||
self.last
|
||||
}
|
||||
}
|
||||
|
||||
/// Packet-loss ratio over a rolling window of ping/pong outcomes.
|
||||
#[derive(Debug)]
|
||||
pub struct LossCalc {
|
||||
win: VecDeque<bool>,
|
||||
cap: usize,
|
||||
}
|
||||
|
||||
impl Default for LossCalc {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl LossCalc {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
win: VecDeque::new(),
|
||||
cap: 20,
|
||||
}
|
||||
}
|
||||
|
||||
/// Record one outcome (`true` = the pong was lost); returns the loss ratio.
|
||||
pub fn record(&mut self, lost: bool) -> f32 {
|
||||
self.win.push_back(lost);
|
||||
while self.win.len() > self.cap {
|
||||
self.win.pop_front();
|
||||
}
|
||||
let n = self.win.len().max(1);
|
||||
self.win.iter().filter(|&&l| l).count() as f32 / n as f32
|
||||
}
|
||||
}
|
||||
|
||||
/// Viewer-side frame-arrival jitter: the mean absolute deviation of recent
|
||||
/// inter-arrival intervals (ms). fps-independent (it measures spread, not the
|
||||
/// absolute cadence), so lowering the frame rate does not read as congestion —
|
||||
/// only bursty / stalled delivery does.
|
||||
#[derive(Debug)]
|
||||
pub struct Jitter {
|
||||
intervals: VecDeque<u32>,
|
||||
cap: usize,
|
||||
last_ms: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for Jitter {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Jitter {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
intervals: VecDeque::new(),
|
||||
cap: 30,
|
||||
last_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Record a frame arrival at monotonic `now_ms`; returns the current jitter.
|
||||
pub fn frame(&mut self, now_ms: u64) -> u32 {
|
||||
if let Some(prev) = self.last_ms {
|
||||
// Clamp absurd gaps (a static screen sends nothing for a while) so a
|
||||
// long idle doesn't poison the jitter once frames resume.
|
||||
let dt = now_ms.saturating_sub(prev).min(2_000) as u32;
|
||||
self.intervals.push_back(dt);
|
||||
while self.intervals.len() > self.cap {
|
||||
self.intervals.pop_front();
|
||||
}
|
||||
}
|
||||
self.last_ms = Some(now_ms);
|
||||
self.jitter()
|
||||
}
|
||||
|
||||
/// Current jitter (mean absolute deviation of intervals, ms).
|
||||
pub fn jitter(&self) -> u32 {
|
||||
let n = self.intervals.len();
|
||||
if n < 2 {
|
||||
return 0;
|
||||
}
|
||||
let mean = self.intervals.iter().sum::<u32>() as f64 / n as f64;
|
||||
let mad = self
|
||||
.intervals
|
||||
.iter()
|
||||
.map(|&x| (x as f64 - mean).abs())
|
||||
.sum::<f64>()
|
||||
/ n as f64;
|
||||
mad.round() as u32
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn rtt_baseline_reports_delay_above_min() {
|
||||
let mut r = RttCalc::new();
|
||||
assert_eq!(r.add(200), 0, "first sample is its own baseline");
|
||||
assert_eq!(r.add(200), 0);
|
||||
assert_eq!(r.add(205), 5, "5 ms above the 200 ms floor");
|
||||
assert_eq!(r.add(350), 150, "a spike reads 150 ms above baseline");
|
||||
assert_eq!(r.add(200), 0, "back at the floor -> no delay");
|
||||
assert_eq!(r.delay(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rtt_baseline_rises_when_the_floor_rises() {
|
||||
let mut r = RttCalc::new();
|
||||
// Fill the window so old low samples roll out.
|
||||
for _ in 0..30 {
|
||||
r.add(100);
|
||||
}
|
||||
// A sustained worse link: once the 100s roll out, 300 becomes the floor.
|
||||
for _ in 0..30 {
|
||||
r.add(300);
|
||||
}
|
||||
assert_eq!(r.add(300), 0, "a sustained new floor reads ~0 delay again");
|
||||
assert_eq!(r.add(450), 150, "spike above the new 300 ms floor");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loss_ratio_tracks_misses() {
|
||||
let mut l = LossCalc::new();
|
||||
for _ in 0..10 {
|
||||
l.record(false);
|
||||
}
|
||||
assert_eq!(l.record(false), 0.0, "all delivered -> 0 loss");
|
||||
// 2 misses in a 13-sample window.
|
||||
l.record(true);
|
||||
let loss = l.record(true);
|
||||
assert!((loss - 2.0 / 13.0).abs() < 1e-6, "2/13 lost, got {loss}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jitter_is_low_for_steady_delivery_and_high_for_bursts() {
|
||||
let mut steady = Jitter::new();
|
||||
let mut t = 0u64;
|
||||
for _ in 0..20 {
|
||||
t += 33; // a clean ~30 fps cadence
|
||||
steady.frame(t);
|
||||
}
|
||||
assert!(steady.jitter() <= 1, "steady cadence -> ~0 jitter, got {}", steady.jitter());
|
||||
|
||||
let mut bursty = Jitter::new();
|
||||
let mut t = 0u64;
|
||||
for i in 0..20 {
|
||||
t += if i % 2 == 0 { 5 } else { 70 }; // frames bunch then gap
|
||||
bursty.frame(t);
|
||||
}
|
||||
assert!(bursty.jitter() > 20, "bursty delivery -> high jitter, got {}", bursty.jitter());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jitter_ignores_a_long_idle_gap() {
|
||||
let mut j = Jitter::new();
|
||||
j.frame(0);
|
||||
j.frame(33);
|
||||
j.frame(66);
|
||||
// Static screen: a 10 s gap is clamped, not counted as 10 s of jitter.
|
||||
j.frame(10_066);
|
||||
assert!(j.jitter() <= 2_000, "idle gap clamped");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,831 @@
|
||||
//! Relay-first P2P multiplexer with seamless fallback.
|
||||
//!
|
||||
//! [`upgrade_relay`] wraps any [`Transport`] (the WSS relay) and returns another
|
||||
//! `Transport` that the session loop ([`crate::signaling::connect_with`] /
|
||||
//! [`crate::host::run_host_with`]) uses completely unchanged. All the P2P logic
|
||||
//! runs in a background engine task that owns both the relay and, once
|
||||
//! negotiated, the WebRTC DataChannel:
|
||||
//!
|
||||
//! * The relay is LIVE from the first frame and is NEVER torn down while P2P is
|
||||
//! up (warm standby), so the first frame never waits on ICE and a P2P drop is
|
||||
//! invisible to the session.
|
||||
//! * The engine SNIFFS the session `sid` out of the handshake messages flowing
|
||||
//! through it (so neither loop has to hand it over), then drives the
|
||||
//! relay-brokered `rtc.*` handshake ([`crate::rtc::upgrade_offerer`] /
|
||||
//! [`crate::rtc::answer_offer`]).
|
||||
//! * UP edge: when the DataChannel opens AND the nominated ICE pair is DIRECT
|
||||
//! (non-TURN — [`crate::rtc::selected_pair_is_direct`]), media + the
|
||||
//! latency-critical input/keyframe control are routed over the channel; a
|
||||
//! TURN-relayed or unknown pair keeps media on the relay.
|
||||
//! * DOWN edge: if the peer connection fails/closes, media reroutes to the
|
||||
//! still-open relay with no `Ended` and no frame loss.
|
||||
//!
|
||||
//! `rtc.*` and `session.*` ALWAYS ride the relay; only `CHAN_VIDEO` / file
|
||||
//! binary and `i.mouse` / `i.key` / `keyframe` move to P2P when it is direct.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::Bytes;
|
||||
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
use webrtc::api::API;
|
||||
use webrtc::data_channel::RTCDataChannel;
|
||||
use webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState;
|
||||
use webrtc::peer_connection::RTCPeerConnection;
|
||||
use webrtc::stats::StatsReportType;
|
||||
|
||||
use crate::error::{NdError, Result};
|
||||
use crate::rtc;
|
||||
use crate::rtc::{RtcIceServerJson, WebDetached};
|
||||
use crate::transport::{install_reassembler, send_fragmented, Lane, Transport, TransportMsg};
|
||||
|
||||
/// A negotiated P2P connection: the native single reliable channel (16 KiB
|
||||
/// app-fragmented) or the web 3-channel topology (media/ctl/move, detached,
|
||||
/// whole frames) used to interop with the browser / Python host.
|
||||
enum P2pConn {
|
||||
Single(Arc<RTCDataChannel>),
|
||||
Web(Arc<WebDetached>),
|
||||
}
|
||||
|
||||
/// High-frequency, latency-critical control that rides the web `ctl` channel
|
||||
/// (the web viewer's send override routes exactly these; other control stays on
|
||||
/// the relay). Cheap prefix checks, no JSON parse on the hot path.
|
||||
fn is_fast_path_text(t: &str) -> bool {
|
||||
t.contains("\"t\":\"i.mouse\"")
|
||||
|| t.contains("\"t\":\"i.key\"")
|
||||
|| t.contains("\"t\":\"keyframe\"")
|
||||
}
|
||||
|
||||
/// Whether the direct-P2P fast path is enabled for this process. DEFAULT-ON:
|
||||
/// the native-peer capability gate keeps P2P to native<->native pairs and every
|
||||
/// other pairing stays on the relay, so default-on cannot break the existing
|
||||
/// web/Python peers. Opt OUT with `NETDESK_P2P=0` (also `false`/`off`/`no`).
|
||||
pub fn p2p_enabled() -> bool {
|
||||
!matches!(
|
||||
std::env::var("NETDESK_P2P").as_deref(),
|
||||
Ok("0") | Ok("false") | Ok("off") | Ok("no") | Ok("OFF") | Ok("FALSE") | Ok("NO")
|
||||
)
|
||||
}
|
||||
|
||||
/// Performance-based path health for the active DIRECT P2P pair.
|
||||
///
|
||||
/// The relay-first state machine already prefers a direct pair and falls back on
|
||||
/// hard failure. This adds the other half of the hybrid: abandon a direct pair
|
||||
/// that is measurably WORSE than a relay hop would be — not just one that
|
||||
/// dropped. It is deliberately conservative (only a *pathologically* high,
|
||||
/// *sustained* RTT trips it; a VPS relay hop is almost always below that), so it
|
||||
/// never fights the normal "direct is faster" case. Feed it the nominated
|
||||
/// candidate pair's RTT (ms) every ~[`HEALTH_INTERVAL`]; `on_rtt` returns true
|
||||
/// once the path has been bad long enough to downgrade.
|
||||
#[derive(Default)]
|
||||
pub struct P2pHealth {
|
||||
bad_streak: u32,
|
||||
}
|
||||
|
||||
/// How often the engine samples the active pair's RTT.
|
||||
pub const HEALTH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(2);
|
||||
/// A direct-pair RTT at/above this (ms) is "bad" — far worse than a relay hop.
|
||||
pub const P2P_RTT_BAD_MS: u32 = 500;
|
||||
/// Consecutive bad samples before downgrading (~10 s at [`HEALTH_INTERVAL`]).
|
||||
pub const P2P_RTT_CONFIRM: u32 = 5;
|
||||
|
||||
impl P2pHealth {
|
||||
/// Record one RTT sample (ms). `0` means "unknown" (not yet measured) and is
|
||||
/// ignored. Returns true when a downgrade to the relay is warranted.
|
||||
pub fn on_rtt(&mut self, rtt_ms: u32) -> bool {
|
||||
if rtt_ms == 0 {
|
||||
return false;
|
||||
}
|
||||
if rtt_ms >= P2P_RTT_BAD_MS {
|
||||
self.bad_streak += 1;
|
||||
} else {
|
||||
self.bad_streak = 0;
|
||||
}
|
||||
self.bad_streak >= P2P_RTT_CONFIRM
|
||||
}
|
||||
}
|
||||
|
||||
/// The nominated ICE candidate-pair RTT in ms, or 0 if unknown/not yet measured.
|
||||
async fn p2p_pair_rtt_ms(pc: &Arc<RTCPeerConnection>) -> u32 {
|
||||
let report = pc.get_stats().await;
|
||||
for v in report.reports.values() {
|
||||
if let StatsReportType::CandidatePair(cp) = v {
|
||||
if cp.nominated && cp.current_round_trip_time > 0.0 {
|
||||
return (cp.current_round_trip_time * 1000.0).round() as u32;
|
||||
}
|
||||
}
|
||||
}
|
||||
0
|
||||
}
|
||||
|
||||
/// Which side of the upgrade handshake this peer drives.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub enum P2pRole {
|
||||
/// Viewer — sends `rtc.config.req` + the offer.
|
||||
Offerer,
|
||||
/// Host — answers an inbound `rtc.offer`.
|
||||
Answerer,
|
||||
}
|
||||
|
||||
/// SCTP send-queue depth (bytes) past which a fresh MEDIA (video) frame is
|
||||
/// DROPPED rather than queued — fresh beats late. Only the drop-tolerant Media
|
||||
/// lane drops; Ctl (file bytes / control) and Move are always sent reliably.
|
||||
/// Matches the web host's 256 KiB `_send_bin` threshold.
|
||||
const P2P_MEDIA_CAP: usize = 262144;
|
||||
|
||||
/// A handle to force the session back onto the relay (the plan's "force-relay
|
||||
/// fast path", and the test hook for the DOWN edge). Cloneable; outlives the
|
||||
/// transport move into the session loop.
|
||||
#[derive(Clone)]
|
||||
pub struct P2pControl {
|
||||
down: mpsc::UnboundedSender<()>,
|
||||
}
|
||||
|
||||
impl P2pControl {
|
||||
/// Drop any active P2P path and route all media back over the warm relay.
|
||||
pub fn force_relay(&self) {
|
||||
let _ = self.down.send(());
|
||||
}
|
||||
}
|
||||
|
||||
/// The app-facing transport the session loop uses unchanged — just the mailbox
|
||||
/// into the engine task.
|
||||
pub struct P2pRelay {
|
||||
out: mpsc::UnboundedSender<(TransportMsg, Lane)>,
|
||||
inn: mpsc::UnboundedReceiver<Result<TransportMsg>>,
|
||||
down: mpsc::UnboundedSender<()>,
|
||||
_engine: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl P2pRelay {
|
||||
/// A control handle to force media back onto the relay.
|
||||
pub fn control(&self) -> P2pControl {
|
||||
P2pControl {
|
||||
down: self.down.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Transport for P2pRelay {
|
||||
async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
|
||||
// No lane hint -> treat as reliable Ctl (never dropped). Callers that
|
||||
// know it is video use send_binary_lane(.., Media).
|
||||
self.out
|
||||
.send((TransportMsg::Binary(data), Lane::Ctl))
|
||||
.map_err(|_| NdError::Connect("p2p engine stopped".into()))
|
||||
}
|
||||
async fn send_binary_lane(&mut self, data: Vec<u8>, lane: Lane) -> Result<()> {
|
||||
self.out
|
||||
.send((TransportMsg::Binary(data), lane))
|
||||
.map_err(|_| NdError::Connect("p2p engine stopped".into()))
|
||||
}
|
||||
async fn send_text(&mut self, text: String) -> Result<()> {
|
||||
self.out
|
||||
.send((TransportMsg::Text(text), Lane::Ctl))
|
||||
.map_err(|_| NdError::Connect("p2p engine stopped".into()))
|
||||
}
|
||||
async fn recv(&mut self) -> Option<Result<TransportMsg>> {
|
||||
self.inn.recv().await
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap `relay` with transparent P2P upgrade + fallback. `api` builds the peer
|
||||
/// connection(s). The returned transport is a drop-in for the relay.
|
||||
pub fn upgrade_relay(relay: Box<dyn Transport>, role: P2pRole, api: Arc<API>) -> P2pRelay {
|
||||
let (app_out_tx, app_out_rx) = mpsc::unbounded_channel::<(TransportMsg, Lane)>();
|
||||
let (app_in_tx, app_in_rx) = mpsc::unbounded_channel::<Result<TransportMsg>>();
|
||||
let (down_tx, down_rx) = mpsc::unbounded_channel::<()>();
|
||||
let engine = tokio::spawn(engine(
|
||||
relay,
|
||||
role,
|
||||
api,
|
||||
app_out_rx,
|
||||
app_in_tx,
|
||||
down_tx.clone(),
|
||||
down_rx,
|
||||
));
|
||||
P2pRelay {
|
||||
out: app_out_tx,
|
||||
inn: app_in_rx,
|
||||
down: down_tx,
|
||||
_engine: engine,
|
||||
}
|
||||
}
|
||||
|
||||
/// Pull `sid` out of a control-text message, if present and non-empty.
|
||||
fn sniff_sid_text(t: &str) -> Option<String> {
|
||||
serde_json::from_str::<Value>(t)
|
||||
.ok()
|
||||
.and_then(|v| v.get("sid").and_then(|x| x.as_str()).map(|s| s.to_string()))
|
||||
.filter(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
/// Whether this control text rides the P2P channel when it is up. Everything
|
||||
/// does EXCEPT `session.*` (session lifecycle) and `rtc.*` (upgrade signaling),
|
||||
/// which stay on the relay. Keeping ALL app traffic — input, keyframe requests
|
||||
/// AND file.* control — on the one ordered channel alongside the binary data is
|
||||
/// essential: splitting data (P2P) from its `file.done` (relay) lets the text
|
||||
/// race ahead of the slower-fragmented data and truncate the download. Cheap
|
||||
/// substring checks avoid a JSON parse on the hot path (serde emits `t` first).
|
||||
fn route_text_p2p(t: &str) -> bool {
|
||||
!(t.contains("\"t\":\"session.") || t.contains("\"t\":\"rtc."))
|
||||
}
|
||||
|
||||
type DoneTx = mpsc::UnboundedSender<(Arc<RTCPeerConnection>, P2pConn)>;
|
||||
|
||||
fn spawn_offerer(
|
||||
api: Arc<API>,
|
||||
sid: String,
|
||||
rtc_out: mpsc::UnboundedSender<String>,
|
||||
mut rtc_in: mpsc::UnboundedReceiver<Value>,
|
||||
done: DoneTx,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let send = {
|
||||
let tx = rtc_out.clone();
|
||||
move |s: String| {
|
||||
let _ = tx.send(s);
|
||||
}
|
||||
};
|
||||
match rtc::upgrade_offerer(&api, &sid, None, "netdesk", &send, &mut rtc_in).await {
|
||||
Ok((pc, dc)) => {
|
||||
let _ = done.send((pc, P2pConn::Single(dc)));
|
||||
}
|
||||
Err(e) => tracing::debug!("P2P offerer upgrade abandoned: {e}"),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// WEB-topology offerer (interop with the browser / Python host). Builds its own
|
||||
/// detach-enabled API, negotiates media/ctl/move, and reports a `P2pConn::Web`.
|
||||
fn spawn_offerer_web(
|
||||
sid: String,
|
||||
rtc_out: mpsc::UnboundedSender<String>,
|
||||
mut rtc_in: mpsc::UnboundedReceiver<Value>,
|
||||
done: DoneTx,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let api = match rtc::build_api_web() {
|
||||
Ok(a) => a,
|
||||
Err(e) => {
|
||||
tracing::debug!("build_api_web: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let send = {
|
||||
let tx = rtc_out.clone();
|
||||
move |s: String| {
|
||||
let _ = tx.send(s);
|
||||
}
|
||||
};
|
||||
match rtc::upgrade_offerer_web(&api, &sid, None, &send, &mut rtc_in).await {
|
||||
Ok((pc, web)) => {
|
||||
let _ = done.send((pc, P2pConn::Web(Arc::new(web))));
|
||||
}
|
||||
Err(e) => tracing::debug!("P2P web offerer abandoned: {e}"),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn spawn_answerer(
|
||||
api: Arc<API>,
|
||||
sid: String,
|
||||
offer_sdp: String,
|
||||
ice: Vec<RtcIceServerJson>,
|
||||
rtc_out: mpsc::UnboundedSender<String>,
|
||||
done: DoneTx,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let send = {
|
||||
let tx = rtc_out.clone();
|
||||
move |s: String| {
|
||||
let _ = tx.send(s);
|
||||
}
|
||||
};
|
||||
match rtc::answer_offer(&api, &sid, &offer_sdp, ice, &send).await {
|
||||
Ok((pc, dc)) => {
|
||||
let _ = done.send((pc, P2pConn::Single(dc)));
|
||||
}
|
||||
Err(e) => tracing::debug!("P2P answerer abandoned: {e}"),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// WEB-topology answerer (a browser / web viewer connects to this native host).
|
||||
/// Builds its own detach-enabled API and reports a `P2pConn::Web`.
|
||||
fn spawn_answerer_web(
|
||||
sid: String,
|
||||
offer_sdp: String,
|
||||
ice: Vec<RtcIceServerJson>,
|
||||
rtc_out: mpsc::UnboundedSender<String>,
|
||||
done: DoneTx,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let api = match rtc::build_api_web() {
|
||||
Ok(a) => a,
|
||||
Err(e) => {
|
||||
tracing::debug!("build_api_web: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let send = {
|
||||
let tx = rtc_out.clone();
|
||||
move |s: String| {
|
||||
let _ = tx.send(s);
|
||||
}
|
||||
};
|
||||
match rtc::answer_offer_web(&api, &sid, &offer_sdp, ice, &send).await {
|
||||
Ok((pc, web)) => {
|
||||
let _ = done.send((pc, P2pConn::Web(Arc::new(web))));
|
||||
}
|
||||
Err(e) => tracing::debug!("P2P web answerer abandoned: {e}"),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Start one offerer upgrade attempt (the initial one, or a re-upgrade after a
|
||||
/// drop). Creates a fresh rtc_in channel, spawns the Single or Web offerer, and
|
||||
/// returns the sender the relay loop forwards inbound rtc.* to.
|
||||
fn start_offerer_attempt(
|
||||
web: bool,
|
||||
api: Arc<API>,
|
||||
sid: String,
|
||||
rtc_out: mpsc::UnboundedSender<String>,
|
||||
done: DoneTx,
|
||||
) -> mpsc::UnboundedSender<Value> {
|
||||
let (tx, rx) = mpsc::unbounded_channel::<Value>();
|
||||
if web {
|
||||
spawn_offerer_web(sid, rtc_out, rx, done);
|
||||
} else {
|
||||
spawn_offerer(api, sid, rtc_out, rx, done);
|
||||
}
|
||||
tx
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn engine(
|
||||
mut relay: Box<dyn Transport>,
|
||||
role: P2pRole,
|
||||
api: Arc<API>,
|
||||
mut app_out_rx: mpsc::UnboundedReceiver<(TransportMsg, Lane)>,
|
||||
app_in_tx: mpsc::UnboundedSender<Result<TransportMsg>>,
|
||||
down_tx: mpsc::UnboundedSender<()>,
|
||||
mut down_rx: mpsc::UnboundedReceiver<()>,
|
||||
) {
|
||||
let mut sid: Option<String> = None;
|
||||
let mut media_on_p2p = false;
|
||||
let mut p2p_send: Option<P2pConn> = None;
|
||||
let mut _pc: Option<Arc<RTCPeerConnection>> = None;
|
||||
let mut next_frame_id: u32 = 0;
|
||||
let mut offerer_started = false;
|
||||
let mut answerer_started = false;
|
||||
// Performance health of the active direct pair, and a latch: once a direct
|
||||
// pair is abandoned for being too slow, stay on the relay for the rest of
|
||||
// the session (don't re-upgrade into the same bad path and flap).
|
||||
let mut health = P2pHealth::default();
|
||||
let mut health_down = false;
|
||||
|
||||
// Inbound P2P frames (reassembled) land here; held open for the engine's
|
||||
// lifetime so re-use across the session is clean. Fed by install_reassembler
|
||||
// once a channel is established.
|
||||
let (p2p_in_tx, mut p2p_in_rx) = mpsc::unbounded_channel::<TransportMsg>();
|
||||
// rtc.* the upgrade task wants put on the relay.
|
||||
let (rtc_out_tx, mut rtc_out_rx) = mpsc::unbounded_channel::<String>();
|
||||
// Inbound rtc.* forwarded to the CURRENT offerer upgrade attempt. Re-created
|
||||
// per attempt so a re-upgrade after a drop gets a fresh channel.
|
||||
let mut rtc_in_tx: Option<mpsc::UnboundedSender<Value>> = None;
|
||||
// Negotiated (pc, connection).
|
||||
let (done_tx, mut done_rx) =
|
||||
mpsc::unbounded_channel::<(Arc<RTCPeerConnection>, P2pConn)>();
|
||||
// Re-upgrade (hysteresis): after a P2P drop, retry the upgrade once after a
|
||||
// backoff so a transient drop recovers to P2P instead of staying on the
|
||||
// relay. Default 15s; shorter in tests via NETDESK_P2P_REUP_MS.
|
||||
let reup_ms: u64 = std::env::var("NETDESK_P2P_REUP_MS")
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(15000);
|
||||
let (reup_tx, mut reup_rx) = mpsc::unbounded_channel::<()>();
|
||||
// Offerer topology once decided: Some(false)=native single, Some(true)=web,
|
||||
// None=relay-only. Remembered so a re-upgrade uses the same topology.
|
||||
let mut offerer_mode: Option<bool> = None;
|
||||
// Periodic RTT health probe of the active direct pair.
|
||||
let mut health_tick = tokio::time::interval(HEALTH_INTERVAL);
|
||||
health_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
|
||||
enum Ev {
|
||||
AppOut(Option<(TransportMsg, Lane)>),
|
||||
RelayIn(Option<Result<TransportMsg>>),
|
||||
P2pIn(Option<TransportMsg>),
|
||||
RtcOut(Option<String>),
|
||||
Done(Option<(Arc<RTCPeerConnection>, P2pConn)>),
|
||||
ReUpgrade,
|
||||
Down,
|
||||
Health,
|
||||
}
|
||||
|
||||
loop {
|
||||
let ev = tokio::select! {
|
||||
m = app_out_rx.recv() => Ev::AppOut(m),
|
||||
m = relay.recv() => Ev::RelayIn(m),
|
||||
m = p2p_in_rx.recv() => Ev::P2pIn(m),
|
||||
s = rtc_out_rx.recv() => Ev::RtcOut(s),
|
||||
d = done_rx.recv() => Ev::Done(d),
|
||||
_ = reup_rx.recv() => Ev::ReUpgrade,
|
||||
_ = down_rx.recv() => Ev::Down,
|
||||
_ = health_tick.tick() => Ev::Health,
|
||||
};
|
||||
|
||||
match ev {
|
||||
// ---- outbound from the session loop ----
|
||||
Ev::AppOut(None) => break, // the session dropped the transport
|
||||
Ev::AppOut(Some((msg, lane))) => {
|
||||
if sid.is_none() {
|
||||
if let TransportMsg::Text(t) = &msg {
|
||||
if let Some(s) = sniff_sid_text(t) {
|
||||
sid = Some(s);
|
||||
}
|
||||
}
|
||||
}
|
||||
match msg {
|
||||
TransportMsg::Binary(b) => match (media_on_p2p, p2p_send.as_ref()) {
|
||||
// native single channel: 16 KiB app-fragments; Media is
|
||||
// drop-tolerant past the queue cap, Ctl/Move reliable.
|
||||
(true, Some(P2pConn::Single(dc))) => {
|
||||
let dc = dc.clone();
|
||||
if lane == Lane::Media && dc.buffered_amount().await > P2P_MEDIA_CAP {
|
||||
// drop (fresh beats late)
|
||||
} else {
|
||||
let fid = next_frame_id;
|
||||
next_frame_id = next_frame_id.wrapping_add(1);
|
||||
if send_fragmented(&dc, fid, &b).await.is_err() {
|
||||
let _ = relay.send_binary_lane(b, lane).await;
|
||||
let _ = down_tx.send(());
|
||||
}
|
||||
}
|
||||
}
|
||||
// web topology: video (Media) rides the media channel whole;
|
||||
// everything else (e.g. file upload) stays on the relay, like
|
||||
// the web viewer. Media is drop-tolerant past the cap.
|
||||
(true, Some(P2pConn::Web(web))) => {
|
||||
if lane == Lane::Media {
|
||||
if web.media.buffered_amount() <= P2P_MEDIA_CAP
|
||||
&& web
|
||||
.media
|
||||
.write_data_channel(&Bytes::from(b.clone()), false)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
let _ = relay.send_binary_lane(b, lane).await;
|
||||
let _ = down_tx.send(());
|
||||
}
|
||||
} else {
|
||||
let _ = relay.send_binary_lane(b, lane).await;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
let _ = relay.send_binary_lane(b, lane).await;
|
||||
}
|
||||
},
|
||||
TransportMsg::Text(t) => match (media_on_p2p, p2p_send.as_ref()) {
|
||||
(true, Some(P2pConn::Single(dc))) if route_text_p2p(&t) => {
|
||||
let dc = dc.clone();
|
||||
if dc.send_text(t.clone()).await.is_err() {
|
||||
let _ = relay.send_text(t).await;
|
||||
let _ = down_tx.send(());
|
||||
}
|
||||
}
|
||||
// web: only input / keyframe rides ctl; other control stays
|
||||
// on the relay (the web viewer's send policy).
|
||||
(true, Some(P2pConn::Web(web))) if is_fast_path_text(&t) => {
|
||||
if web
|
||||
.ctl
|
||||
.write_data_channel(&Bytes::from(t.clone().into_bytes()), true)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
let _ = relay.send_text(t).await;
|
||||
let _ = down_tx.send(());
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
let _ = relay.send_text(t).await;
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// ---- inbound from the relay ----
|
||||
Ev::RelayIn(None) => break, // relay closed -> session ends (app recv None)
|
||||
Ev::RelayIn(Some(Err(e))) => {
|
||||
let _ = app_in_tx.send(Err(e));
|
||||
break;
|
||||
}
|
||||
Ev::RelayIn(Some(Ok(m))) => {
|
||||
let mut forward = true;
|
||||
if let TransportMsg::Text(t) = &m {
|
||||
if let Ok(v) = serde_json::from_str::<Value>(t) {
|
||||
let tt = v.get("t").and_then(|x| x.as_str()).unwrap_or("").to_string();
|
||||
if sid.is_none() {
|
||||
if let Some(s) = v.get("sid").and_then(|x| x.as_str()) {
|
||||
if !s.is_empty() {
|
||||
sid = Some(s.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
if tt.starts_with("rtc.") {
|
||||
forward = false; // rtc.* is ours, never the app's
|
||||
match role {
|
||||
P2pRole::Offerer => {
|
||||
if let Some(tx) = &rtc_in_tx {
|
||||
let _ = tx.send(v);
|
||||
}
|
||||
}
|
||||
P2pRole::Answerer => {
|
||||
if tt == "rtc.offer" && !answerer_started {
|
||||
answerer_started = true;
|
||||
let sid_s = sid.clone().unwrap_or_default();
|
||||
let is_native =
|
||||
v.get("native").and_then(|x| x.as_bool()) == Some(true);
|
||||
let sdp = v
|
||||
.get("sdp")
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let ice: Vec<RtcIceServerJson> = v
|
||||
.get("ice")
|
||||
.and_then(|x| serde_json::from_value(x.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
if is_native {
|
||||
// native single-channel offerer.
|
||||
spawn_answerer(
|
||||
api.clone(),
|
||||
sid_s,
|
||||
sdp,
|
||||
ice,
|
||||
rtc_out_tx.clone(),
|
||||
done_tx.clone(),
|
||||
);
|
||||
} else if std::env::var("NETDESK_P2P_WEB").is_ok() {
|
||||
// web/3-channel offerer (a browser) — answer in
|
||||
// web topology when opted in.
|
||||
spawn_answerer_web(
|
||||
sid_s,
|
||||
sdp,
|
||||
ice,
|
||||
rtc_out_tx.clone(),
|
||||
done_tx.clone(),
|
||||
);
|
||||
} else {
|
||||
// not opted into web interop: decline so the
|
||||
// pairing stays on the relay.
|
||||
let _ = rtc_out_tx.send(
|
||||
format!(r#"{{"t":"rtc.unavailable","sid":"{sid_s}"}}"#),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == P2pRole::Offerer
|
||||
&& !offerer_started
|
||||
&& tt == "session.accept"
|
||||
{
|
||||
// A NATIVE host (info.p2p == "native") uses the native
|
||||
// single-channel upgrade. A non-native host (browser /
|
||||
// Python) uses the WEB 3-channel topology, but only when
|
||||
// NETDESK_P2P_WEB is opted in — so the default stays the
|
||||
// safe native-only behavior until the live cross-matrix
|
||||
// is confirmed.
|
||||
let native_host = v
|
||||
.get("info")
|
||||
.and_then(|i| i.get("p2p"))
|
||||
.and_then(|p| p.as_str())
|
||||
== Some("native");
|
||||
let web_ok = std::env::var("NETDESK_P2P_WEB").is_ok();
|
||||
if native_host {
|
||||
offerer_started = true;
|
||||
offerer_mode = Some(false);
|
||||
rtc_in_tx = Some(start_offerer_attempt(
|
||||
false,
|
||||
api.clone(),
|
||||
sid.clone().unwrap_or_default(),
|
||||
rtc_out_tx.clone(),
|
||||
done_tx.clone(),
|
||||
));
|
||||
} else if web_ok {
|
||||
offerer_started = true;
|
||||
offerer_mode = Some(true);
|
||||
rtc_in_tx = Some(start_offerer_attempt(
|
||||
true,
|
||||
api.clone(),
|
||||
sid.clone().unwrap_or_default(),
|
||||
rtc_out_tx.clone(),
|
||||
done_tx.clone(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if forward {
|
||||
let _ = app_in_tx.send(Ok(m));
|
||||
}
|
||||
}
|
||||
|
||||
// ---- inbound from the P2P channel ----
|
||||
Ev::P2pIn(Some(m)) => {
|
||||
let _ = app_in_tx.send(Ok(m));
|
||||
}
|
||||
Ev::P2pIn(None) => {} // we hold a sender; this does not fire in practice
|
||||
|
||||
// ---- rtc.* the upgrade task wants on the relay ----
|
||||
Ev::RtcOut(Some(s)) => {
|
||||
let _ = relay.send_text(s).await;
|
||||
}
|
||||
Ev::RtcOut(None) => {}
|
||||
|
||||
// ---- UP edge: a channel was negotiated ----
|
||||
Ev::Done(Some((new_pc, conn))) => {
|
||||
let direct = rtc::selected_pair_is_direct(&new_pc).await;
|
||||
match &conn {
|
||||
P2pConn::Single(dc) => {
|
||||
// native: 16 KiB app-fragments reassembled by install_reassembler.
|
||||
install_reassembler(dc, p2p_in_tx.clone());
|
||||
}
|
||||
P2pConn::Web(web) => {
|
||||
// web: whole frames — spawn a detached read-loop per inbound
|
||||
// channel (media = video, ctl = control). Each read is one
|
||||
// complete message (is_string -> Text, else Binary).
|
||||
for raw in [web.media.clone(), web.ctl.clone()] {
|
||||
let tx = p2p_in_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut buf = vec![0u8; 262144];
|
||||
loop {
|
||||
match raw.read_data_channel(&mut buf).await {
|
||||
Ok((0, _)) => break,
|
||||
Ok((n, is_string)) => {
|
||||
let msg = if is_string {
|
||||
match String::from_utf8(buf[..n].to_vec()) {
|
||||
Ok(s) => TransportMsg::Text(s),
|
||||
Err(_) => continue,
|
||||
}
|
||||
} else {
|
||||
TransportMsg::Binary(buf[..n].to_vec())
|
||||
};
|
||||
let _ = tx.send(msg);
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
let dtx = down_tx.clone();
|
||||
new_pc.on_peer_connection_state_change(Box::new(move |s: RTCPeerConnectionState| {
|
||||
let dtx = dtx.clone();
|
||||
Box::pin(async move {
|
||||
if matches!(
|
||||
s,
|
||||
RTCPeerConnectionState::Failed
|
||||
| RTCPeerConnectionState::Disconnected
|
||||
| RTCPeerConnectionState::Closed
|
||||
) {
|
||||
let _ = dtx.send(());
|
||||
}
|
||||
})
|
||||
}));
|
||||
p2p_send = Some(conn);
|
||||
_pc = Some(new_pc);
|
||||
// Only move media off the relay for a DIRECT pair; a TURN-relayed
|
||||
// hop is no better than our own relay.
|
||||
media_on_p2p = direct;
|
||||
tracing::info!(direct, "P2P DataChannel open");
|
||||
if role == P2pRole::Offerer {
|
||||
let _ = app_in_tx.send(Ok(TransportMsg::Text(format!(
|
||||
r#"{{"t":"p2p.state","active":{media_on_p2p},"direct":{direct}}}"#
|
||||
))));
|
||||
}
|
||||
}
|
||||
Ev::Done(None) => {}
|
||||
|
||||
// ---- periodic health probe: abandon a pathologically slow direct
|
||||
// pair for the relay (the "pick the faster path" half of hybrid). ----
|
||||
Ev::Health => {
|
||||
if media_on_p2p && !health_down {
|
||||
if let Some(pc) = &_pc {
|
||||
let rtt = p2p_pair_rtt_ms(pc).await;
|
||||
if health.on_rtt(rtt) {
|
||||
tracing::warn!(
|
||||
rtt_ms = rtt,
|
||||
"P2P direct pair too slow — downgrading to the relay for the session"
|
||||
);
|
||||
health_down = true;
|
||||
let _ = down_tx.send(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- DOWN edge: pc failed/closed -> fall back to the warm relay ----
|
||||
Ev::Down => {
|
||||
let was_up = media_on_p2p || p2p_send.is_some();
|
||||
if was_up {
|
||||
tracing::info!("P2P dropped — continuing over the relay");
|
||||
}
|
||||
media_on_p2p = false;
|
||||
p2p_send = None;
|
||||
_pc = None;
|
||||
health = P2pHealth::default(); // reset the RTT streak for the next pair
|
||||
if was_up && role == P2pRole::Offerer {
|
||||
let _ = app_in_tx.send(Ok(TransportMsg::Text(
|
||||
r#"{"t":"p2p.state","active":false,"direct":false}"#.to_string(),
|
||||
)));
|
||||
}
|
||||
// Re-upgrade hysteresis: allow a fresh attempt. The answerer
|
||||
// just re-answers the viewer's next rtc.offer; the offerer
|
||||
// schedules one retry after a backoff.
|
||||
if was_up {
|
||||
answerer_started = false;
|
||||
offerer_started = false;
|
||||
rtc_in_tx = None;
|
||||
// Suppress re-upgrade after a health downgrade: the direct
|
||||
// path proved too slow, so stay on the relay (no flapping).
|
||||
if role == P2pRole::Offerer && offerer_mode.is_some() && !health_down {
|
||||
let reup_tx = reup_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(reup_ms)).await;
|
||||
let _ = reup_tx.send(());
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- re-upgrade timer fired: retry the offerer upgrade once ----
|
||||
Ev::ReUpgrade => {
|
||||
if role == P2pRole::Offerer
|
||||
&& !offerer_started
|
||||
&& p2p_send.is_none()
|
||||
&& sid.is_some()
|
||||
{
|
||||
if let Some(web) = offerer_mode {
|
||||
offerer_started = true;
|
||||
rtc_in_tx = Some(start_offerer_attempt(
|
||||
web,
|
||||
api.clone(),
|
||||
sid.clone().unwrap_or_default(),
|
||||
rtc_out_tx.clone(),
|
||||
done_tx.clone(),
|
||||
));
|
||||
tracing::info!("P2P re-upgrade attempt");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod health_tests {
|
||||
use super::{P2pHealth, P2P_RTT_CONFIRM};
|
||||
|
||||
#[test]
|
||||
fn good_rtt_never_downgrades() {
|
||||
let mut h = P2pHealth::default();
|
||||
for _ in 0..50 {
|
||||
assert!(!h.on_rtt(30), "a healthy direct pair must stay on P2P");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_rtt_is_ignored() {
|
||||
let mut h = P2pHealth::default();
|
||||
for _ in 0..50 {
|
||||
assert!(!h.on_rtt(0), "0 = not-yet-measured, never a downgrade");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sustained_bad_rtt_downgrades_after_confirm() {
|
||||
let mut h = P2pHealth::default();
|
||||
// Must take the full confirm streak — not a single spike.
|
||||
for i in 1..P2P_RTT_CONFIRM {
|
||||
assert!(!h.on_rtt(800), "should not trip on sample {i} (< confirm)");
|
||||
}
|
||||
assert!(h.on_rtt(800), "should downgrade once the streak reaches confirm");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_good_sample_resets_the_streak() {
|
||||
let mut h = P2pHealth::default();
|
||||
for _ in 0..(P2P_RTT_CONFIRM - 1) {
|
||||
assert!(!h.on_rtt(900));
|
||||
}
|
||||
assert!(!h.on_rtt(40), "a healthy sample clears the bad streak");
|
||||
// Now it must take another full streak to trip.
|
||||
for i in 1..P2P_RTT_CONFIRM {
|
||||
assert!(!h.on_rtt(900), "streak restarted; sample {i} must not trip");
|
||||
}
|
||||
assert!(h.on_rtt(900));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
//! Wire protocol for the NetDesk signaling server (`/ws`). Mirrors the messages
|
||||
//! the web viewer exchanges, so the native client is a drop-in peer:
|
||||
//! viewer -> session.begin{id}
|
||||
//! host -> session.challenge{sid,salt,nonce,iters,e2e,pub,totp}
|
||||
//! viewer -> session.auth{proof,pub,totp}
|
||||
//! host -> session.accept{sid,info{host_name,e2e,w,h,...}} (then binary frames)
|
||||
//! host -> session.deny{reason}
|
||||
//! Binary frames (after optional AES-GCM): [0]=channel, [1]=flags, [2..]=codec
|
||||
//! bitstream. For CHAN_VIDEO, `flags` packs bit0=keyframe and bits1-3=codec id
|
||||
//! (see [`VideoCodec`] / [`video_flags`]): H.264 Annex-B (0) or VP9 (1).
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub const CHAN_VIDEO: u8 = 0x01;
|
||||
pub const CHAN_AUDIO: u8 = 0xa1;
|
||||
/// Viewer->host microphone for a two-way voice call: `[CHAN_VOICE][opus]`,
|
||||
/// E2E-sealed like the other channels. The host decodes it and plays it on its
|
||||
/// default render endpoint.
|
||||
pub const CHAN_VOICE: u8 = 0xa2;
|
||||
pub const CHAN_FILE: u8 = 0xf1;
|
||||
pub const CHAN_FILE_END: u8 = 0xf2;
|
||||
|
||||
/// Video codec carried in the high bits of a CHAN_VIDEO frame's flags byte.
|
||||
///
|
||||
/// The host tags every frame; the viewer builds the matching decoder and can
|
||||
/// switch mid-session at any keyframe (the adaptive path moves hardware H.264
|
||||
/// <-> software VP9 by link quality — VP9 is far better per-bitrate on slow
|
||||
/// links, H.264 is near-free on fast ones). Values are forward-compatible:
|
||||
/// bits1-3 leave room for VP8/AV1 later.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
#[repr(u8)]
|
||||
pub enum VideoCodec {
|
||||
H264 = 0,
|
||||
Vp9 = 1,
|
||||
}
|
||||
|
||||
impl VideoCodec {
|
||||
/// Decode the codec id from a CHAN_VIDEO flags byte (unknown -> H264).
|
||||
#[inline]
|
||||
pub fn from_bits(flags: u8) -> Self {
|
||||
match (flags >> 1) & 0x7 {
|
||||
1 => VideoCodec::Vp9,
|
||||
_ => VideoCodec::H264,
|
||||
}
|
||||
}
|
||||
#[inline]
|
||||
pub fn as_u8(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
}
|
||||
|
||||
/// Pack a CHAN_VIDEO flags byte: bit0 = keyframe, bits1-3 = codec id.
|
||||
#[inline]
|
||||
pub fn video_flags(keyframe: bool, codec: VideoCodec) -> u8 {
|
||||
(keyframe as u8) | ((codec as u8) << 1)
|
||||
}
|
||||
|
||||
/// Outbound: begin a session against a host id.
|
||||
#[derive(Serialize)]
|
||||
pub struct SessionBegin<'a> {
|
||||
pub t: &'a str, // "session.begin"
|
||||
pub id: &'a str,
|
||||
pub direct: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub takeover: Option<bool>,
|
||||
/// Viewer can decode VP9 — the host only sends VP9 frames when this is set
|
||||
/// (a non-VP9 viewer, e.g. a non-Windows native build or the web client,
|
||||
/// stays on H.264). Omitted when false so older hosts ignore it.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub vp9: Option<bool>,
|
||||
}
|
||||
|
||||
/// Inbound: the host's auth challenge (relayed through the server).
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct Challenge {
|
||||
pub sid: String,
|
||||
pub salt: String, // base64
|
||||
pub nonce: String, // opaque string, concatenated before hashing
|
||||
pub iters: u32,
|
||||
#[serde(default)]
|
||||
pub e2e: bool,
|
||||
#[serde(default, rename = "pub")]
|
||||
pub pub_: Option<String>, // host ECDH P-256 public key, base64 raw (65 bytes)
|
||||
#[serde(default)]
|
||||
pub totp: bool,
|
||||
}
|
||||
|
||||
/// Outbound: proof of the password + our ECDH public key.
|
||||
#[derive(Serialize)]
|
||||
pub struct SessionAuth<'a> {
|
||||
pub t: &'a str, // "session.auth"
|
||||
pub proof: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none", rename = "pub")]
|
||||
pub pub_: Option<String>,
|
||||
/// Channel-bound proof (MITM hardening): `hex(SHA256(nonce||verifier||host_pub))`.
|
||||
/// Sent alongside `proof` whenever E2E is active, so the host can detect a
|
||||
/// relay that swapped its ephemeral key. Omitted on non-E2E sessions; older
|
||||
/// hosts ignore the extra field.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub proof_bind: Option<String>,
|
||||
pub totp: String,
|
||||
pub sid: String,
|
||||
/// Viewer can decode VP9. Carried in `session.auth` (not `session.begin`,
|
||||
/// whose custom fields the signaling server strips — it does relay auth,
|
||||
/// since the host reads our `pub` here for E2E). The host only sends VP9
|
||||
/// when this is set. Omitted when false so older hosts ignore it.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub vp9: Option<bool>,
|
||||
}
|
||||
|
||||
/// Inbound: accepted. `info` carries host name, dims, e2e flag, etc.
|
||||
#[derive(Deserialize, Debug, Clone, Default)]
|
||||
pub struct AcceptInfo {
|
||||
#[serde(default)]
|
||||
pub host_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub e2e: bool,
|
||||
#[serde(default)]
|
||||
pub w: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub h: Option<u32>,
|
||||
}
|
||||
|
||||
/// A keyboard/pointer input event sent viewer -> host.
|
||||
#[derive(Serialize)]
|
||||
pub struct MouseInput<'a> {
|
||||
pub t: &'a str, // "i.mouse"
|
||||
pub x: i32,
|
||||
pub y: i32,
|
||||
pub b: u32,
|
||||
pub w: f32,
|
||||
pub hw: f32,
|
||||
pub sid: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct KeyInput<'a> {
|
||||
pub t: &'a str, // "i.key"
|
||||
// The host (netdesk_host.py) and the web viewer both use the short wire
|
||||
// names `c` (DOM code) and `d` (down). The native viewer MUST match them or
|
||||
// the host reads an empty code with down=false and injects nothing — i.e.
|
||||
// the keyboard is silently dead ("connects but it's just a view").
|
||||
#[serde(rename = "c")]
|
||||
pub code: String,
|
||||
#[serde(rename = "d")]
|
||||
pub down: bool,
|
||||
pub sid: String,
|
||||
}
|
||||
|
||||
/// Split an H.264 Annex-B buffer into NAL units (strips 00 00 01 / 00 00 00 01 start codes).
|
||||
pub fn split_nals(buf: &[u8]) -> Vec<&[u8]> {
|
||||
let mut out = Vec::new();
|
||||
let mut i = 0usize;
|
||||
let n = buf.len();
|
||||
let mut start: Option<usize> = None;
|
||||
while i + 3 <= n {
|
||||
let sc3 = buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1;
|
||||
let sc4 = i + 4 <= n && buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 0 && buf[i + 3] == 1;
|
||||
if sc3 || sc4 {
|
||||
if let Some(s) = start {
|
||||
out.push(&buf[s..i]);
|
||||
}
|
||||
i += if sc4 { 4 } else { 3 };
|
||||
start = Some(i);
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
if let Some(s) = start {
|
||||
if s < n {
|
||||
out.push(&buf[s..n]);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// NAL unit type (lower 5 bits of the first byte). 7=SPS, 8=PPS, 5=IDR, 1=non-IDR.
|
||||
pub fn nal_type(nal: &[u8]) -> u8 {
|
||||
if nal.is_empty() {
|
||||
0
|
||||
} else {
|
||||
nal[0] & 0x1f
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
//! Session recording — repackage the received H.264 Annex-B access units into a
|
||||
//! playable `.mp4` with **no re-encode** (just remuxing). The viewer tees the
|
||||
//! same access units it decodes for display into this recorder; the host never
|
||||
//! has to do anything. Variable-framerate is preserved by timing each sample
|
||||
//! from its arrival timestamp (the host only sends a frame when the screen
|
||||
//! changes, so fixed-rate muxing would speed up idle stretches).
|
||||
//!
|
||||
//! Developed by Akronlogic.
|
||||
use crate::error::{NdError, Result};
|
||||
use crate::proto::{nal_type, split_nals};
|
||||
use std::fs::File;
|
||||
use std::io::{BufReader, BufWriter};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn other(msg: impl std::fmt::Display) -> NdError {
|
||||
NdError::Other(anyhow::anyhow!("recorder: {msg}"))
|
||||
}
|
||||
|
||||
/// Muxes H.264 access units into an MP4 file on disk.
|
||||
pub struct Mp4Recorder {
|
||||
path: PathBuf,
|
||||
width: u16,
|
||||
height: u16,
|
||||
writer: Option<mp4::Mp4Writer<BufWriter<File>>>,
|
||||
sps: Option<Vec<u8>>,
|
||||
pps: Option<Vec<u8>>,
|
||||
start_ms: Option<u64>,
|
||||
/// One buffered sample so its duration can be set from the NEXT sample's
|
||||
/// timestamp (MP4 samples carry how long they are shown).
|
||||
pending: Option<(Vec<u8>, bool, u64)>, // (avcc bytes, is_sync, rel_ms)
|
||||
samples: u64,
|
||||
}
|
||||
|
||||
const TIMESCALE: u32 = 1000; // milliseconds
|
||||
|
||||
impl Mp4Recorder {
|
||||
/// Create a recorder that will write to `path`. `width`/`height` are the
|
||||
/// encoded frame size (from `session.info`); the file is only created once
|
||||
/// the first keyframe (SPS+PPS) arrives.
|
||||
pub fn new(path: impl AsRef<Path>, width: u16, height: u16) -> Self {
|
||||
Self {
|
||||
path: path.as_ref().to_path_buf(),
|
||||
width: width.max(2),
|
||||
height: height.max(2),
|
||||
writer: None,
|
||||
sps: None,
|
||||
pps: None,
|
||||
start_ms: None,
|
||||
pending: None,
|
||||
samples: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed one access unit (Annex-B). `ts_ms` is a monotonically increasing
|
||||
/// arrival time in milliseconds. SPS/PPS are harvested automatically and the
|
||||
/// file starts on the first keyframe.
|
||||
pub fn push(&mut self, au: &[u8], keyframe: bool, ts_ms: u64) -> Result<()> {
|
||||
let nals = split_nals(au);
|
||||
|
||||
// Harvest parameter sets (needed for the avcC box).
|
||||
for nal in &nals {
|
||||
match nal_type(nal) {
|
||||
7 if self.sps.is_none() => self.sps = Some(nal.to_vec()),
|
||||
8 if self.pps.is_none() => self.pps = Some(nal.to_vec()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Lazily open the file once we can describe the track.
|
||||
if self.writer.is_none() {
|
||||
let (sps, pps) = match (self.sps.clone(), self.pps.clone()) {
|
||||
(Some(s), Some(p)) => (s, p),
|
||||
_ => return Ok(()), // wait for the first keyframe's SPS/PPS
|
||||
};
|
||||
let file = BufWriter::new(File::create(&self.path).map_err(other)?);
|
||||
let config = mp4::Mp4Config {
|
||||
major_brand: mp4::FourCC::from(*b"isom"),
|
||||
minor_version: 512,
|
||||
compatible_brands: vec![
|
||||
mp4::FourCC::from(*b"isom"),
|
||||
mp4::FourCC::from(*b"iso2"),
|
||||
mp4::FourCC::from(*b"avc1"),
|
||||
mp4::FourCC::from(*b"mp41"),
|
||||
],
|
||||
timescale: TIMESCALE,
|
||||
};
|
||||
let mut w = mp4::Mp4Writer::write_start(file, &config).map_err(other)?;
|
||||
let track = mp4::TrackConfig {
|
||||
track_type: mp4::TrackType::Video,
|
||||
timescale: TIMESCALE,
|
||||
language: "und".to_string(),
|
||||
media_conf: mp4::MediaConfig::AvcConfig(mp4::AvcConfig {
|
||||
width: self.width,
|
||||
height: self.height,
|
||||
seq_param_set: sps,
|
||||
pic_param_set: pps,
|
||||
}),
|
||||
};
|
||||
w.add_track(&track).map_err(other)?;
|
||||
self.writer = Some(w);
|
||||
}
|
||||
|
||||
// Convert the Annex-B access unit to AVCC (4-byte length-prefixed NALs),
|
||||
// dropping parameter sets (already in avcC) and access-unit delimiters.
|
||||
let mut avcc = Vec::with_capacity(au.len());
|
||||
for nal in &nals {
|
||||
if nal.is_empty() {
|
||||
continue;
|
||||
}
|
||||
match nal_type(nal) {
|
||||
7 | 8 | 9 => continue, // SPS, PPS, AUD
|
||||
_ => {}
|
||||
}
|
||||
avcc.extend_from_slice(&(nal.len() as u32).to_be_bytes());
|
||||
avcc.extend_from_slice(nal);
|
||||
}
|
||||
if avcc.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let rel_ms = {
|
||||
let start = *self.start_ms.get_or_insert(ts_ms);
|
||||
ts_ms.saturating_sub(start)
|
||||
};
|
||||
|
||||
// Flush the previously buffered sample now that we know its duration.
|
||||
if let Some((bytes, is_sync, prev_rel)) = self.pending.take() {
|
||||
let dur = rel_ms.saturating_sub(prev_rel).clamp(1, 10_000) as u32;
|
||||
self.write_sample(bytes, is_sync, prev_rel, dur)?;
|
||||
}
|
||||
self.pending = Some((avcc, keyframe, rel_ms));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_sample(&mut self, bytes: Vec<u8>, is_sync: bool, start: u64, dur: u32) -> Result<()> {
|
||||
if let Some(w) = self.writer.as_mut() {
|
||||
let sample = mp4::Mp4Sample {
|
||||
start_time: start,
|
||||
duration: dur,
|
||||
rendering_offset: 0,
|
||||
is_sync,
|
||||
bytes: mp4::Bytes::from(bytes),
|
||||
};
|
||||
w.write_sample(1, &sample).map_err(other)?;
|
||||
self.samples += 1;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Flush the last sample and finalize the MP4. Returns (path, sample count).
|
||||
pub fn finish(mut self) -> Result<(PathBuf, u64)> {
|
||||
if let Some((bytes, is_sync, rel)) = self.pending.take() {
|
||||
self.write_sample(bytes, is_sync, rel, 33)?; // nominal last-frame hold
|
||||
}
|
||||
if let Some(mut w) = self.writer.take() {
|
||||
w.write_end().map_err(other)?;
|
||||
}
|
||||
Ok((self.path.clone(), self.samples))
|
||||
}
|
||||
|
||||
pub fn samples(&self) -> u64 {
|
||||
self.samples
|
||||
}
|
||||
}
|
||||
|
||||
/// Read an MP4 back and return (video track count, total video samples, width,
|
||||
/// height). Used to verify a recording is a valid, parseable, playable file.
|
||||
pub fn probe_mp4(path: impl AsRef<Path>) -> Result<(u32, u64, u16, u16)> {
|
||||
let f = File::open(path.as_ref()).map_err(other)?;
|
||||
let size = f.metadata().map_err(other)?.len();
|
||||
let reader = mp4::Mp4Reader::read_header(BufReader::new(f), size).map_err(other)?;
|
||||
let (mut vtracks, mut samples, mut w, mut h) = (0u32, 0u64, 0u16, 0u16);
|
||||
for track in reader.tracks().values() {
|
||||
if matches!(track.track_type(), Ok(mp4::TrackType::Video)) {
|
||||
vtracks += 1;
|
||||
samples += track.sample_count() as u64;
|
||||
w = track.width();
|
||||
h = track.height();
|
||||
}
|
||||
}
|
||||
Ok((vtracks, samples, w, h))
|
||||
}
|
||||
@@ -0,0 +1,585 @@
|
||||
//! WebRTC peer-connection helpers for the native P2P fast path (webrtc 0.17.2,
|
||||
//! the classic tokio-callback API). This module builds the `API` + a
|
||||
//! `RTCPeerConnection` from the ICE servers the relay hands us in `rtc.config`
|
||||
//! / `rtc.offer.ice`. The DataChannel transport ([`crate::transport`]) and the
|
||||
//! offer/answer signaling are layered on top of these.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::mpsc;
|
||||
use webrtc::api::interceptor_registry::register_default_interceptors;
|
||||
use webrtc::api::media_engine::MediaEngine;
|
||||
use webrtc::api::{APIBuilder, API};
|
||||
use webrtc::data_channel::data_channel_init::RTCDataChannelInit;
|
||||
use webrtc::data_channel::RTCDataChannel;
|
||||
use webrtc::ice_transport::ice_candidate_type::RTCIceCandidateType;
|
||||
use webrtc::ice_transport::ice_server::RTCIceServer;
|
||||
use webrtc::interceptor::registry::Registry;
|
||||
use webrtc::peer_connection::configuration::RTCConfiguration;
|
||||
use webrtc::peer_connection::sdp::session_description::RTCSessionDescription;
|
||||
use webrtc::peer_connection::RTCPeerConnection;
|
||||
|
||||
/// One ICE server exactly as it arrives in the relay's `rtc.config` reply and
|
||||
/// is forwarded in `rtc.offer.ice`: `{"urls":[...],"username":...,"credential":...}`.
|
||||
/// coturn REST creds (username = "<expiry>:nd", credential = base64 HMAC-SHA1)
|
||||
/// drop straight into webrtc's long-term TURN credential fields.
|
||||
#[derive(Debug, Clone, Deserialize, serde::Serialize)]
|
||||
pub struct RtcIceServerJson {
|
||||
#[serde(default)]
|
||||
pub urls: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub username: Option<String>,
|
||||
#[serde(default)]
|
||||
pub credential: Option<String>,
|
||||
}
|
||||
|
||||
/// Map the relay's ICE JSON into webrtc `RTCIceServer`s (0.17.2 has only
|
||||
/// urls/username/credential).
|
||||
pub fn ice_servers_from_json(ice: &[RtcIceServerJson]) -> Vec<RTCIceServer> {
|
||||
ice.iter()
|
||||
.filter(|s| !s.urls.is_empty())
|
||||
.map(|s| RTCIceServer {
|
||||
urls: s.urls.clone(),
|
||||
username: s.username.clone().unwrap_or_default(),
|
||||
credential: s.credential.clone().unwrap_or_default(),
|
||||
..Default::default()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Build a webrtc `API` with the default media engine + interceptors. One per
|
||||
/// peer connection is fine (cheap).
|
||||
///
|
||||
/// Raises the SCTP max send message size to 256 KiB (webrtc-rs defaults to
|
||||
/// 64 KiB), matching the web client's advertised `a=max-message-size:262144`.
|
||||
/// Without this a whole ~200 KiB video keyframe is rejected with
|
||||
/// `ErrOutboundPacketTooLarge`; with it the native peer can send whole E2E
|
||||
/// frames exactly like the browser / aiortc, which is what web/Python interop
|
||||
/// on the `media` channel requires. (Harmless for the native single channel,
|
||||
/// which fragments at 16 KiB.)
|
||||
pub fn build_api() -> Result<API> {
|
||||
use webrtc::api::setting_engine::{SctpMaxMessageSize, SettingEngine};
|
||||
let mut media = MediaEngine::default();
|
||||
let mut registry = Registry::new();
|
||||
registry = register_default_interceptors(registry, &mut media)?;
|
||||
let mut se = SettingEngine::default();
|
||||
se.set_sctp_max_message_size_can_send(SctpMaxMessageSize::Bounded(262144));
|
||||
Ok(APIBuilder::new()
|
||||
.with_media_engine(media)
|
||||
.with_interceptor_registry(registry)
|
||||
.with_setting_engine(se)
|
||||
.build())
|
||||
}
|
||||
|
||||
/// A peer-connection configuration from a list of ICE servers.
|
||||
pub fn rtc_config(ice_servers: Vec<RTCIceServer>) -> RTCConfiguration {
|
||||
RTCConfiguration {
|
||||
ice_servers,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new peer connection on `api` with `cfg`.
|
||||
pub async fn new_pc(api: &API, cfg: RTCConfiguration) -> Result<Arc<RTCPeerConnection>> {
|
||||
Ok(Arc::new(api.new_peer_connection(cfg).await?))
|
||||
}
|
||||
|
||||
/// Cap on how long we wait for ICE gathering to finish before sending the SDP.
|
||||
/// The web/Python peers are NON-TRICKLE (candidates must be embedded in the
|
||||
/// offer/answer), so we gather to completion — but never hang on a stuck
|
||||
/// network interface: 2.5 s matches the reference signaling timeout.
|
||||
pub const GATHER_CAP: Duration = Duration::from_millis(2500);
|
||||
|
||||
/// The max DataChannel message size we advertise (bytes) — matches the web
|
||||
/// client's `a=max-message-size:262144`.
|
||||
pub const MAX_MESSAGE_SIZE: u32 = 262144;
|
||||
|
||||
/// Ensure the SDP's `m=application` section advertises our max-message-size.
|
||||
/// webrtc-rs never emits `a=max-message-size`, so a peer defaults our receive
|
||||
/// cap to 64 KiB and refuses to send us a whole ~200 KiB keyframe. Injecting it
|
||||
/// (right after `a=sctp-port:`) raises the PEER's send-to-us cap to 256 KiB —
|
||||
/// and is exactly what aiortc / the browser advertise, so it is also how we
|
||||
/// receive whole video frames from the Python host over the `media` channel.
|
||||
pub fn advertise_max_message_size(sdp: String) -> String {
|
||||
if sdp.contains("a=max-message-size:") {
|
||||
return sdp; // already present (e.g. echoed from a web peer)
|
||||
}
|
||||
let nl = if sdp.contains("\r\n") { "\r\n" } else { "\n" };
|
||||
let needle = "a=sctp-port:";
|
||||
if let Some(pos) = sdp.find(needle) {
|
||||
// insert the attribute on its own line after the sctp-port line
|
||||
if let Some(eol) = sdp[pos..].find('\n') {
|
||||
let insert_at = pos + eol + 1;
|
||||
let mut out = String::with_capacity(sdp.len() + 32);
|
||||
out.push_str(&sdp[..insert_at]);
|
||||
out.push_str(&format!("a=max-message-size:{MAX_MESSAGE_SIZE}{nl}"));
|
||||
out.push_str(&sdp[insert_at..]);
|
||||
return out;
|
||||
}
|
||||
}
|
||||
sdp
|
||||
}
|
||||
|
||||
/// Offerer (viewer): create the offer, then gather ICE **to completion** and
|
||||
/// return the full local SDP for `rtc.offer.sdp`. Non-trickle: every candidate
|
||||
/// is embedded in the returned SDP, because the Python host has no trickle
|
||||
/// candidate handler. The caller must create its DataChannel(s) on `pc` BEFORE
|
||||
/// calling (they shape the offer). `gathering_complete_promise` is taken before
|
||||
/// `set_local_description`, which is what actually starts gathering.
|
||||
pub async fn offer_nontrickle(pc: &Arc<RTCPeerConnection>, gather_cap: Duration) -> Result<String> {
|
||||
let offer = pc.create_offer(None).await?;
|
||||
let mut gather = pc.gathering_complete_promise().await;
|
||||
pc.set_local_description(offer).await?;
|
||||
let _ = tokio::time::timeout(gather_cap, gather.recv()).await;
|
||||
let local = pc
|
||||
.local_description()
|
||||
.await
|
||||
.ok_or_else(|| anyhow!("no local description after ICE gathering"))?;
|
||||
Ok(advertise_max_message_size(local.sdp))
|
||||
}
|
||||
|
||||
/// Answerer (host): apply the remote offer SDP, create the answer, gather ICE
|
||||
/// to completion (non-trickle) and return the full local SDP for
|
||||
/// `rtc.answer.sdp`. Register `on_data_channel` on `pc` BEFORE calling so the
|
||||
/// viewer-created channel is captured.
|
||||
pub async fn answer_nontrickle(
|
||||
pc: &Arc<RTCPeerConnection>,
|
||||
offer_sdp: &str,
|
||||
gather_cap: Duration,
|
||||
) -> Result<String> {
|
||||
let offer = RTCSessionDescription::offer(offer_sdp.to_string())?;
|
||||
pc.set_remote_description(offer).await?;
|
||||
let answer = pc.create_answer(None).await?;
|
||||
let mut gather = pc.gathering_complete_promise().await;
|
||||
pc.set_local_description(answer).await?;
|
||||
let _ = tokio::time::timeout(gather_cap, gather.recv()).await;
|
||||
let local = pc
|
||||
.local_description()
|
||||
.await
|
||||
.ok_or_else(|| anyhow!("no local description after ICE gathering"))?;
|
||||
Ok(advertise_max_message_size(local.sdp))
|
||||
}
|
||||
|
||||
/// Offerer (viewer): apply the remote answer SDP received in `rtc.answer`.
|
||||
pub async fn apply_answer(pc: &Arc<RTCPeerConnection>, answer_sdp: &str) -> Result<()> {
|
||||
let answer = RTCSessionDescription::answer(answer_sdp.to_string())?;
|
||||
pc.set_remote_description(answer).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Wait for the server to return rtc.config before falling back to a default
|
||||
/// STUN server. 2.5 s matches the reference signaling timeout.
|
||||
pub const CONFIG_CAP: Duration = Duration::from_millis(2500);
|
||||
/// How long to wait for the peer's rtc.answer.
|
||||
pub const ANSWER_CAP: Duration = Duration::from_secs(10);
|
||||
/// How long to wait for the pc to connect + the DataChannel to open.
|
||||
pub const CONNECT_CAP: Duration = Duration::from_secs(15);
|
||||
|
||||
/// Default ICE if the server does not answer rtc.config in time. A public STUN
|
||||
/// is enough to still find a server-reflexive path; host/LAN pairs connect
|
||||
/// without it (and loopback needs no ICE server at all).
|
||||
pub fn default_ice() -> Vec<RtcIceServerJson> {
|
||||
vec![RtcIceServerJson {
|
||||
urls: vec!["stun:stun.l.google.com:19302".into()],
|
||||
username: None,
|
||||
credential: None,
|
||||
}]
|
||||
}
|
||||
|
||||
/// Pull rtc.* Values from `rtc_in` until one matches `want` (dropping any other
|
||||
/// rtc.* that arrives first), capped by `cap`. rtc.unavailable short-circuits to
|
||||
/// an error so the caller abandons the upgrade and stays on the relay.
|
||||
async fn wait_rtc(
|
||||
rtc_in: &mut mpsc::UnboundedReceiver<Value>,
|
||||
want: &str,
|
||||
cap: Duration,
|
||||
) -> Result<Value> {
|
||||
let deadline = tokio::time::Instant::now() + cap;
|
||||
loop {
|
||||
match tokio::time::timeout_at(deadline, rtc_in.recv()).await {
|
||||
Ok(Some(v)) => {
|
||||
let t = v.get("t").and_then(|t| t.as_str()).unwrap_or("");
|
||||
if t == want {
|
||||
return Ok(v);
|
||||
}
|
||||
if t == "rtc.unavailable" {
|
||||
return Err(anyhow!("peer declined P2P (rtc.unavailable)"));
|
||||
}
|
||||
// ignore any other rtc.* and keep waiting
|
||||
}
|
||||
Ok(None) => return Err(anyhow!("rtc signaling channel closed")),
|
||||
Err(_) => return Err(anyhow!("timed out waiting for {want}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Viewer offerer: run the relay-brokered P2P upgrade handshake and return the
|
||||
/// connected pc + its DataChannel once the channel opens. `send_text` enqueues
|
||||
/// an rtc.* text message onto the relay (the caller supplies one that carries
|
||||
/// `sid`); `rtc_in` yields inbound rtc.config / rtc.answer / rtc.unavailable
|
||||
/// Values the relay loop forwards. When `ice_override` is `None` the offerer
|
||||
/// asks the server for ICE via rtc.config.req (default STUN on timeout); tests
|
||||
/// pass `Some(vec![])` for a loopback host-candidate-only pc.
|
||||
pub async fn upgrade_offerer(
|
||||
api: &API,
|
||||
sid: &str,
|
||||
ice_override: Option<Vec<RtcIceServerJson>>,
|
||||
channel_label: &str,
|
||||
send_text: &(dyn Fn(String) + Send + Sync),
|
||||
rtc_in: &mut mpsc::UnboundedReceiver<Value>,
|
||||
) -> Result<(Arc<RTCPeerConnection>, Arc<RTCDataChannel>)> {
|
||||
let ice = match ice_override {
|
||||
Some(i) => i,
|
||||
None => {
|
||||
send_text(json!({ "t": "rtc.config.req", "sid": sid }).to_string());
|
||||
match wait_rtc(rtc_in, "rtc.config", CONFIG_CAP).await {
|
||||
Ok(v) => v
|
||||
.get("iceServers")
|
||||
.and_then(|x| serde_json::from_value::<Vec<RtcIceServerJson>>(x.clone()).ok())
|
||||
.unwrap_or_else(default_ice),
|
||||
Err(_) => default_ice(),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let pc = new_pc(api, rtc_config(ice_servers_from_json(&ice))).await?;
|
||||
let (open_tx, mut open_rx) = mpsc::channel::<()>(1);
|
||||
let dc = pc
|
||||
.create_data_channel(
|
||||
channel_label,
|
||||
Some(RTCDataChannelInit {
|
||||
ordered: Some(true),
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
{
|
||||
let tx = open_tx.clone();
|
||||
dc.on_open(Box::new(move || {
|
||||
let tx = tx.clone();
|
||||
Box::pin(async move {
|
||||
let _ = tx.send(()).await;
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
// Non-trickle: gather fully, then send the offer with the ICE config echoed
|
||||
// through so a TURN-only carrier-NAT host can reach us.
|
||||
let offer_sdp = offer_nontrickle(&pc, GATHER_CAP).await?;
|
||||
// `native:true` marks a native single-channel offerer: a native host only
|
||||
// answers these, and declines (rtc.unavailable) a web/3-channel offerer so
|
||||
// that pairing safely stays on the relay.
|
||||
send_text(
|
||||
json!({ "t": "rtc.offer", "sid": sid, "sdp": offer_sdp, "ice": ice, "native": true })
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
let answer = wait_rtc(rtc_in, "rtc.answer", ANSWER_CAP).await?;
|
||||
let answer_sdp = answer
|
||||
.get("sdp")
|
||||
.and_then(|s| s.as_str())
|
||||
.ok_or_else(|| anyhow!("rtc.answer missing sdp"))?;
|
||||
apply_answer(&pc, answer_sdp).await?;
|
||||
|
||||
tokio::time::timeout(CONNECT_CAP, open_rx.recv())
|
||||
.await
|
||||
.map_err(|_| anyhow!("P2P DataChannel did not open"))?
|
||||
.ok_or_else(|| anyhow!("P2P open signal dropped"))?;
|
||||
Ok((pc, dc))
|
||||
}
|
||||
|
||||
/// Host answerer: handle one inbound rtc.offer. Build the answerer pc from the
|
||||
/// offer plus the `ice` the offerer echoed, capture the viewer-created channel
|
||||
/// via on_data_channel, create+send the answer via `send_text` (carrying
|
||||
/// `sid`), and return the connected pc + channel once it opens.
|
||||
pub async fn answer_offer(
|
||||
api: &API,
|
||||
sid: &str,
|
||||
offer_sdp: &str,
|
||||
ice: Vec<RtcIceServerJson>,
|
||||
send_text: &(dyn Fn(String) + Send + Sync),
|
||||
) -> Result<(Arc<RTCPeerConnection>, Arc<RTCDataChannel>)> {
|
||||
let pc = new_pc(api, rtc_config(ice_servers_from_json(&ice))).await?;
|
||||
let (dc_tx, mut dc_rx) = mpsc::channel::<Arc<RTCDataChannel>>(1);
|
||||
let (open_tx, mut open_rx) = mpsc::channel::<()>(1);
|
||||
{
|
||||
let dc_tx = dc_tx.clone();
|
||||
let open_tx = open_tx.clone();
|
||||
pc.on_data_channel(Box::new(move |d: Arc<RTCDataChannel>| {
|
||||
let dc_tx = dc_tx.clone();
|
||||
let open_tx = open_tx.clone();
|
||||
Box::pin(async move {
|
||||
let ot = open_tx.clone();
|
||||
d.on_open(Box::new(move || {
|
||||
let ot = ot.clone();
|
||||
Box::pin(async move {
|
||||
let _ = ot.send(()).await;
|
||||
})
|
||||
}));
|
||||
let _ = dc_tx.send(d).await;
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
let answer_sdp = answer_nontrickle(&pc, offer_sdp, GATHER_CAP).await?;
|
||||
send_text(json!({ "t": "rtc.answer", "sid": sid, "sdp": answer_sdp }).to_string());
|
||||
|
||||
let dc = tokio::time::timeout(CONNECT_CAP, dc_rx.recv())
|
||||
.await
|
||||
.map_err(|_| anyhow!("no DataChannel from the offerer"))?
|
||||
.ok_or_else(|| anyhow!("DataChannel signal dropped"))?;
|
||||
tokio::time::timeout(CONNECT_CAP, open_rx.recv())
|
||||
.await
|
||||
.map_err(|_| anyhow!("P2P DataChannel did not open (host)"))?
|
||||
.ok_or_else(|| anyhow!("host open signal dropped"))?;
|
||||
Ok((pc, dc))
|
||||
}
|
||||
|
||||
/// True when the pc's nominated ICE pair is a DIRECT path (host / server-
|
||||
/// reflexive on both ends) rather than a TURN relay. The relay-first state
|
||||
/// machine only reports "Direct P2P" and moves media off the WSS relay for a
|
||||
/// direct pair: a TURN-relayed hop is no faster than our own relay, and a
|
||||
/// not-yet-selected pair is unknown, so both keep media on the relay.
|
||||
pub async fn selected_pair_is_direct(pc: &Arc<RTCPeerConnection>) -> bool {
|
||||
let dtls = pc.sctp().transport();
|
||||
let ice = dtls.ice_transport();
|
||||
match ice.get_selected_candidate_pair().await {
|
||||
Some(pair) => {
|
||||
pair.local.typ != RTCIceCandidateType::Relay
|
||||
&& pair.remote.typ != RTCIceCandidateType::Relay
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// The web/Python 3-channel DataChannel topology. Labels + options match
|
||||
/// `app.js` exactly so a native peer interops with the browser and the Python
|
||||
/// host: `media` carries video (unreliable, 150 ms lifetime — drop beats late),
|
||||
/// `ctl` carries input/clipboard/file/keyframe (reliable, ordered), `move`
|
||||
/// carries pointer motion (unreliable, no retransmit).
|
||||
pub struct WebChannels {
|
||||
pub media: Arc<RTCDataChannel>,
|
||||
pub ctl: Arc<RTCDataChannel>,
|
||||
pub mov: Arc<RTCDataChannel>,
|
||||
}
|
||||
|
||||
/// Offerer: create the three web channels on `pc`. Call BEFORE `offer_nontrickle`
|
||||
/// so they are announced in the SDP. The option set is byte-for-byte the web
|
||||
/// client's (`media`{ordered:false,maxPacketLifeTime:150},
|
||||
/// `ctl`{ordered:true}, `move`{ordered:false,maxRetransmits:0}).
|
||||
pub async fn create_web_channels(pc: &Arc<RTCPeerConnection>) -> Result<WebChannels> {
|
||||
let media = pc
|
||||
.create_data_channel(
|
||||
"media",
|
||||
Some(RTCDataChannelInit {
|
||||
ordered: Some(false),
|
||||
max_packet_life_time: Some(150),
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let ctl = pc
|
||||
.create_data_channel(
|
||||
"ctl",
|
||||
Some(RTCDataChannelInit {
|
||||
ordered: Some(true),
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let mov = pc
|
||||
.create_data_channel(
|
||||
"move",
|
||||
Some(RTCDataChannelInit {
|
||||
ordered: Some(false),
|
||||
max_retransmits: Some(0),
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
Ok(WebChannels { media, ctl, mov })
|
||||
}
|
||||
|
||||
/// Like [`build_api`] but with DETACHED DataChannels enabled. In detach mode we
|
||||
/// read each channel with our OWN buffer instead of webrtc-rs's fixed 64 KiB
|
||||
/// `read_loop` buffer — which is the ONLY way to receive a whole >64 KiB
|
||||
/// DataChannel message on webrtc-rs 0.17.2. This is required for the web/Python
|
||||
/// `media` channel, where the browser / aiortc send whole ~100-200 KiB video
|
||||
/// frames. Use this API for the web 3-channel topology; [`build_api`] (on_message,
|
||||
/// 64 KiB) stays the default for the native single channel, which fragments at
|
||||
/// 16 KiB and never needs a large receive.
|
||||
pub fn build_api_web() -> Result<API> {
|
||||
use webrtc::api::setting_engine::{SctpMaxMessageSize, SettingEngine};
|
||||
let mut media = MediaEngine::default();
|
||||
let mut registry = Registry::new();
|
||||
registry = register_default_interceptors(registry, &mut media)?;
|
||||
let mut se = SettingEngine::default();
|
||||
se.set_sctp_max_message_size_can_send(SctpMaxMessageSize::Bounded(MAX_MESSAGE_SIZE));
|
||||
se.detach_data_channels();
|
||||
Ok(APIBuilder::new()
|
||||
.with_media_engine(media)
|
||||
.with_interceptor_registry(registry)
|
||||
.with_setting_engine(se)
|
||||
.build())
|
||||
}
|
||||
|
||||
/// The web DataChannels the viewer actually uses, after detaching (raw
|
||||
/// read/write handles with our own buffer — the only way to receive whole
|
||||
/// >64KiB frames on webrtc-rs 0.17.2). `move` is created for topology/interop
|
||||
/// but the viewer never reads/writes it, so it is not detached (fewer detaches
|
||||
/// = more reliable).
|
||||
pub struct WebDetached {
|
||||
pub media: Arc<webrtc::data::data_channel::DataChannel>,
|
||||
pub ctl: Arc<webrtc::data::data_channel::DataChannel>,
|
||||
}
|
||||
|
||||
/// Viewer offerer, WEB topology (interop with the browser / Python host). Builds
|
||||
/// the 3 channels (media/ctl/move, labels+options per app.js), runs the
|
||||
/// non-trickle rtc.offer/answer handshake over `send_text`+`rtc_in` (offer does
|
||||
/// NOT carry `native` — it must look like a standard web offer so the Python
|
||||
/// host accepts it), waits for all 3 to open, then detaches them SEQUENTIALLY
|
||||
/// (concurrent multi-channel detach deadlocks) and returns the raw handles. The
|
||||
/// pc must be built from [`build_api_web`] (detach enabled).
|
||||
pub async fn upgrade_offerer_web(
|
||||
api: &API,
|
||||
sid: &str,
|
||||
ice_override: Option<Vec<RtcIceServerJson>>,
|
||||
send_text: &(dyn Fn(String) + Send + Sync),
|
||||
rtc_in: &mut mpsc::UnboundedReceiver<Value>,
|
||||
) -> Result<(Arc<RTCPeerConnection>, WebDetached)> {
|
||||
let ice = match ice_override {
|
||||
Some(i) => i,
|
||||
None => {
|
||||
send_text(json!({ "t": "rtc.config.req", "sid": sid }).to_string());
|
||||
match wait_rtc(rtc_in, "rtc.config", CONFIG_CAP).await {
|
||||
Ok(v) => v
|
||||
.get("iceServers")
|
||||
.and_then(|x| serde_json::from_value::<Vec<RtcIceServerJson>>(x.clone()).ok())
|
||||
.unwrap_or_else(default_ice),
|
||||
Err(_) => default_ice(),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let pc = new_pc(api, rtc_config(ice_servers_from_json(&ice))).await?;
|
||||
let web = create_web_channels(&pc).await?;
|
||||
|
||||
// open signals for the three channels
|
||||
let (open_tx, mut open_rx) = mpsc::unbounded_channel::<&'static str>();
|
||||
for (lbl, ch) in [
|
||||
("media", &web.media),
|
||||
("ctl", &web.ctl),
|
||||
("move", &web.mov),
|
||||
] {
|
||||
let tx = open_tx.clone();
|
||||
ch.on_open(Box::new(move || {
|
||||
let tx = tx.clone();
|
||||
Box::pin(async move {
|
||||
let _ = tx.send(lbl);
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
let offer_sdp = offer_nontrickle(&pc, GATHER_CAP).await?;
|
||||
// NOTE: no `native` marker — this is a standard web offer.
|
||||
send_text(json!({ "t": "rtc.offer", "sid": sid, "sdp": offer_sdp, "ice": ice }).to_string());
|
||||
let answer = wait_rtc(rtc_in, "rtc.answer", ANSWER_CAP).await?;
|
||||
let answer_sdp = answer
|
||||
.get("sdp")
|
||||
.and_then(|s| s.as_str())
|
||||
.ok_or_else(|| anyhow!("rtc.answer missing sdp"))?;
|
||||
apply_answer(&pc, answer_sdp).await?;
|
||||
|
||||
// wait for the two channels we use (media+ctl) to open, capped.
|
||||
let mut opened = std::collections::HashSet::new();
|
||||
let deadline = tokio::time::Instant::now() + CONNECT_CAP;
|
||||
while !(opened.contains("media") && opened.contains("ctl")) {
|
||||
match tokio::time::timeout_at(deadline, open_rx.recv()).await {
|
||||
Ok(Some(l)) => {
|
||||
opened.insert(l);
|
||||
}
|
||||
_ => return Err(anyhow!("web channels did not open")),
|
||||
}
|
||||
}
|
||||
|
||||
// detach SEQUENTIALLY (concurrent detach deadlocks) — only media + ctl.
|
||||
let media = web.media.detach().await?;
|
||||
let ctl = web.ctl.detach().await?;
|
||||
Ok((pc, WebDetached { media, ctl }))
|
||||
}
|
||||
|
||||
/// Host answerer, WEB topology (a browser / web viewer connects to a native
|
||||
/// host). Builds the answerer pc, binds the viewer-created media/ctl/move
|
||||
/// channels by label, runs the non-trickle answer, waits for media+ctl to open,
|
||||
/// then detaches them SEQUENTIALLY. The pc must be built from [`build_api_web`].
|
||||
/// The host then SENDS video on media (whole frames, cap 256KiB) and READS
|
||||
/// input on ctl.
|
||||
pub async fn answer_offer_web(
|
||||
api: &API,
|
||||
sid: &str,
|
||||
offer_sdp: &str,
|
||||
ice: Vec<RtcIceServerJson>,
|
||||
send_text: &(dyn Fn(String) + Send + Sync),
|
||||
) -> Result<(Arc<RTCPeerConnection>, WebDetached)> {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
|
||||
let pc = new_pc(api, rtc_config(ice_servers_from_json(&ice))).await?;
|
||||
|
||||
// bind channels by label + signal opens
|
||||
let chans: Arc<StdMutex<HashMap<String, Arc<RTCDataChannel>>>> =
|
||||
Arc::new(StdMutex::new(HashMap::new()));
|
||||
let (open_tx, mut open_rx) = mpsc::unbounded_channel::<String>();
|
||||
{
|
||||
let chans = chans.clone();
|
||||
let open_tx = open_tx.clone();
|
||||
pc.on_data_channel(Box::new(move |d: Arc<RTCDataChannel>| {
|
||||
let chans = chans.clone();
|
||||
let open_tx = open_tx.clone();
|
||||
let label = d.label().to_string();
|
||||
Box::pin(async move {
|
||||
chans.lock().unwrap().insert(label.clone(), d.clone());
|
||||
let tx = open_tx.clone();
|
||||
let l = label.clone();
|
||||
d.on_open(Box::new(move || {
|
||||
let tx = tx.clone();
|
||||
let l = l.clone();
|
||||
Box::pin(async move {
|
||||
let _ = tx.send(l);
|
||||
})
|
||||
}));
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
let answer_sdp = answer_nontrickle(&pc, offer_sdp, GATHER_CAP).await?;
|
||||
send_text(json!({ "t": "rtc.answer", "sid": sid, "sdp": answer_sdp }).to_string());
|
||||
|
||||
// wait for media + ctl to open, capped
|
||||
let mut opened = std::collections::HashSet::new();
|
||||
let deadline = tokio::time::Instant::now() + CONNECT_CAP;
|
||||
while !(opened.contains("media") && opened.contains("ctl")) {
|
||||
match tokio::time::timeout_at(deadline, open_rx.recv()).await {
|
||||
Ok(Some(l)) => {
|
||||
opened.insert(l);
|
||||
}
|
||||
_ => return Err(anyhow!("web channels (media/ctl) did not open")),
|
||||
}
|
||||
}
|
||||
|
||||
// grab wrappers, then detach SEQUENTIALLY (concurrent detach deadlocks).
|
||||
let (media_w, ctl_w) = {
|
||||
let m = chans.lock().unwrap();
|
||||
(
|
||||
m.get("media").cloned().ok_or_else(|| anyhow!("no media channel"))?,
|
||||
m.get("ctl").cloned().ok_or_else(|| anyhow!("no ctl channel"))?,
|
||||
)
|
||||
};
|
||||
let media = media_w.detach().await?;
|
||||
let ctl = ctl_w.detach().await?;
|
||||
Ok((pc, WebDetached { media, ctl }))
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,471 @@
|
||||
//! Pluggable transport for the session protocol. The E2E AES-GCM framing sits
|
||||
//! ABOVE this layer on both the viewer and host, so a transport only ever moves
|
||||
//! opaque frames (binary media/file chunks, or text control JSON). The current
|
||||
//! WSS relay ([`WsRelay`]) and the P2P WebRTC DataChannel (added later) both
|
||||
//! implement [`Transport`], so `connect_with` / `run_host_with` carry either and
|
||||
//! the relay→P2P upgrade / downgrade can swap transports without touching the
|
||||
//! protocol, crypto, or framing.
|
||||
|
||||
use crate::error::{NdError, Result};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
|
||||
|
||||
use bytes::Bytes;
|
||||
use webrtc::data_channel::data_channel_message::DataChannelMessage;
|
||||
use webrtc::data_channel::RTCDataChannel;
|
||||
|
||||
/// One message crossing a transport. Opaque here; the caller encrypts/decrypts
|
||||
/// and dispatches by channel byte above this layer.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum TransportMsg {
|
||||
Binary(Vec<u8>),
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Which logical stream a binary frame belongs to. The caller knows this from
|
||||
/// the (pre-encryption) channel byte; the transport cannot see it because the
|
||||
/// byte is inside the E2E ciphertext. It decides per-lane delivery: `Media` is
|
||||
/// drop-tolerant (video — drop a fresh frame rather than queue it late), `Ctl`
|
||||
/// must be reliable (file bytes, control), `Move` is drop-tolerant pointer
|
||||
/// motion. Maps 1:1 to the web/Python `media` / `ctl` / `move` DataChannels;
|
||||
/// the relay and the native single channel use it only for back-pressure.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Lane {
|
||||
Media,
|
||||
Ctl,
|
||||
Move,
|
||||
}
|
||||
|
||||
/// A bidirectional message transport. Used from a single task (the viewer's
|
||||
/// socket task / the host's `run_host` loop): `recv` is awaited in one
|
||||
/// `tokio::select!` arm and `send_*` in other arms' bodies, so `&mut self` for
|
||||
/// all three is sound (select drops the non-chosen futures first).
|
||||
#[async_trait]
|
||||
pub trait Transport: Send + 'static {
|
||||
async fn send_binary(&mut self, data: Vec<u8>) -> Result<()>;
|
||||
async fn send_text(&mut self, text: String) -> Result<()>;
|
||||
/// Next inbound message, or `None` when the transport has closed.
|
||||
async fn recv(&mut self) -> Option<Result<TransportMsg>>;
|
||||
|
||||
/// Send a binary frame tagged with its [`Lane`]. The default ignores the
|
||||
/// lane (relay, native single channel today), so existing transports need
|
||||
/// no change; the P2P multiplexer overrides it to pick the channel /
|
||||
/// back-pressure policy. Callers that know the channel (video vs file)
|
||||
/// should prefer this over [`Transport::send_binary`].
|
||||
async fn send_binary_lane(&mut self, data: Vec<u8>, lane: Lane) -> Result<()> {
|
||||
let _ = lane;
|
||||
self.send_binary(data).await
|
||||
}
|
||||
}
|
||||
|
||||
/// The WSS relay transport — the current path, behaviourally identical to the
|
||||
/// pre-seam code. Wraps the tungstenite WebSocket (TLS to the relay); the server
|
||||
/// relays frames opaquely between the two session peers.
|
||||
pub struct WsRelay {
|
||||
ws: WebSocketStream<MaybeTlsStream<TcpStream>>,
|
||||
}
|
||||
|
||||
impl WsRelay {
|
||||
/// Dial the relay (TLS). Mirrors the old `connect_async` + error mapping.
|
||||
pub async fn connect(ws_url: &str) -> Result<Self> {
|
||||
let (ws, _resp) = connect_async(ws_url)
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))?;
|
||||
Ok(Self { ws })
|
||||
}
|
||||
|
||||
/// Wrap an already-connected WebSocket (e.g. one opened during a handshake
|
||||
/// elsewhere).
|
||||
pub fn from_ws(ws: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
|
||||
Self { ws }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for WsRelay {
|
||||
async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
|
||||
self.ws
|
||||
.send(Message::Binary(data))
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))
|
||||
}
|
||||
|
||||
async fn send_text(&mut self, text: String) -> Result<()> {
|
||||
self.ws
|
||||
.send(Message::Text(text))
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))
|
||||
}
|
||||
|
||||
async fn recv(&mut self) -> Option<Result<TransportMsg>> {
|
||||
loop {
|
||||
match self.ws.next().await {
|
||||
Some(Ok(Message::Binary(b))) => return Some(Ok(TransportMsg::Binary(b))),
|
||||
Some(Ok(Message::Text(t))) => return Some(Ok(TransportMsg::Text(t))),
|
||||
Some(Ok(Message::Close(_))) | None => return None,
|
||||
// Ping/Pong/Frame are handled by tungstenite; skip at this layer.
|
||||
Some(Ok(_)) => continue,
|
||||
Some(Err(e)) => return Some(Err(NdError::Connect(e.to_string()))),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================ P2P WebRTC DataChannel ============================
|
||||
|
||||
/// 9-byte fragment header prepended to every binary chunk:
|
||||
/// `[frame_id:u32 BE][idx:u16 BE][count:u16 BE][flags:u8]`. `flags` bit0 is
|
||||
/// reserved for a keyframe hint (set by a higher layer in a later phase); the
|
||||
/// reassembler ignores it.
|
||||
const FRAG_HEADER: usize = 9;
|
||||
/// Each SCTP message stays at or below 16 KiB (header + payload) — safely under
|
||||
/// the negotiated max-message-size, and small enough that one slow frame never
|
||||
/// head-of-line-blocks the queue for long.
|
||||
const MAX_CHUNK_PAYLOAD: usize = 16 * 1024 - FRAG_HEADER;
|
||||
/// Pause sending once this much is queued in the SCTP send buffer...
|
||||
const HIGH_WATER: usize = 1024 * 1024;
|
||||
/// ...and resume when it drains to here (`on_buffered_amount_low` fires).
|
||||
const LOW_WATER: usize = 512 * 1024;
|
||||
|
||||
/// Reassembly accumulator for one multi-chunk binary frame.
|
||||
struct FrameAsm {
|
||||
count: u16,
|
||||
got: u16,
|
||||
parts: Vec<Option<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl FrameAsm {
|
||||
fn new(count: u16) -> Self {
|
||||
Self {
|
||||
count,
|
||||
got: 0,
|
||||
parts: vec![None; count as usize],
|
||||
}
|
||||
}
|
||||
fn put(&mut self, idx: u16, payload: &[u8]) {
|
||||
let i = idx as usize;
|
||||
if i < self.parts.len() && self.parts[i].is_none() {
|
||||
self.parts[i] = Some(payload.to_vec());
|
||||
self.got += 1;
|
||||
}
|
||||
}
|
||||
fn complete(&self) -> bool {
|
||||
self.got == self.count
|
||||
}
|
||||
/// Concatenate the chunks in index order. Because the chunks are contiguous
|
||||
/// slices of the original nonce||ciphertext||tag blob, the rejoined bytes
|
||||
/// are byte-identical — a mis-ordered or missing chunk would corrupt the
|
||||
/// blob and the viewer AES-GCM decrypt_frame would fail the auth tag, which
|
||||
/// is exactly the integrity guarantee we want.
|
||||
fn concat(self) -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
for p in self.parts.into_iter().flatten() {
|
||||
out.extend_from_slice(&p);
|
||||
}
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// Install the inbound reassembly handler on `dc`: SCTP string messages become
|
||||
/// `TransportMsg::Text`; fragmented binary chunks are rejoined by `frame_id`
|
||||
/// and pushed as exactly one contiguous `TransportMsg::Binary` (so the E2E
|
||||
/// `nonce||ct||tag` blob stays intact for AES-GCM). Shared by `WebRtcChannel`
|
||||
/// (direct P2P) and the relay/P2P multiplexer so both frame the wire identically.
|
||||
pub fn install_reassembler(dc: &Arc<RTCDataChannel>, out: mpsc::UnboundedSender<TransportMsg>) {
|
||||
let reasm: Arc<Mutex<HashMap<u32, FrameAsm>>> = Arc::new(Mutex::new(HashMap::new()));
|
||||
dc.on_message(Box::new(move |msg: DataChannelMessage| {
|
||||
let out = out.clone();
|
||||
let reasm = reasm.clone();
|
||||
Box::pin(async move {
|
||||
if msg.is_string {
|
||||
if let Ok(s) = String::from_utf8(msg.data.to_vec()) {
|
||||
let _ = out.send(TransportMsg::Text(s));
|
||||
}
|
||||
return;
|
||||
}
|
||||
let data = &msg.data;
|
||||
if data.len() < FRAG_HEADER {
|
||||
return; // malformed; drop
|
||||
}
|
||||
let frame_id = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
|
||||
let idx = u16::from_be_bytes([data[4], data[5]]);
|
||||
let count = u16::from_be_bytes([data[6], data[7]]);
|
||||
// data[8] = flags (bit0 = keyframe), reserved; unused for reassembly.
|
||||
let payload = &data[FRAG_HEADER..];
|
||||
match count {
|
||||
0 => {}
|
||||
1 => {
|
||||
let _ = out.send(TransportMsg::Binary(payload.to_vec()));
|
||||
}
|
||||
_ => {
|
||||
let mut map = reasm.lock().unwrap();
|
||||
let entry = map.entry(frame_id).or_insert_with(|| FrameAsm::new(count));
|
||||
entry.put(idx, payload);
|
||||
if entry.complete() {
|
||||
if let Some(asm) = map.remove(&frame_id) {
|
||||
drop(map);
|
||||
let _ = out.send(TransportMsg::Binary(asm.concat()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
/// Fragment `data` into `<=16 KiB` SCTP messages (9-byte header each) under one
|
||||
/// `frame_id` and send them in order on `dc`. No back-pressure here — the caller
|
||||
/// applies its own policy (wait vs drop-oldest) around this. Returns the error
|
||||
/// from the first failed `dc.send`.
|
||||
pub async fn send_fragmented(dc: &Arc<RTCDataChannel>, frame_id: u32, data: &[u8]) -> Result<()> {
|
||||
let total = data.len();
|
||||
let count = ((total + MAX_CHUNK_PAYLOAD - 1) / MAX_CHUNK_PAYLOAD).max(1);
|
||||
if count > u16::MAX as usize {
|
||||
return Err(NdError::Connect(format!(
|
||||
"frame too large to fragment: {total} bytes"
|
||||
)));
|
||||
}
|
||||
for idx in 0..count {
|
||||
let start = idx * MAX_CHUNK_PAYLOAD;
|
||||
let end = (start + MAX_CHUNK_PAYLOAD).min(total);
|
||||
let chunk = &data[start..end];
|
||||
let mut framed = Vec::with_capacity(FRAG_HEADER + chunk.len());
|
||||
framed.extend_from_slice(&frame_id.to_be_bytes());
|
||||
framed.extend_from_slice(&(idx as u16).to_be_bytes());
|
||||
framed.extend_from_slice(&(count as u16).to_be_bytes());
|
||||
framed.push(0u8); // flags (bit0 = keyframe) reserved
|
||||
framed.extend_from_slice(chunk);
|
||||
dc.send(&Bytes::from(framed))
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A [`Transport`] over a single reliable+ordered WebRTC DataChannel — the
|
||||
/// native-to-native direct P2P fast path. The E2E AES-GCM framing is unchanged
|
||||
/// above this layer; this channel only moves opaque binary frames (fragmented
|
||||
/// to fit SCTP, reassembled whole so the GCM blob stays contiguous) and opaque
|
||||
/// text control JSON (sent as SCTP string messages, so the peer tells the two
|
||||
/// apart with `DataChannelMessage::is_string`).
|
||||
pub struct WebRtcChannel {
|
||||
dc: Arc<RTCDataChannel>,
|
||||
inbound: mpsc::UnboundedReceiver<TransportMsg>,
|
||||
/// Drains one token whenever the SCTP send buffer falls below `LOW_WATER`.
|
||||
low_water_rx: mpsc::Receiver<()>,
|
||||
/// Set true by `on_close`; makes `recv` return `None` once inbound is empty.
|
||||
closed_flag: Arc<AtomicBool>,
|
||||
closed_rx: mpsc::Receiver<()>,
|
||||
next_frame_id: u32,
|
||||
}
|
||||
|
||||
impl WebRtcChannel {
|
||||
/// Wrap an open DataChannel. Registers the inbound reassembly handler, the
|
||||
/// back-pressure low-water signal, and the close signal. Call once the
|
||||
/// channel is open (sends before open error at the SCTP layer).
|
||||
pub async fn new(dc: Arc<RTCDataChannel>) -> Self {
|
||||
let (in_tx, in_rx) = mpsc::unbounded_channel::<TransportMsg>();
|
||||
|
||||
// --- inbound: reassemble binary frames, pass text straight through.
|
||||
install_reassembler(&dc, in_tx);
|
||||
|
||||
// --- back-pressure: a size-1 mpsc is race-free vs Notify here. The
|
||||
// low-water event only fires when the buffer crosses DOWN through the
|
||||
// threshold, so a token can be pending at most once; try_send drops
|
||||
// extras.
|
||||
let (low_tx, low_rx) = mpsc::channel::<()>(1);
|
||||
dc.set_buffered_amount_low_threshold(LOW_WATER).await;
|
||||
dc.on_buffered_amount_low(Box::new(move || {
|
||||
let low_tx = low_tx.clone();
|
||||
Box::pin(async move {
|
||||
let _ = low_tx.try_send(());
|
||||
})
|
||||
}))
|
||||
.await;
|
||||
|
||||
// --- close signal so recv() can return None on a dropped channel.
|
||||
let closed_flag = Arc::new(AtomicBool::new(false));
|
||||
let (closed_tx, closed_rx) = mpsc::channel::<()>(1);
|
||||
let cf = closed_flag.clone();
|
||||
dc.on_close(Box::new(move || {
|
||||
let cf = cf.clone();
|
||||
let closed_tx = closed_tx.clone();
|
||||
Box::pin(async move {
|
||||
cf.store(true, Ordering::SeqCst);
|
||||
let _ = closed_tx.try_send(());
|
||||
})
|
||||
}));
|
||||
|
||||
Self {
|
||||
dc,
|
||||
inbound: in_rx,
|
||||
low_water_rx: low_rx,
|
||||
closed_flag,
|
||||
closed_rx,
|
||||
next_frame_id: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// The underlying channel, for stats / state inspection (upgrade logic).
|
||||
pub fn data_channel(&self) -> &Arc<RTCDataChannel> {
|
||||
&self.dc
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for WebRtcChannel {
|
||||
async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
|
||||
// Pause if the SCTP send queue is already deep: wait for it to drain
|
||||
// below LOW_WATER (fresh frames never pile up unbounded in RAM).
|
||||
while self.dc.buffered_amount().await + FRAG_HEADER + data.len() > HIGH_WATER {
|
||||
if self.low_water_rx.recv().await.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let frame_id = self.next_frame_id;
|
||||
self.next_frame_id = self.next_frame_id.wrapping_add(1);
|
||||
send_fragmented(&self.dc, frame_id, &data).await
|
||||
}
|
||||
|
||||
async fn send_text(&mut self, text: String) -> Result<()> {
|
||||
self.dc
|
||||
.send_text(text)
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn recv(&mut self) -> Option<Result<TransportMsg>> {
|
||||
if self.closed_flag.load(Ordering::SeqCst) {
|
||||
// Drain anything already reassembled, then report closed.
|
||||
return self.inbound.try_recv().ok().map(Ok);
|
||||
}
|
||||
tokio::select! {
|
||||
biased;
|
||||
msg = self.inbound.recv() => msg.map(Ok),
|
||||
_ = self.closed_rx.recv() => self.inbound.try_recv().ok().map(Ok),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Direct-IP transport — a raw TCP socket instead of the relay. The viewer dials
|
||||
/// the host's `IP:port` and the session runs the SAME challenge/auth + E2E
|
||||
/// AES-GCM above this layer (so direct-IP is encrypted, unlike RustDesk's).
|
||||
/// Framing is `[type:u8][len:u32 BE][payload]` (type 0 = binary, 1 = text). A
|
||||
/// dedicated reader task parses frames into a channel, so `recv` is
|
||||
/// cancellation-safe in the caller's `select!` (a bare `read_exact` is not).
|
||||
pub struct DirectTcp {
|
||||
writer: tokio::net::tcp::OwnedWriteHalf,
|
||||
rx: mpsc::UnboundedReceiver<Result<TransportMsg>>,
|
||||
}
|
||||
|
||||
/// Cap on a single framed message (generous for a 4K keyframe).
|
||||
const DIRECT_MAX_FRAME: usize = 64 * 1024 * 1024;
|
||||
|
||||
impl DirectTcp {
|
||||
/// Wrap an accepted / dialed TCP stream.
|
||||
pub fn new(stream: TcpStream) -> Self {
|
||||
let _ = stream.set_nodelay(true);
|
||||
let (mut read, write) = stream.into_split();
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
tokio::spawn(async move {
|
||||
use tokio::io::AsyncReadExt;
|
||||
let mut hdr = [0u8; 5];
|
||||
loop {
|
||||
if read.read_exact(&mut hdr).await.is_err() {
|
||||
break;
|
||||
}
|
||||
let ty = hdr[0];
|
||||
let len = u32::from_be_bytes([hdr[1], hdr[2], hdr[3], hdr[4]]) as usize;
|
||||
if len > DIRECT_MAX_FRAME {
|
||||
break;
|
||||
}
|
||||
let mut buf = vec![0u8; len];
|
||||
if read.read_exact(&mut buf).await.is_err() {
|
||||
break;
|
||||
}
|
||||
let msg = if ty == 1 {
|
||||
TransportMsg::Text(String::from_utf8_lossy(&buf).into_owned())
|
||||
} else {
|
||||
TransportMsg::Binary(buf)
|
||||
};
|
||||
if tx.send(Ok(msg)).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
DirectTcp { writer: write, rx }
|
||||
}
|
||||
|
||||
async fn write_frame(&mut self, ty: u8, payload: &[u8]) -> Result<()> {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let mut hdr = [0u8; 5];
|
||||
hdr[0] = ty;
|
||||
hdr[1..5].copy_from_slice(&(payload.len() as u32).to_be_bytes());
|
||||
self.writer
|
||||
.write_all(&hdr)
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))?;
|
||||
self.writer
|
||||
.write_all(payload)
|
||||
.await
|
||||
.map_err(|e| NdError::Connect(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for DirectTcp {
|
||||
async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
|
||||
self.write_frame(0, &data).await
|
||||
}
|
||||
async fn send_text(&mut self, text: String) -> Result<()> {
|
||||
self.write_frame(1, text.as_bytes()).await
|
||||
}
|
||||
async fn recv(&mut self) -> Option<Result<TransportMsg>> {
|
||||
self.rx.recv().await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod direct_tests {
|
||||
use super::*;
|
||||
|
||||
// DirectTcp framing must round-trip text + binary, in order, over a real
|
||||
// loopback socket (the foundation of encrypted direct-IP access).
|
||||
#[tokio::test]
|
||||
async fn direct_tcp_roundtrips_text_and_binary_in_order() {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let srv = tokio::spawn(async move {
|
||||
let (s, _) = listener.accept().await.unwrap();
|
||||
let mut t = DirectTcp::new(s);
|
||||
for _ in 0..3 {
|
||||
match t.recv().await.unwrap().unwrap() {
|
||||
TransportMsg::Text(s) => t.send_text(format!("echo:{s}")).await.unwrap(),
|
||||
TransportMsg::Binary(b) => {
|
||||
let mut r = b.clone();
|
||||
r.reverse();
|
||||
t.send_binary(r).await.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let cs = tokio::net::TcpStream::connect(addr).await.unwrap();
|
||||
let mut c = DirectTcp::new(cs);
|
||||
c.send_text("hello".into()).await.unwrap();
|
||||
c.send_binary(vec![1, 2, 3, 4]).await.unwrap();
|
||||
c.send_text("world".into()).await.unwrap();
|
||||
assert!(matches!(c.recv().await.unwrap().unwrap(), TransportMsg::Text(s) if s == "echo:hello"));
|
||||
assert!(matches!(c.recv().await.unwrap().unwrap(), TransportMsg::Binary(b) if b == vec![4, 3, 2, 1]));
|
||||
assert!(matches!(c.recv().await.unwrap().unwrap(), TransportMsg::Text(s) if s == "echo:world"));
|
||||
srv.await.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
//! Self-update support: a tiny release manifest the client fetches to learn the
|
||||
//! latest version, where to download it, and its SHA-256 — so an update is only
|
||||
//! ever applied after the bytes are verified. RustDesk-style sha256-checked
|
||||
//! auto-update, distilled to pure, testable helpers.
|
||||
//!
|
||||
//! The server publishes `version.json` alongside the installer (served by the
|
||||
//! existing `/download/<name>` static route — no server code change needed):
|
||||
//! ```json
|
||||
//! { "version": "0.2.0",
|
||||
//! "url": "https://netdesk.space/download/NetDesk-Setup.exe",
|
||||
//! "sha256": "<hex>", "notes": "what changed" }
|
||||
//! ```
|
||||
|
||||
use crate::error::{NdError, Result};
|
||||
use serde::Deserialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// A parsed release manifest.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct Manifest {
|
||||
pub version: String,
|
||||
pub url: String,
|
||||
#[serde(default)]
|
||||
pub sha256: String,
|
||||
#[serde(default)]
|
||||
pub notes: String,
|
||||
}
|
||||
|
||||
/// Parse the `version.json` manifest.
|
||||
pub fn parse_manifest(json: &str) -> Result<Manifest> {
|
||||
serde_json::from_str(json).map_err(|e| NdError::Other(anyhow::anyhow!("manifest: {e}")))
|
||||
}
|
||||
|
||||
/// Parse a dotted `x.y.z` version into a comparable tuple. Extra components are
|
||||
/// ignored; a missing/garbage component is treated as 0 (so a malformed version
|
||||
/// never reads as "newer" and triggers a spurious update).
|
||||
fn parse_ver(v: &str) -> (u64, u64, u64) {
|
||||
let mut it = v
|
||||
.trim()
|
||||
.trim_start_matches('v')
|
||||
.split(|c| c == '.' || c == '-' || c == '+')
|
||||
.map(|p| p.parse::<u64>().unwrap_or(0));
|
||||
(
|
||||
it.next().unwrap_or(0),
|
||||
it.next().unwrap_or(0),
|
||||
it.next().unwrap_or(0),
|
||||
)
|
||||
}
|
||||
|
||||
/// True when `latest` is strictly newer than `current` (semantic x.y.z compare).
|
||||
/// Equal or older → false, so re-checking never re-triggers.
|
||||
pub fn is_newer(current: &str, latest: &str) -> bool {
|
||||
parse_ver(latest) > parse_ver(current)
|
||||
}
|
||||
|
||||
/// Verify a downloaded artifact against the manifest's SHA-256 (hex, any case).
|
||||
/// An empty expected hash fails closed (we never install unverified bytes).
|
||||
pub fn verify_sha256(bytes: &[u8], expected_hex: &str) -> bool {
|
||||
if expected_hex.trim().is_empty() {
|
||||
return false;
|
||||
}
|
||||
let mut h = Sha256::new();
|
||||
h.update(bytes);
|
||||
let got = hex::encode(h.finalize());
|
||||
// length-independent constant-time-ish compare on the lowercased hex
|
||||
let exp = expected_hex.trim().to_lowercase();
|
||||
got.len() == exp.len()
|
||||
&& got
|
||||
.bytes()
|
||||
.zip(exp.bytes())
|
||||
.fold(0u8, |acc, (a, b)| acc | (a ^ b))
|
||||
== 0
|
||||
}
|
||||
|
||||
/// Decide whether to update given the running version + a fetched manifest.
|
||||
/// Returns Some(manifest) only when the manifest names a strictly newer version
|
||||
/// AND carries a sha256 (we refuse to "update" to an unverifiable artifact).
|
||||
pub fn update_available(current: &str, manifest: &Manifest) -> bool {
|
||||
!manifest.sha256.trim().is_empty() && is_newer(current, &manifest.version)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn version_compare() {
|
||||
assert!(is_newer("0.1.0", "0.2.0"));
|
||||
assert!(is_newer("0.1.0", "0.1.1"));
|
||||
assert!(is_newer("1.9.9", "2.0.0"));
|
||||
assert!(!is_newer("0.2.0", "0.2.0")); // equal
|
||||
assert!(!is_newer("0.2.0", "0.1.9")); // older
|
||||
assert!(!is_newer("0.2.0", "garbage")); // malformed -> not newer
|
||||
assert!(is_newer("v0.1.0", "v0.1.1")); // tolerate a 'v' prefix
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sha256_verify_roundtrip() {
|
||||
let data = b"NetDesk-Setup.exe bytes";
|
||||
let mut h = Sha256::new();
|
||||
h.update(data);
|
||||
let good = hex::encode(h.finalize());
|
||||
assert!(verify_sha256(data, &good));
|
||||
assert!(verify_sha256(data, &good.to_uppercase())); // case-insensitive
|
||||
assert!(!verify_sha256(data, "")); // empty fails closed
|
||||
assert!(!verify_sha256(data, &"0".repeat(64))); // wrong hash
|
||||
assert!(!verify_sha256(b"tampered", &good)); // tampered bytes
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_available_needs_newer_and_hash() {
|
||||
let m = |ver: &str, sha: &str| Manifest {
|
||||
version: ver.into(),
|
||||
url: "https://x/NetDesk-Setup.exe".into(),
|
||||
sha256: sha.into(),
|
||||
notes: String::new(),
|
||||
};
|
||||
assert!(update_available("0.1.0", &m("0.2.0", "abcd")));
|
||||
assert!(!update_available("0.1.0", &m("0.2.0", "")), "no hash -> refuse");
|
||||
assert!(!update_available("0.2.0", &m("0.2.0", "abcd")), "not newer");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_parses() {
|
||||
let m = parse_manifest(
|
||||
r#"{"version":"0.3.0","url":"https://netdesk.space/download/NetDesk-Setup.exe","sha256":"deadbeef","notes":"x"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(m.version, "0.3.0");
|
||||
assert!(m.url.ends_with("NetDesk-Setup.exe"));
|
||||
assert_eq!(m.sha256, "deadbeef");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,420 @@
|
||||
//! Pure-Rust WebM (Matroska/EBML) muxer for session recording — remuxes the
|
||||
//! received VP9 video frames + Opus audio packets into a playable `.webm` with
|
||||
//! **no re-encode** (just like `recorder::Mp4Recorder` does for H.264 → MP4).
|
||||
//!
|
||||
//! Why WebM and not MP4: our audio is Opus and the `mp4` crate is AAC-only (and
|
||||
//! we have no AAC encoder), so Opus can't ride an MP4 track. VP9 + Opus are both
|
||||
//! WebM-native, so a WebM container tees both received streams with zero
|
||||
//! transcoding. This is a hand-rolled EBML writer (no C++ libwebm / no mingw
|
||||
//! dependency): an EBML header, a streaming Segment (unknown size), Info,
|
||||
//! Tracks (V_VP9 + optional A_OPUS), then live Clusters of SimpleBlocks.
|
||||
//!
|
||||
//! Developed by Akronlogic.
|
||||
#![allow(clippy::unusual_byte_groupings)]
|
||||
|
||||
use crate::error::{NdError, Result};
|
||||
use std::fs::File;
|
||||
use std::io::{BufWriter, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn other(msg: impl std::fmt::Display) -> NdError {
|
||||
NdError::Other(anyhow::anyhow!("webm: {msg}"))
|
||||
}
|
||||
|
||||
// ---- EBML element IDs (written as their literal bytes; the ID already encodes
|
||||
// its own length in the leading bits) ----
|
||||
const ID_EBML: &[u8] = &[0x1A, 0x45, 0xDF, 0xA3];
|
||||
const ID_EBMLVERSION: &[u8] = &[0x42, 0x86];
|
||||
const ID_EBMLREADVERSION: &[u8] = &[0x42, 0xF7];
|
||||
const ID_EBMLMAXIDLENGTH: &[u8] = &[0x42, 0xF2];
|
||||
const ID_EBMLMAXSIZELENGTH: &[u8] = &[0x42, 0xF3];
|
||||
const ID_DOCTYPE: &[u8] = &[0x42, 0x82];
|
||||
const ID_DOCTYPEVERSION: &[u8] = &[0x42, 0x87];
|
||||
const ID_DOCTYPEREADVERSION: &[u8] = &[0x42, 0x85];
|
||||
const ID_SEGMENT: &[u8] = &[0x18, 0x53, 0x80, 0x67];
|
||||
const ID_INFO: &[u8] = &[0x15, 0x49, 0xA9, 0x66];
|
||||
const ID_TIMECODESCALE: &[u8] = &[0x2A, 0xD7, 0xB1];
|
||||
const ID_MUXINGAPP: &[u8] = &[0x4D, 0x80];
|
||||
const ID_WRITINGAPP: &[u8] = &[0x57, 0x41];
|
||||
const ID_TRACKS: &[u8] = &[0x16, 0x54, 0xAE, 0x6B];
|
||||
const ID_TRACKENTRY: &[u8] = &[0xAE];
|
||||
const ID_TRACKNUMBER: &[u8] = &[0xD7];
|
||||
const ID_TRACKUID: &[u8] = &[0x73, 0xC5];
|
||||
const ID_TRACKTYPE: &[u8] = &[0x83];
|
||||
const ID_CODECID: &[u8] = &[0x86];
|
||||
const ID_CODECPRIVATE: &[u8] = &[0x63, 0xA2];
|
||||
const ID_VIDEO: &[u8] = &[0xE0];
|
||||
const ID_PIXELWIDTH: &[u8] = &[0xB0];
|
||||
const ID_PIXELHEIGHT: &[u8] = &[0xBA];
|
||||
const ID_AUDIO: &[u8] = &[0xE1];
|
||||
const ID_SAMPLINGFREQ: &[u8] = &[0xB5];
|
||||
const ID_CHANNELS: &[u8] = &[0x9F];
|
||||
const ID_CLUSTER: &[u8] = &[0x1F, 0x43, 0xB6, 0x75];
|
||||
const ID_TIMECODE: &[u8] = &[0xE7];
|
||||
const ID_SIMPLEBLOCK: &[u8] = &[0xA3];
|
||||
|
||||
/// EBML variable-length integer encoding `value` as a size/number field.
|
||||
fn vint(value: u64) -> Vec<u8> {
|
||||
for n in 1..=8u32 {
|
||||
if value < (1u64 << (7 * n)) - 1 {
|
||||
let v = value | (1u64 << (7 * n));
|
||||
return v.to_be_bytes()[(8 - n as usize)..].to_vec();
|
||||
}
|
||||
}
|
||||
let v = (value & ((1u64 << 56) - 1)) | (1u64 << 56);
|
||||
v.to_be_bytes()[1..].to_vec()
|
||||
}
|
||||
|
||||
/// Minimal big-endian unsigned (EBML uint element body).
|
||||
fn uint_be(v: u64) -> Vec<u8> {
|
||||
if v == 0 {
|
||||
return vec![0u8];
|
||||
}
|
||||
let b = v.to_be_bytes();
|
||||
let first = b.iter().position(|&x| x != 0).unwrap();
|
||||
b[first..].to_vec()
|
||||
}
|
||||
|
||||
fn elem(id: &[u8], data: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(id.len() + 8 + data.len());
|
||||
out.extend_from_slice(id);
|
||||
out.extend_from_slice(&vint(data.len() as u64));
|
||||
out.extend_from_slice(data);
|
||||
out
|
||||
}
|
||||
|
||||
fn elem_uint(id: &[u8], v: u64) -> Vec<u8> {
|
||||
elem(id, &uint_be(v))
|
||||
}
|
||||
|
||||
fn ebml_header() -> Vec<u8> {
|
||||
let mut h = Vec::new();
|
||||
h.extend(elem_uint(ID_EBMLVERSION, 1));
|
||||
h.extend(elem_uint(ID_EBMLREADVERSION, 1));
|
||||
h.extend(elem_uint(ID_EBMLMAXIDLENGTH, 4));
|
||||
h.extend(elem_uint(ID_EBMLMAXSIZELENGTH, 8));
|
||||
h.extend(elem(ID_DOCTYPE, b"webm"));
|
||||
h.extend(elem_uint(ID_DOCTYPEVERSION, 2));
|
||||
h.extend(elem_uint(ID_DOCTYPEREADVERSION, 2));
|
||||
elem(ID_EBML, &h)
|
||||
}
|
||||
|
||||
fn info() -> Vec<u8> {
|
||||
let mut i = Vec::new();
|
||||
i.extend(elem_uint(ID_TIMECODESCALE, 1_000_000)); // 1 tick = 1 ms
|
||||
i.extend(elem(ID_MUXINGAPP, b"NetDesk"));
|
||||
i.extend(elem(ID_WRITINGAPP, b"NetDesk"));
|
||||
elem(ID_INFO, &i)
|
||||
}
|
||||
|
||||
fn opus_head(channels: u8) -> Vec<u8> {
|
||||
let mut h = Vec::new();
|
||||
h.extend_from_slice(b"OpusHead");
|
||||
h.push(1); // version
|
||||
h.push(channels.max(1));
|
||||
h.extend_from_slice(&3840u16.to_le_bytes()); // pre-skip
|
||||
h.extend_from_slice(&48000u32.to_le_bytes()); // input sample rate
|
||||
h.extend_from_slice(&0i16.to_le_bytes()); // output gain
|
||||
h.push(0); // channel-mapping family 0
|
||||
h
|
||||
}
|
||||
|
||||
fn tracks(width: u16, height: u16, channels: u8) -> Vec<u8> {
|
||||
let mut v = Vec::new();
|
||||
v.extend(elem_uint(ID_TRACKNUMBER, 1));
|
||||
v.extend(elem_uint(ID_TRACKUID, 1));
|
||||
v.extend(elem_uint(ID_TRACKTYPE, 1)); // video
|
||||
v.extend(elem(ID_CODECID, b"V_VP9"));
|
||||
let mut vid = Vec::new();
|
||||
vid.extend(elem_uint(ID_PIXELWIDTH, width as u64));
|
||||
vid.extend(elem_uint(ID_PIXELHEIGHT, height as u64));
|
||||
v.extend(elem(ID_VIDEO, &vid));
|
||||
let mut entries = elem(ID_TRACKENTRY, &v);
|
||||
|
||||
if channels > 0 {
|
||||
let mut a = Vec::new();
|
||||
a.extend(elem_uint(ID_TRACKNUMBER, 2));
|
||||
a.extend(elem_uint(ID_TRACKUID, 2));
|
||||
a.extend(elem_uint(ID_TRACKTYPE, 2)); // audio
|
||||
a.extend(elem(ID_CODECID, b"A_OPUS"));
|
||||
a.extend(elem(ID_CODECPRIVATE, &opus_head(channels)));
|
||||
let mut aud = Vec::new();
|
||||
aud.extend(elem(ID_SAMPLINGFREQ, &48000.0f64.to_be_bytes()));
|
||||
aud.extend(elem_uint(ID_CHANNELS, channels as u64));
|
||||
a.extend(elem(ID_AUDIO, &aud));
|
||||
entries.extend(elem(ID_TRACKENTRY, &a));
|
||||
}
|
||||
elem(ID_TRACKS, &entries)
|
||||
}
|
||||
|
||||
fn simpleblock(track: u64, rel_tc: i16, keyframe: bool, frame: &[u8]) -> Vec<u8> {
|
||||
let mut d = Vec::with_capacity(4 + frame.len());
|
||||
d.extend_from_slice(&vint(track)); // track number (EBML vint, 1-based)
|
||||
d.extend_from_slice(&rel_tc.to_be_bytes()); // signed 16-bit relative timecode
|
||||
d.push(if keyframe { 0x80 } else { 0x00 }); // flags: 0x80 = keyframe
|
||||
d.extend_from_slice(frame);
|
||||
elem(ID_SIMPLEBLOCK, &d)
|
||||
}
|
||||
|
||||
/// Muxes received VP9 video (+ optional Opus audio) into a `.webm`. No decode /
|
||||
/// re-encode: the frames are stored exactly as they arrive from the session.
|
||||
pub struct WebmRecorder {
|
||||
path: PathBuf,
|
||||
width: u16,
|
||||
height: u16,
|
||||
channels: u8, // 0 = video-only (no audio track)
|
||||
writer: Option<BufWriter<File>>,
|
||||
start_ms: Option<u64>,
|
||||
cluster_base: u64,
|
||||
cluster: Vec<u8>,
|
||||
open_cluster: bool,
|
||||
frames: u64,
|
||||
}
|
||||
|
||||
impl WebmRecorder {
|
||||
/// `channels` 0 records video only; >0 adds an Opus audio track.
|
||||
pub fn new(path: impl AsRef<Path>, width: u16, height: u16, channels: u8) -> Self {
|
||||
Self {
|
||||
path: path.as_ref().to_path_buf(),
|
||||
width: width.max(2),
|
||||
height: height.max(2),
|
||||
channels,
|
||||
writer: None,
|
||||
start_ms: None,
|
||||
cluster_base: 0,
|
||||
cluster: Vec::new(),
|
||||
open_cluster: false,
|
||||
frames: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_open(&mut self) -> Result<()> {
|
||||
if self.writer.is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
let mut w = BufWriter::new(File::create(&self.path).map_err(other)?);
|
||||
w.write_all(&ebml_header()).map_err(other)?;
|
||||
// Streaming Segment: unknown size (all data bits set in an 8-octet vint).
|
||||
w.write_all(ID_SEGMENT).map_err(other)?;
|
||||
w.write_all(&[0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]).map_err(other)?;
|
||||
w.write_all(&info()).map_err(other)?;
|
||||
w.write_all(&tracks(self.width, self.height, self.channels)).map_err(other)?;
|
||||
self.writer = Some(w);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush_cluster(&mut self) -> Result<()> {
|
||||
if self.open_cluster && !self.cluster.is_empty() {
|
||||
if let Some(w) = self.writer.as_mut() {
|
||||
let mut body = Vec::with_capacity(self.cluster.len() + 8);
|
||||
body.extend(elem_uint(ID_TIMECODE, self.cluster_base));
|
||||
body.extend_from_slice(&self.cluster);
|
||||
w.write_all(ID_CLUSTER).map_err(other)?;
|
||||
w.write_all(&vint(body.len() as u64)).map_err(other)?;
|
||||
w.write_all(&body).map_err(other)?;
|
||||
}
|
||||
}
|
||||
self.cluster.clear();
|
||||
self.open_cluster = false;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn add_block(&mut self, track: u64, data: &[u8], keyframe: bool, ts_ms: u64) -> Result<()> {
|
||||
let start = *self.start_ms.get_or_insert(ts_ms);
|
||||
let rel_total = ts_ms.saturating_sub(start);
|
||||
// New cluster on a video keyframe (so seeking works) or before the
|
||||
// relative timecode could overflow the signed-16-bit SimpleBlock field.
|
||||
if !self.open_cluster
|
||||
|| (track == 1 && keyframe)
|
||||
|| rel_total.saturating_sub(self.cluster_base) > 30_000
|
||||
{
|
||||
self.flush_cluster()?;
|
||||
self.cluster_base = rel_total;
|
||||
self.open_cluster = true;
|
||||
}
|
||||
let rel = rel_total.saturating_sub(self.cluster_base).min(32_767) as i16;
|
||||
self.cluster.extend(simpleblock(track, rel, keyframe, data));
|
||||
self.frames += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Feed one VP9 frame. The file opens on the first keyframe.
|
||||
pub fn push_video(&mut self, data: &[u8], keyframe: bool, ts_ms: u64) -> Result<()> {
|
||||
if self.writer.is_none() {
|
||||
if !keyframe {
|
||||
return Ok(()); // wait for the first keyframe to start the file
|
||||
}
|
||||
self.ensure_open()?;
|
||||
}
|
||||
if data.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
self.add_block(1, data, keyframe, ts_ms)
|
||||
}
|
||||
|
||||
/// Feed one Opus packet (ignored until the video track + file exist, and
|
||||
/// when the recorder was created video-only).
|
||||
pub fn push_audio(&mut self, data: &[u8], ts_ms: u64) -> Result<()> {
|
||||
if self.writer.is_none() || self.channels == 0 || data.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
self.add_block(2, data, true, ts_ms) // every Opus frame is independent
|
||||
}
|
||||
|
||||
/// Flush the final cluster and close the file. Returns (path, block count).
|
||||
pub fn finish(mut self) -> Result<(PathBuf, u64)> {
|
||||
self.flush_cluster()?;
|
||||
if let Some(mut w) = self.writer.take() {
|
||||
w.flush().map_err(other)?;
|
||||
}
|
||||
Ok((self.path.clone(), self.frames))
|
||||
}
|
||||
|
||||
pub fn frames(&self) -> u64 {
|
||||
self.frames
|
||||
}
|
||||
}
|
||||
|
||||
// ---- verification: a minimal EBML walk that re-reads a muxed file ----
|
||||
fn find_sub(hay: &[u8], needle: &[u8]) -> bool {
|
||||
hay.windows(needle.len()).any(|w| w == needle)
|
||||
}
|
||||
|
||||
fn read_id(d: &[u8], p: &mut usize) -> Option<Vec<u8>> {
|
||||
let b0 = *d.get(*p)?;
|
||||
let n = (0..4).find(|i| b0 & (0x80 >> i) != 0).map(|i| i + 1)?;
|
||||
if *p + n > d.len() {
|
||||
return None;
|
||||
}
|
||||
let id = d[*p..*p + n].to_vec();
|
||||
*p += n;
|
||||
Some(id)
|
||||
}
|
||||
|
||||
fn read_size(d: &[u8], p: &mut usize) -> Option<u64> {
|
||||
let b0 = *d.get(*p)?;
|
||||
let n = (0..8).find(|i| b0 & (0x80 >> i) != 0).map(|i| i + 1)?;
|
||||
if *p + n > d.len() {
|
||||
return None;
|
||||
}
|
||||
let mut v = (b0 & (0xFF >> n)) as u64;
|
||||
for i in 1..n {
|
||||
v = (v << 8) | d[*p + i] as u64;
|
||||
}
|
||||
*p += n;
|
||||
Some(v)
|
||||
}
|
||||
|
||||
/// Re-read a `.webm` and confirm its structure: returns
|
||||
/// (is_webm_doctype, has_vp9, has_opus, simpleblock_count). Walks the top-level
|
||||
/// Segment children and descends into each Cluster to count SimpleBlocks for
|
||||
/// real (no byte-scan false positives for the block count).
|
||||
pub fn probe_webm(path: impl AsRef<Path>) -> Result<(bool, bool, bool, u64)> {
|
||||
let data = std::fs::read(path.as_ref()).map_err(other)?;
|
||||
if data.len() < 4 || &data[0..4] != ID_EBML {
|
||||
return Err(other("missing EBML magic"));
|
||||
}
|
||||
let webm = find_sub(&data, b"webm");
|
||||
let mut p = 0usize;
|
||||
// EBML header element (skip).
|
||||
let _ = read_id(&data, &mut p).ok_or_else(|| other("no EBML id"))?;
|
||||
let hsz = read_size(&data, &mut p).ok_or_else(|| other("no EBML size"))? as usize;
|
||||
p += hsz;
|
||||
// Segment (unknown/huge size) — walk its children to EOF.
|
||||
let _seg = read_id(&data, &mut p).ok_or_else(|| other("no Segment id"))?;
|
||||
let _ = read_size(&data, &mut p).ok_or_else(|| other("no Segment size"))?;
|
||||
let (mut vp9, mut opus, mut blocks) = (false, false, 0u64);
|
||||
while p < data.len() {
|
||||
let id = match read_id(&data, &mut p) {
|
||||
Some(i) => i,
|
||||
None => break,
|
||||
};
|
||||
let sz = match read_size(&data, &mut p) {
|
||||
Some(s) => s as usize,
|
||||
None => break,
|
||||
};
|
||||
let end = (p + sz).min(data.len());
|
||||
if id == ID_CLUSTER {
|
||||
let mut q = p;
|
||||
while q < end {
|
||||
let cid = match read_id(&data, &mut q) {
|
||||
Some(i) => i,
|
||||
None => break,
|
||||
};
|
||||
let csz = match read_size(&data, &mut q) {
|
||||
Some(s) => s as usize,
|
||||
None => break,
|
||||
};
|
||||
if cid == ID_SIMPLEBLOCK {
|
||||
blocks += 1;
|
||||
}
|
||||
q += csz;
|
||||
}
|
||||
} else if id == ID_TRACKS {
|
||||
let seg = &data[p..end];
|
||||
vp9 = find_sub(seg, b"V_VP9");
|
||||
opus = find_sub(seg, b"A_OPUS");
|
||||
}
|
||||
p = end;
|
||||
}
|
||||
Ok((webm, vp9, opus, blocks))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn mux_vp9_opus_webm_roundtrips() {
|
||||
// Codec-agnostic remux: the payload bytes are arbitrary (the muxer does
|
||||
// not parse VP9/Opus — the live session supplies real frames). We prove
|
||||
// the CONTAINER: a valid WebM with a V_VP9 + A_OPUS track and the fed
|
||||
// frames stored as SimpleBlocks, re-read by probe_webm.
|
||||
let dir = std::env::temp_dir();
|
||||
let path = dir.join("netdesk_webm_mux_test.webm");
|
||||
let _ = std::fs::remove_file(&path);
|
||||
|
||||
let mut rec = WebmRecorder::new(&path, 320, 240, 2);
|
||||
// audio before the first keyframe is dropped (file not open yet)
|
||||
rec.push_audio(&[0xAA; 40], 0).unwrap();
|
||||
// first keyframe opens the file
|
||||
rec.push_video(&[0x11; 300], true, 0).unwrap();
|
||||
rec.push_audio(&[0xBB; 40], 0).unwrap();
|
||||
rec.push_video(&[0x22; 120], false, 33).unwrap();
|
||||
rec.push_audio(&[0xCC; 40], 20).unwrap();
|
||||
rec.push_video(&[0x33; 110], false, 66).unwrap();
|
||||
rec.push_video(&[0x44; 280], true, 100).unwrap(); // 2nd keyframe => new cluster
|
||||
rec.push_audio(&[0xDD; 40], 100).unwrap();
|
||||
let (out, frames) = rec.finish().unwrap();
|
||||
|
||||
assert!(out.exists(), "webm file not written");
|
||||
let sz = std::fs::metadata(&out).unwrap().len();
|
||||
let (is_webm, has_vp9, has_opus, blocks) = probe_webm(&out).unwrap();
|
||||
println!(
|
||||
"WEBM_OK file={} bytes={} frames={} vp9={} opus={} simpleblocks={}",
|
||||
out.display(), sz, frames, has_vp9, has_opus, blocks
|
||||
);
|
||||
assert!(is_webm, "DocType is not webm");
|
||||
assert!(has_vp9, "no V_VP9 track");
|
||||
assert!(has_opus, "no A_OPUS track");
|
||||
// 4 video + 4 audio pushed, but the 1st audio (pre-keyframe) was dropped.
|
||||
assert_eq!(blocks, 7, "expected 7 SimpleBlocks, got {blocks}");
|
||||
assert_eq!(frames, 7, "recorder frame count");
|
||||
let _ = std::fs::remove_file(&out);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn video_only_webm_has_no_audio_track() {
|
||||
let path = std::env::temp_dir().join("netdesk_webm_vid_only.webm");
|
||||
let _ = std::fs::remove_file(&path);
|
||||
let mut rec = WebmRecorder::new(&path, 64, 64, 0);
|
||||
rec.push_audio(&[0x01; 20], 0).unwrap(); // ignored (video-only)
|
||||
rec.push_video(&[0x11; 50], true, 0).unwrap();
|
||||
rec.push_video(&[0x22; 40], false, 33).unwrap();
|
||||
let (out, frames) = rec.finish().unwrap();
|
||||
let (is_webm, has_vp9, has_opus, blocks) = probe_webm(&out).unwrap();
|
||||
assert!(is_webm && has_vp9 && !has_opus, "video-only must have VP9 and no Opus");
|
||||
assert_eq!((frames, blocks), (2, 2));
|
||||
let _ = std::fs::remove_file(&out);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
//! End-to-end: the core HOST (host.rs) serves a session that the core VIEWER
|
||||
//! (signaling.rs connect) authenticates to, over the LIVE server. Proves the
|
||||
//! Android host's protocol (register -> challenge -> auth/E2E -> accept) and
|
||||
//! input routing interoperate before any Kotlin is wired.
|
||||
//!
|
||||
//! Network test -> #[ignore] by default. Run explicitly:
|
||||
//! cargo test -p netdesk-core --target x86_64-pc-windows-gnu --test host_viewer -- --ignored --nocapture
|
||||
|
||||
use netdesk_core::crypto::{gen_host_id, new_host_credentials};
|
||||
use netdesk_core::host::{run_host, HostConfig, HostEvent, HostFrame};
|
||||
use netdesk_core::{connect, ConnectConfig, SessionEvent};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::timeout;
|
||||
|
||||
async fn wait_host<F: Fn(&HostEvent) -> bool>(
|
||||
rx: &mut mpsc::UnboundedReceiver<HostEvent>,
|
||||
pred: F,
|
||||
secs: u64,
|
||||
) -> Option<HostEvent> {
|
||||
let deadline = Duration::from_secs(secs);
|
||||
loop {
|
||||
match timeout(deadline, rx.recv()).await {
|
||||
Ok(Some(ev)) => {
|
||||
if pred(&ev) {
|
||||
return Some(ev);
|
||||
}
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_view<F: Fn(&SessionEvent) -> bool>(
|
||||
rx: &mut mpsc::UnboundedReceiver<SessionEvent>,
|
||||
pred: F,
|
||||
secs: u64,
|
||||
) -> Option<SessionEvent> {
|
||||
let deadline = Duration::from_secs(secs);
|
||||
loop {
|
||||
match timeout(deadline, rx.recv()).await {
|
||||
Ok(Some(ev)) => {
|
||||
if pred(&ev) {
|
||||
return Some(ev);
|
||||
}
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn host_viewer_handshake_and_input() {
|
||||
let ws = "wss://netdesk.space/ws";
|
||||
let id = gen_host_id();
|
||||
let password = "e2e-host-test-pw";
|
||||
let (salt_b64, iters, verifier_hex) = new_host_credentials(password);
|
||||
|
||||
let (frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let hcfg = HostConfig {
|
||||
ws_url: ws.into(),
|
||||
host_id: id.clone(),
|
||||
name: "TestAndroidHost".into(),
|
||||
salt_b64,
|
||||
iters,
|
||||
verifier_hex,
|
||||
width: 1080,
|
||||
height: 1920,
|
||||
audio: false,
|
||||
monitors: Vec::new(),
|
||||
monitor: 0,
|
||||
perms: Default::default(),
|
||||
approve: false,
|
||||
totp_secret: String::new(),
|
||||
};
|
||||
// Outbound control channel (unused by this test; drop the sender).
|
||||
let (_hctl_tx, hctl_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let host_task =
|
||||
tokio::spawn(async move { run_host(hcfg, frame_rx, hev_tx, hctl_rx).await });
|
||||
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Registered), 10)
|
||||
.await
|
||||
.is_some(),
|
||||
"host did not register"
|
||||
);
|
||||
// let the registration settle on the server before the viewer dials in
|
||||
tokio::time::sleep(Duration::from_millis(600)).await;
|
||||
|
||||
// viewer connects to our host id
|
||||
let (vtx, mut vrx) = mpsc::unbounded_channel::<SessionEvent>();
|
||||
let vcfg = ConnectConfig {
|
||||
ws_url: ws.into(),
|
||||
host_id: id.clone(),
|
||||
password: password.into(),
|
||||
totp: String::new(),
|
||||
takeover: false,
|
||||
};
|
||||
let session = connect(vcfg, vtx).await.expect("viewer connect+handshake");
|
||||
|
||||
// host sees the viewer; viewer sees Connected with our advertised size
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::ViewerConnected { .. }), 10)
|
||||
.await
|
||||
.is_some(),
|
||||
"host never saw ViewerConnected"
|
||||
);
|
||||
let connected = wait_view(&mut vrx, |e| matches!(e, SessionEvent::Connected { .. }), 10).await;
|
||||
match connected {
|
||||
Some(SessionEvent::Connected { w, h, e2e, .. }) => {
|
||||
assert_eq!((w, h), (1080, 1920), "advertised resolution");
|
||||
assert!(e2e, "E2E negotiated");
|
||||
}
|
||||
_ => panic!("viewer never got Connected"),
|
||||
}
|
||||
|
||||
// viewer input must surface on the host side with matching coordinates
|
||||
session.mouse(123, 456, 1, 0.0, 0.0);
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Mouse { .. }), 10).await {
|
||||
Some(HostEvent::Mouse { x, y, buttons, .. }) => {
|
||||
assert_eq!((x, y, buttons), (123, 456, 1), "mouse routed to host");
|
||||
}
|
||||
_ => panic!("host never received the viewer's mouse"),
|
||||
}
|
||||
|
||||
session.key("KeyA", true);
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Key { .. }), 10).await {
|
||||
Some(HostEvent::Key { code, down }) => {
|
||||
assert_eq!((code.as_str(), down), ("KeyA", true), "key routed to host");
|
||||
}
|
||||
_ => panic!("host never received the viewer's key"),
|
||||
}
|
||||
|
||||
// a frame fed to the host reaches the viewer as a binary message (we only
|
||||
// assert the host accepts + seals it; decode is covered elsewhere).
|
||||
let _ = frame_tx.send(HostFrame {
|
||||
data: vec![0, 0, 0, 1, 0x67, 0x42, 0, 0x28],
|
||||
keyframe: true,
|
||||
codec: 0,
|
||||
});
|
||||
tokio::time::sleep(Duration::from_millis(300)).await;
|
||||
|
||||
session.close();
|
||||
drop(frame_tx);
|
||||
let _ = timeout(Duration::from_secs(3), host_task).await;
|
||||
println!("HOST_VIEWER_E2E_OK id={id}");
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,739 @@
|
||||
//! Host-enforced permission ACL + view-only + approve-before-connect — protocol
|
||||
//! tests over the in-memory relay (same harness as p2p_loopback.rs): the FULL
|
||||
//! `connect_with` / `run_host_with` stacks run server-free, and the assertions
|
||||
//! are on what the ENGINE emits/sends — i.e. what a (possibly modified) viewer
|
||||
//! could actually make the host do, not what a well-behaved UI chooses to send.
|
||||
//!
|
||||
//! Proves the Tier-0 security fix: with a permission disabled, the engine never
|
||||
//! emits the corresponding platform event (input/clipboard/file/audio), replies
|
||||
//! `file.err` to file requests, advertises the REAL perms in `session.accept`,
|
||||
//! enforces viewer-requested view-only host-side, and holds a password-passed
|
||||
//! session for Accept/Reject when approve mode is on.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use netdesk_core::crypto::{
|
||||
gen_host_id, new_host_credentials, password_proof, password_proof_bound,
|
||||
};
|
||||
use netdesk_core::host::{run_host_with, HostConfig, HostControl, HostEvent, HostFrame, HostPerms};
|
||||
use netdesk_core::signaling::connect_with;
|
||||
use netdesk_core::{ConnectConfig, NdError, SessionEvent, Transport, TransportMsg};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// In-memory relay pair (the p2p_loopback.rs harness, minus the drop toggles).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct MemTransport {
|
||||
tx: mpsc::UnboundedSender<TransportMsg>,
|
||||
rx: mpsc::UnboundedReceiver<TransportMsg>,
|
||||
}
|
||||
|
||||
fn mem_relay_pair() -> (MemTransport, MemTransport) {
|
||||
let (a_tx, a_rx) = mpsc::unbounded_channel::<TransportMsg>();
|
||||
let (b_tx, b_rx) = mpsc::unbounded_channel::<TransportMsg>();
|
||||
(MemTransport { tx: b_tx, rx: a_rx }, MemTransport { tx: a_tx, rx: b_rx })
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for MemTransport {
|
||||
async fn send_binary(&mut self, data: Vec<u8>) -> netdesk_core::Result<()> {
|
||||
self.tx
|
||||
.send(TransportMsg::Binary(data))
|
||||
.map_err(|_| NdError::Connect("mem relay peer gone".into()))
|
||||
}
|
||||
async fn send_text(&mut self, text: String) -> netdesk_core::Result<()> {
|
||||
self.tx
|
||||
.send(TransportMsg::Text(text))
|
||||
.map_err(|_| NdError::Connect("mem relay peer gone".into()))
|
||||
}
|
||||
async fn recv(&mut self) -> Option<netdesk_core::Result<TransportMsg>> {
|
||||
self.rx.recv().await.map(Ok)
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_host<F: Fn(&HostEvent) -> bool>(
|
||||
rx: &mut mpsc::UnboundedReceiver<HostEvent>,
|
||||
pred: F,
|
||||
secs: u64,
|
||||
) -> Option<HostEvent> {
|
||||
let deadline = Duration::from_secs(secs);
|
||||
loop {
|
||||
match timeout(deadline, rx.recv()).await {
|
||||
Ok(Some(ev)) => {
|
||||
if pred(&ev) {
|
||||
return Some(ev);
|
||||
}
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_view<F: Fn(&SessionEvent) -> bool>(
|
||||
rx: &mut mpsc::UnboundedReceiver<SessionEvent>,
|
||||
pred: F,
|
||||
secs: u64,
|
||||
) -> Option<SessionEvent> {
|
||||
let deadline = Duration::from_secs(secs);
|
||||
loop {
|
||||
match timeout(deadline, rx.recv()).await {
|
||||
Ok(Some(ev)) => {
|
||||
if pred(&ev) {
|
||||
return Some(ev);
|
||||
}
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn host_cfg(perms: HostPerms, approve: bool) -> (HostConfig, String) {
|
||||
host_cfg_totp(perms, approve, "")
|
||||
}
|
||||
|
||||
fn host_cfg_totp(perms: HostPerms, approve: bool, totp_secret: &str) -> (HostConfig, String) {
|
||||
let password = "perms-test-pw".to_string();
|
||||
let (salt_b64, iters, verifier_hex) = new_host_credentials(&password);
|
||||
let cfg = HostConfig {
|
||||
ws_url: "mem://relay".into(),
|
||||
host_id: gen_host_id(),
|
||||
name: "PermsHost".into(),
|
||||
salt_b64,
|
||||
iters,
|
||||
verifier_hex,
|
||||
width: 1280,
|
||||
height: 720,
|
||||
audio: true,
|
||||
monitors: Vec::new(),
|
||||
monitor: 0,
|
||||
perms,
|
||||
approve,
|
||||
totp_secret: totp_secret.to_string(),
|
||||
};
|
||||
(cfg, password)
|
||||
}
|
||||
|
||||
/// Collect host events until `marker` matches (inclusive), with a timeout.
|
||||
/// Returns everything seen before the marker.
|
||||
async fn collect_until_marker<F: Fn(&HostEvent) -> bool>(
|
||||
rx: &mut mpsc::UnboundedReceiver<HostEvent>,
|
||||
marker: F,
|
||||
secs: u64,
|
||||
) -> Vec<HostEvent> {
|
||||
let mut seen = Vec::new();
|
||||
let deadline = Duration::from_secs(secs);
|
||||
loop {
|
||||
match timeout(deadline, rx.recv()).await {
|
||||
Ok(Some(ev)) => {
|
||||
if marker(&ev) {
|
||||
return seen;
|
||||
}
|
||||
seen.push(ev);
|
||||
}
|
||||
_ => panic!("marker event never arrived; saw {seen:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 1) All perms DISABLED: the engine must not emit input/clipboard/file/audio
|
||||
// events no matter what the viewer sends, and file requests get file.err.
|
||||
// ---------------------------------------------------------------------------
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn perms_disabled_block_all_channels() {
|
||||
let (viewer_relay, host_relay) = mem_relay_pair();
|
||||
let perms = HostPerms { input: false, clipboard: false, files: false, audio: false };
|
||||
let (hcfg, password) = host_cfg(perms, false);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
|
||||
let (frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
let host_task = tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let (vtx, mut vrx) = mpsc::unbounded_channel::<SessionEvent>();
|
||||
let vcfg = ConnectConfig {
|
||||
ws_url: "mem://relay".into(),
|
||||
host_id,
|
||||
password,
|
||||
totp: String::new(),
|
||||
takeover: false,
|
||||
};
|
||||
let session = connect_with(vcfg, vtx, Box::new(viewer_relay))
|
||||
.await
|
||||
.expect("viewer connects (password auth still works with perms off)");
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::ViewerConnected { .. }), 10)
|
||||
.await
|
||||
.is_some(),
|
||||
"host saw the viewer"
|
||||
);
|
||||
|
||||
// Fire one of EVERYTHING a hostile viewer could send on disabled channels…
|
||||
session.mouse(100, 100, 1, 0.0, 0.0);
|
||||
session.key("F15", true);
|
||||
session.key("F15", false);
|
||||
session.clipboard("must-not-reach-the-platform");
|
||||
session.audio_set(1);
|
||||
session.list_roots();
|
||||
// …then flip view-only as an ORDERED MARKER (its handler is independent of
|
||||
// the four perms, so it must still emit — everything sent before it was
|
||||
// processed first on the same ordered channel).
|
||||
session.view_only(true);
|
||||
|
||||
let before = collect_until_marker(
|
||||
&mut hev_rx,
|
||||
|e| matches!(e, HostEvent::ViewOnly { on: true }),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
for ev in &before {
|
||||
assert!(
|
||||
!matches!(
|
||||
ev,
|
||||
HostEvent::Mouse { .. }
|
||||
| HostEvent::Key { .. }
|
||||
| HostEvent::Clipboard { .. }
|
||||
| HostEvent::File { .. }
|
||||
| HostEvent::FileChunk { .. }
|
||||
| HostEvent::AudioSource { .. }
|
||||
| HostEvent::Voice { .. }
|
||||
),
|
||||
"engine leaked a blocked event to the platform: {ev:?}"
|
||||
);
|
||||
}
|
||||
|
||||
// The denied file request must answer with a clear file.err.
|
||||
match wait_view(&mut vrx, |e| matches!(e, SessionEvent::FileError { .. }), 10).await {
|
||||
Some(SessionEvent::FileError { msg }) => {
|
||||
assert!(
|
||||
msg.contains("disabled"),
|
||||
"file.err should say the channel is disabled, got: {msg}"
|
||||
);
|
||||
}
|
||||
_ => panic!("viewer never got file.err for the denied file.roots"),
|
||||
}
|
||||
|
||||
session.close();
|
||||
drop(frame_tx);
|
||||
let _ = timeout(Duration::from_secs(5), host_task).await;
|
||||
println!("PERMS_DENY_OK all four channels blocked at the engine + file.err reply");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 2) Default perms: every channel flows (no regression from the gates).
|
||||
// ---------------------------------------------------------------------------
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn perms_default_channels_flow() {
|
||||
let (viewer_relay, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg(HostPerms::default(), false);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
|
||||
let (frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
let host_task = tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let (vtx, _vrx) = mpsc::unbounded_channel::<SessionEvent>();
|
||||
let vcfg = ConnectConfig {
|
||||
ws_url: "mem://relay".into(),
|
||||
host_id,
|
||||
password,
|
||||
totp: String::new(),
|
||||
takeover: false,
|
||||
};
|
||||
let session = connect_with(vcfg, vtx, Box::new(viewer_relay)).await.expect("connect");
|
||||
|
||||
session.mouse(321, 654, 2, 0.0, 0.0);
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Mouse { .. }), 10).await {
|
||||
Some(HostEvent::Mouse { x, y, buttons, .. }) => assert_eq!((x, y, buttons), (321, 654, 2)),
|
||||
_ => panic!("mouse blocked despite perms on"),
|
||||
}
|
||||
// Relative (game-mouse) motion flows through the same input gate.
|
||||
session.mouse_rel(17, -9, 1, 0.0);
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::MouseRel { .. }), 10).await {
|
||||
Some(HostEvent::MouseRel { dx, dy, buttons, .. }) => {
|
||||
assert_eq!((dx, dy, buttons), (17, -9, 1))
|
||||
}
|
||||
_ => panic!("relative mouse blocked despite perms on"),
|
||||
}
|
||||
// Translate-keyboard Unicode text flows through the same input gate.
|
||||
session.text("h\u{00e9}llo");
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Text { .. }), 10).await {
|
||||
Some(HostEvent::Text { text }) => assert_eq!(text, "h\u{00e9}llo"),
|
||||
_ => panic!("text input blocked despite perms on"),
|
||||
}
|
||||
// Annotation strokes flow through the same input gate and reach the host.
|
||||
session.annotate_stroke(0xFF_0000, 3, &[(10, 20), (30, 40)]);
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Annot { .. }), 10).await {
|
||||
Some(HostEvent::Annot { json }) => assert!(json.contains("stroke"), "annot json: {json}"),
|
||||
_ => panic!("annotation blocked despite perms on"),
|
||||
}
|
||||
session.key("KeyA", true);
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Key { down: true, .. }), 10)
|
||||
.await
|
||||
.is_some(),
|
||||
"key blocked despite perms on"
|
||||
);
|
||||
session.clipboard("hello");
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Clipboard { .. }), 10).await {
|
||||
Some(HostEvent::Clipboard { text }) => assert_eq!(text, "hello"),
|
||||
_ => panic!("clipboard blocked despite perms on"),
|
||||
}
|
||||
session.audio_set(1);
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::AudioSource { .. }), 10).await {
|
||||
Some(HostEvent::AudioSource { src }) => assert_eq!(src, 1),
|
||||
_ => panic!("audio.set blocked despite perms on"),
|
||||
}
|
||||
session.list_roots();
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::File { .. }), 10)
|
||||
.await
|
||||
.is_some(),
|
||||
"file.roots blocked despite perms on"
|
||||
);
|
||||
|
||||
session.close();
|
||||
drop(frame_tx);
|
||||
let _ = timeout(Duration::from_secs(5), host_task).await;
|
||||
println!("PERMS_ALLOW_OK all channels flow with default perms");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 3) View-only: input gated OFF and back ON mid-session, enforced host-side.
|
||||
// ---------------------------------------------------------------------------
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn view_only_gates_input_live() {
|
||||
let (viewer_relay, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg(HostPerms::default(), false);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
|
||||
let (frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
let host_task = tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let (vtx, _vrx) = mpsc::unbounded_channel::<SessionEvent>();
|
||||
let vcfg = ConnectConfig {
|
||||
ws_url: "mem://relay".into(),
|
||||
host_id,
|
||||
password,
|
||||
totp: String::new(),
|
||||
takeover: false,
|
||||
};
|
||||
let session = connect_with(vcfg, vtx, Box::new(viewer_relay)).await.expect("connect");
|
||||
|
||||
// Input flows before view-only…
|
||||
session.mouse(10, 10, 0, 0.0, 0.0);
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Mouse { .. }), 10).await.is_some(),
|
||||
"input should flow before view-only"
|
||||
);
|
||||
|
||||
// …is dropped while it is on…
|
||||
session.view_only(true);
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::ViewOnly { on: true }), 10)
|
||||
.await
|
||||
.is_some()
|
||||
);
|
||||
session.mouse(20, 20, 1, 0.0, 0.0);
|
||||
session.key("KeyB", true);
|
||||
session.view_only(false); // ordered marker: everything above processed first
|
||||
let between = collect_until_marker(
|
||||
&mut hev_rx,
|
||||
|e| matches!(e, HostEvent::ViewOnly { on: false }),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
for ev in &between {
|
||||
assert!(
|
||||
!matches!(ev, HostEvent::Mouse { .. } | HostEvent::Key { .. }),
|
||||
"input leaked through host-enforced view-only: {ev:?}"
|
||||
);
|
||||
}
|
||||
|
||||
// …and flows again after it is off.
|
||||
session.mouse(30, 30, 0, 0.0, 0.0);
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::Mouse { .. }), 10).await.is_some(),
|
||||
"input should flow again after view-only off"
|
||||
);
|
||||
|
||||
session.close();
|
||||
drop(frame_tx);
|
||||
let _ = timeout(Duration::from_secs(5), host_task).await;
|
||||
println!("VIEW_ONLY_OK host-enforced input gate on/off mid-session");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 4) Wire-level: session.accept advertises the REAL perms; approve mode holds
|
||||
// the session and Accept/Reject finalizes/denies it. Drives the raw
|
||||
// protocol by hand so the actual JSON on the wire is asserted.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A parsed session.challenge from the host.
|
||||
struct Challenge {
|
||||
sid: String,
|
||||
salt: String,
|
||||
iters: u32,
|
||||
nonce: String,
|
||||
pub_b64: String,
|
||||
totp: bool,
|
||||
}
|
||||
|
||||
/// Drive begin → challenge on a raw viewer-end transport; return the challenge.
|
||||
async fn raw_challenge(viewer: &mut MemTransport, host_id: &str) -> Challenge {
|
||||
viewer
|
||||
.send_text(format!(r#"{{"t":"session.begin","id":"{host_id}","peer":"203.0.113.7","user":"tester@example.com"}}"#))
|
||||
.await
|
||||
.unwrap();
|
||||
loop {
|
||||
match timeout(Duration::from_secs(10), viewer.recv()).await {
|
||||
Ok(Some(Ok(TransportMsg::Text(t)))) => {
|
||||
let v: Value = serde_json::from_str(&t).unwrap();
|
||||
if v["t"] == "session.challenge" {
|
||||
return Challenge {
|
||||
sid: v["sid"].as_str().unwrap().to_string(),
|
||||
salt: v["salt"].as_str().unwrap().to_string(),
|
||||
iters: v["iters"].as_u64().unwrap() as u32,
|
||||
nonce: v["nonce"].as_str().unwrap().to_string(),
|
||||
pub_b64: v["pub"].as_str().unwrap_or("").to_string(),
|
||||
totp: v["totp"].as_bool().unwrap_or(false),
|
||||
};
|
||||
}
|
||||
}
|
||||
_ => panic!("no session.challenge"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drive begin → challenge → auth (plain proof only, legacy viewer). Returns sid.
|
||||
async fn raw_auth(viewer: &mut MemTransport, host_id: &str, password: &str) -> String {
|
||||
let ch = raw_challenge(viewer, host_id).await;
|
||||
let proof = password_proof(password, &ch.salt, &ch.nonce, ch.iters).unwrap();
|
||||
viewer
|
||||
.send_text(format!(r#"{{"t":"session.auth","sid":"{}","proof":"{proof}"}}"#, ch.sid))
|
||||
.await
|
||||
.unwrap();
|
||||
ch.sid
|
||||
}
|
||||
|
||||
async fn next_text(viewer: &mut MemTransport, secs: u64) -> Option<Value> {
|
||||
loop {
|
||||
match timeout(Duration::from_secs(secs), viewer.recv()).await {
|
||||
Ok(Some(Ok(TransportMsg::Text(t)))) => {
|
||||
return serde_json::from_str(&t).ok();
|
||||
}
|
||||
Ok(Some(Ok(TransportMsg::Binary(_)))) => continue,
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn accept_advertises_real_perms() {
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let perms = HostPerms { input: true, clipboard: false, files: true, audio: false };
|
||||
let (hcfg, password) = host_cfg(perms, false);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, _hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
let _host = tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
raw_auth(&mut viewer, &host_id, &password).await;
|
||||
let accept = next_text(&mut viewer, 10).await.expect("session.accept");
|
||||
assert_eq!(accept["t"], "session.accept");
|
||||
let p = &accept["info"]["perms"];
|
||||
assert_eq!(p["input"], true, "perms.input must be advertised truthfully");
|
||||
assert_eq!(p["clipboard"], false, "perms.clipboard must be advertised truthfully");
|
||||
assert_eq!(p["files"], true, "perms.files must be advertised truthfully");
|
||||
assert_eq!(p["audio"], false, "perms.audio must be advertised truthfully");
|
||||
// Disabled audio also clears the audio capability flag.
|
||||
assert_eq!(accept["info"]["audio"], false);
|
||||
println!("PERMS_ADVERTISE_OK session.accept carries the enforced perms");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn approve_holds_then_rejects() {
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg(HostPerms::default(), true);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
let _host = tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let sid = raw_auth(&mut viewer, &host_id, &password).await;
|
||||
|
||||
// The platform gets the approve request with the relay-reported identity…
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::ApproveRequest { .. }), 10).await {
|
||||
Some(HostEvent::ApproveRequest { sid: s, peer, user }) => {
|
||||
assert_eq!(s, sid);
|
||||
assert_eq!(peer, "203.0.113.7");
|
||||
assert_eq!(user, "tester@example.com");
|
||||
}
|
||||
_ => panic!("no ApproveRequest emitted in approve mode"),
|
||||
}
|
||||
// …and the wire stays SILENT (no session.accept while held).
|
||||
assert!(
|
||||
timeout(Duration::from_millis(900), viewer.recv()).await.is_err(),
|
||||
"session must be held (no accept) until the platform answers"
|
||||
);
|
||||
|
||||
// Reject → session.deny with a clear reason.
|
||||
hctl_tx
|
||||
.send(HostControl::ApproveReply { sid: sid.clone(), accept: false })
|
||||
.unwrap();
|
||||
let deny = next_text(&mut viewer, 10).await.expect("session.deny");
|
||||
assert_eq!(deny["t"], "session.deny");
|
||||
assert_eq!(deny["sid"], sid.as_str());
|
||||
assert_eq!(deny["reason"], "not accepted by host");
|
||||
println!("APPROVE_REJECT_OK held session denied with 'not accepted by host'");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn approve_accept_finalizes_session() {
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg(HostPerms::default(), true);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
let _host = tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let sid = raw_auth(&mut viewer, &host_id, &password).await;
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::ApproveRequest { .. }), 10)
|
||||
.await
|
||||
.is_some()
|
||||
);
|
||||
hctl_tx
|
||||
.send(HostControl::ApproveReply { sid: sid.clone(), accept: true })
|
||||
.unwrap();
|
||||
let accept = next_text(&mut viewer, 10).await.expect("session.accept");
|
||||
assert_eq!(accept["t"], "session.accept");
|
||||
assert_eq!(accept["sid"], sid.as_str());
|
||||
assert_eq!(accept["info"]["perms"]["input"], true);
|
||||
assert!(
|
||||
wait_host(&mut hev_rx, |e| matches!(e, HostEvent::ViewerConnected { .. }), 10)
|
||||
.await
|
||||
.is_some(),
|
||||
"accepting must finalize the session (ViewerConnected)"
|
||||
);
|
||||
println!("APPROVE_ACCEPT_OK held session finalized on Accept");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 5) Device-key binding (MITM hardening): a viewer that binds the host's REAL
|
||||
// ephemeral pub is accepted; one that binds a SWAPPED pub (what a malicious
|
||||
// relay would do) is denied "key verification failed". Legacy viewers that
|
||||
// send no proof_bind still connect (covered by the approve_* tests above).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn bound_proof_accepts_matching_key() {
|
||||
// Rebuild host_cfg inline so we keep both host_id and password.
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg(HostPerms::default(), false);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, _hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let ch = raw_challenge(&mut viewer, &host_id).await;
|
||||
assert!(!ch.pub_b64.is_empty(), "host must offer an ECDH pub (E2E)");
|
||||
let proof = password_proof(&password, &ch.salt, &ch.nonce, ch.iters).unwrap();
|
||||
// Bind the host's REAL pub (what an honest, un-swapped path yields).
|
||||
let bind = password_proof_bound(&password, &ch.salt, &ch.nonce, ch.iters, &ch.pub_b64).unwrap();
|
||||
viewer
|
||||
.send_text(format!(
|
||||
r#"{{"t":"session.auth","sid":"{}","proof":"{proof}","proof_bind":"{bind}","pub":"{}"}}"#,
|
||||
ch.sid, ch.pub_b64
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let accept = next_text(&mut viewer, 10).await.expect("reply");
|
||||
assert_eq!(accept["t"], "session.accept", "matching bound proof must be accepted");
|
||||
println!("BIND_OK matching channel-bound proof accepted");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn bound_proof_rejects_swapped_key() {
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg(HostPerms::default(), false);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_frame_tx, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, _hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
let ch = raw_challenge(&mut viewer, &host_id).await;
|
||||
let proof = password_proof(&password, &ch.salt, &ch.nonce, ch.iters).unwrap();
|
||||
// A malicious relay swapped the host's pub: the viewer binds the ATTACKER's
|
||||
// pub, but the host verifies against its own real pub → mismatch → deny.
|
||||
// (A valid-looking but different P-256 point, base64 of 65 bytes 0x04||…)
|
||||
let swapped = {
|
||||
use netdesk_core::crypto::HostE2e;
|
||||
HostE2e::new().our_pub_b64 // a different, real ECDH pub
|
||||
};
|
||||
assert_ne!(swapped, ch.pub_b64, "the swapped pub must differ from the host's");
|
||||
let bind = password_proof_bound(&password, &ch.salt, &ch.nonce, ch.iters, &swapped).unwrap();
|
||||
viewer
|
||||
.send_text(format!(
|
||||
r#"{{"t":"session.auth","sid":"{}","proof":"{proof}","proof_bind":"{bind}","pub":"{}"}}"#,
|
||||
ch.sid, swapped
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let deny = next_text(&mut viewer, 10).await.expect("reply");
|
||||
assert_eq!(deny["t"], "session.deny", "a swapped-key bound proof must be denied");
|
||||
assert_eq!(deny["reason"], "key verification failed");
|
||||
println!("BIND_SWAP_OK relay key-swap detected + denied ('key verification failed')");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 6) TOTP 2FA on the native host: the challenge advertises totp:true, a wrong/
|
||||
// missing code is denied "bad 2FA code", and a valid code is accepted.
|
||||
// ---------------------------------------------------------------------------
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn totp_required_rejects_missing_and_accepts_valid() {
|
||||
use netdesk_core::crypto::{password_proof, totp_at};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let secret = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"; // RFC 6238 test secret
|
||||
|
||||
// --- wrong/missing code is denied ---
|
||||
{
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg_totp(HostPerms::default(), false, secret);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_ft, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, _hev) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hc, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
let ch = raw_challenge(&mut viewer, &host_id).await;
|
||||
assert!(ch.totp, "challenge must advertise totp:true when a secret is set");
|
||||
let proof = password_proof(&password, &ch.salt, &ch.nonce, ch.iters).unwrap();
|
||||
// Correct password, NO totp code → bad 2FA.
|
||||
viewer
|
||||
.send_text(format!(r#"{{"t":"session.auth","sid":"{}","proof":"{proof}"}}"#, ch.sid))
|
||||
.await
|
||||
.unwrap();
|
||||
let deny = next_text(&mut viewer, 10).await.expect("reply");
|
||||
assert_eq!(deny["t"], "session.deny");
|
||||
assert_eq!(deny["reason"], "bad 2FA code");
|
||||
}
|
||||
|
||||
// --- a valid, current code is accepted ---
|
||||
{
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, password) = host_cfg_totp(HostPerms::default(), false, secret);
|
||||
let host_id = hcfg.host_id.clone();
|
||||
let (_ft, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, _hev) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (_hc, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
let ch = raw_challenge(&mut viewer, &host_id).await;
|
||||
let proof = password_proof(&password, &ch.salt, &ch.nonce, ch.iters).unwrap();
|
||||
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
|
||||
let code = totp_at(secret, now, 30, 6).unwrap();
|
||||
viewer
|
||||
.send_text(format!(
|
||||
r#"{{"t":"session.auth","sid":"{}","proof":"{proof}","totp":"{code}"}}"#,
|
||||
ch.sid
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let accept = next_text(&mut viewer, 10).await.expect("reply");
|
||||
assert_eq!(accept["t"], "session.accept", "a valid TOTP code must be accepted");
|
||||
}
|
||||
println!("TOTP_OK native host advertises + enforces TOTP (deny bad, accept valid)");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 7) TCP tunnel routing (core): tunnel.open -> HostEvent::TunnelOpen; a raw
|
||||
// [0xF3][tid][payload] frame (NOT E2E) -> HostEvent::TunnelData ONLY when
|
||||
// that tid is open; HostControl::TunnelData -> a [0xF3] frame on the wire.
|
||||
// Proves the server-relayed tunnel plumbing + the 0xF3 disambiguation.
|
||||
// ---------------------------------------------------------------------------
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn tunnel_routing_core() {
|
||||
let (mut viewer, host_relay) = mem_relay_pair();
|
||||
let (hcfg, _password) = host_cfg(HostPerms::default(), false);
|
||||
let (_ft, frame_rx) = mpsc::unbounded_channel::<HostFrame>();
|
||||
let (hev_tx, mut hev_rx) = mpsc::unbounded_channel::<HostEvent>();
|
||||
let (hctl_tx, hctl_rx) = mpsc::unbounded_channel::<HostControl>();
|
||||
tokio::spawn(async move {
|
||||
run_host_with(hcfg, frame_rx, hev_tx, hctl_rx, Box::new(host_relay)).await
|
||||
});
|
||||
|
||||
// The server asks the host to open tunnel tid=7 (no video session needed).
|
||||
viewer
|
||||
.send_text(r#"{"t":"tunnel.open","tid":7,"thost":"127.0.0.1","port":5555}"#.into())
|
||||
.await
|
||||
.unwrap();
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::TunnelOpen { .. }), 10).await {
|
||||
Some(HostEvent::TunnelOpen { tid, thost, port }) => {
|
||||
assert_eq!((tid, thost.as_str(), port), (7, "127.0.0.1", 5555));
|
||||
}
|
||||
_ => panic!("no TunnelOpen"),
|
||||
}
|
||||
|
||||
// Raw [0xF3][tid=7 LE][payload] -> TunnelData for an OPEN tid.
|
||||
let mut frame = vec![0xF3u8];
|
||||
frame.extend_from_slice(&7u32.to_le_bytes());
|
||||
frame.extend_from_slice(b"hello-tcp");
|
||||
viewer.send_binary(frame).await.unwrap();
|
||||
match wait_host(&mut hev_rx, |e| matches!(e, HostEvent::TunnelData { .. }), 10).await {
|
||||
Some(HostEvent::TunnelData { tid, data }) => {
|
||||
assert_eq!(tid, 7);
|
||||
assert_eq!(&data, b"hello-tcp");
|
||||
}
|
||||
_ => panic!("no TunnelData for an open tid"),
|
||||
}
|
||||
|
||||
// Outbound: HostControl::TunnelData -> a raw [0xF3][tid][data] on the wire.
|
||||
hctl_tx
|
||||
.send(HostControl::TunnelData { tid: 7, data: b"reply42".to_vec() })
|
||||
.unwrap();
|
||||
loop {
|
||||
match timeout(Duration::from_secs(10), viewer.recv()).await {
|
||||
Ok(Some(Ok(TransportMsg::Binary(b)))) => {
|
||||
assert_eq!(b[0], 0xF3);
|
||||
assert_eq!(u32::from_le_bytes([b[1], b[2], b[3], b[4]]), 7);
|
||||
assert_eq!(&b[5..], b"reply42");
|
||||
break;
|
||||
}
|
||||
Ok(Some(Ok(TransportMsg::Text(_)))) => continue,
|
||||
_ => panic!("no outbound tunnel frame"),
|
||||
}
|
||||
}
|
||||
println!("TUNNEL_CORE_OK tunnel.open + [0xF3] route in/out, tid-gated");
|
||||
}
|
||||
Reference in new issue
Block a user