NetDesk 1.5.0 — source release

This commit is contained in:
NetDesk committed 2026-10-18 02:02:46 +00:00
commit 40318aafdd
227 files changed
+80798

No files matched your search

+82
View File
@@ -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"
+214
View File
@@ -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);
}
}
+88
View File
@@ -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"));
}
}
+450
View File
@@ -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()
}
+254
View File
@@ -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()
}
+512
View File
@@ -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;
}
}
+646
View File
@@ -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()))
}
+287
View File
@@ -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()))
}
+539
View File
@@ -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;
}
}
+23
View File
@@ -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
+205
View File
@@ -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"));
}
}
+34
View File
@@ -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";
+223
View File
@@ -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");
}
}
+831
View File
@@ -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));
}
}
+184
View File
@@ -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
}
}
+182
View File
@@ -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))
}
+585
View File
@@ -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
+471
View File
@@ -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();
}
}
+133
View File
@@ -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");
}
}
+420
View File
@@ -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);
}
}
+149
View File
@@ -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
+739
View File
@@ -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");
}