NetDesk 1.5.0 — source release
This commit is contained in:
NetDesk committed
commit
40318aafdd
227 files changed
+80798
No files matched your search
@@ -0,0 +1,81 @@
|
||||
# netdesk-host — the NATIVE RUST host (screen capture + H.264 encode + input
|
||||
# injection) that feeds netdesk_core::host::run_host. This replaces the Python
|
||||
# host engine so the Windows app is one unified Rust binary (host + viewer).
|
||||
# Developed by Akronlogic.
|
||||
[package]
|
||||
name = "netdesk-host"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
netdesk-core = { path = "../core" }
|
||||
tokio = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
log = "0.4"
|
||||
# OS clipboard for viewer->host clipboard sync (RustDesk-style).
|
||||
arboard = "3"
|
||||
# Encode extracted custom-cursor bitmaps to PNG for cursor.shape (the viewers
|
||||
# render data:image/png). Base64 of the bytes is done in core on send.
|
||||
png = "0.17"
|
||||
env_logger = "0.11"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { workspace = true }
|
||||
openh264 = { workspace = true }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
# VP9 software encode (libvpx) for constrained links — the constrained-link
|
||||
# codec alongside the hardware H.264 path.
|
||||
netdesk-vpx = { path = "../vpx" }
|
||||
# Opus encode for system-audio capture (WASAPI loopback -> Opus -> CHAN_AUDIO).
|
||||
netdesk-opus = { path = "../opus" }
|
||||
windows = { version = "0.58", features = [
|
||||
"Win32_Media_Audio",
|
||||
"Win32_Media_KernelStreaming",
|
||||
"Win32_System_Com_StructuredStorage",
|
||||
"Win32_Foundation",
|
||||
"Win32_System_Com",
|
||||
"Win32_System_Threading",
|
||||
"Win32_System_Console",
|
||||
"Win32_System_Pipes",
|
||||
"Win32_System_IO",
|
||||
"Win32_Security",
|
||||
"Win32_System_Power",
|
||||
"Win32_Storage_FileSystem",
|
||||
"Win32_System_SystemInformation",
|
||||
"Win32_System_RemoteDesktop",
|
||||
"Win32_System_Services",
|
||||
"Win32_System_StationsAndDesktops",
|
||||
"Win32_Security_Authentication_Identity",
|
||||
"Win32_System_ProcessStatus",
|
||||
"Win32_System_Diagnostics_ToolHelp",
|
||||
"Win32_Media_MediaFoundation",
|
||||
"Win32_Graphics_Direct3D",
|
||||
"Win32_Graphics_Direct3D11",
|
||||
"Win32_Graphics_Dxgi",
|
||||
"Win32_Graphics_Dxgi_Common",
|
||||
"Win32_Graphics_Gdi",
|
||||
"Win32_UI_Input_KeyboardAndMouse",
|
||||
"Win32_UI_WindowsAndMessaging",
|
||||
"Win32_UI_HiDpi",
|
||||
] }
|
||||
|
||||
[[bin]]
|
||||
name = "netdesk-host-probe"
|
||||
path = "src/bin/probe.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "netdesk-host-encbench"
|
||||
path = "src/bin/encbench.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "netdesk-host-engine"
|
||||
path = "src/bin/engine.rs"
|
||||
|
||||
# Privacy-mode feasibility probe: does WDA_EXCLUDEFROMCAPTURE hide a window from
|
||||
# this host's own DXGI + GDI capture? (determines the privacy-mode mechanism)
|
||||
[[bin]]
|
||||
name = "netdesk-privacy-probe"
|
||||
path = "src/bin/privacy_probe.rs"
|
||||
@@ -0,0 +1,261 @@
|
||||
//! Collaborative annotation overlay. A fullscreen, topmost, click-through,
|
||||
//! color-keyed LAYERED window on which the remote viewer's ink is drawn, so the
|
||||
//! person AT the host sees the arrows/circles the operator draws (RustDesk-style
|
||||
//! "annotate"). Only the strokes are opaque; the rest of the window is a magenta
|
||||
//! color-key that is fully transparent, so the real desktop shows through and
|
||||
//! the host user's mouse/keyboard pass straight to it (WS_EX_TRANSPARENT).
|
||||
//!
|
||||
//! Marked `WDA_EXCLUDEFROMCAPTURE` so the overlay does NOT appear in the host's
|
||||
//! own capture — the drawing viewer already renders its ink locally, so this
|
||||
//! prevents a double image on that viewer. (Relaying strokes to OTHER viewers so
|
||||
//! they also see the ink is a follow-on.) Modeled on `privacy.rs`; same
|
||||
//! own-thread + message-pump lifecycle, dropped to clear + remove the overlay.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
use std::thread::JoinHandle;
|
||||
use std::time::Duration;
|
||||
|
||||
use windows::core::w;
|
||||
use windows::Win32::Foundation::{COLORREF, HINSTANCE, HWND, LPARAM, LRESULT, WPARAM};
|
||||
use windows::Win32::Graphics::Gdi::{
|
||||
BeginPaint, CreatePen, CreateSolidBrush, DeleteObject, EndPaint, FillRect, InvalidateRect,
|
||||
LineTo, MoveToEx, SelectObject, HGDIOBJ, PAINTSTRUCT, PS_SOLID,
|
||||
};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
CreateWindowExW, DefWindowProcW, DestroyWindow, DispatchMessageW, GetSystemMetrics,
|
||||
PeekMessageW, RegisterClassW, SetLayeredWindowAttributes, SetWindowDisplayAffinity, ShowWindow,
|
||||
TranslateMessage, LWA_COLORKEY, MSG, PM_REMOVE, SM_CXVIRTUALSCREEN, SM_CYVIRTUALSCREEN,
|
||||
SM_XVIRTUALSCREEN, SM_YVIRTUALSCREEN, SW_SHOWNA, WDA_EXCLUDEFROMCAPTURE, WM_PAINT, WNDCLASSW,
|
||||
WS_EX_LAYERED, WS_EX_NOACTIVATE, WS_EX_TOOLWINDOW, WS_EX_TOPMOST, WS_EX_TRANSPARENT, WS_POPUP,
|
||||
WS_VISIBLE,
|
||||
};
|
||||
|
||||
/// One pen stroke in VIRTUAL-DESKTOP pixel coordinates.
|
||||
#[derive(Clone)]
|
||||
pub struct Stroke {
|
||||
pub color: u32, // 0xRRGGBB
|
||||
pub width: i32,
|
||||
pub points: Vec<(i32, i32)>,
|
||||
}
|
||||
|
||||
// The magenta transparent color-key (COLORREF is 0x00BBGGRR: B=FF, G=00, R=FF).
|
||||
const KEY_BGR: u32 = 0x00FF_00FF;
|
||||
|
||||
static STROKES: OnceLock<Mutex<Vec<Stroke>>> = OnceLock::new();
|
||||
static ORIGIN: OnceLock<(i32, i32)> = OnceLock::new(); // virtual-desktop top-left
|
||||
|
||||
fn strokes() -> &'static Mutex<Vec<Stroke>> {
|
||||
STROKES.get_or_init(|| Mutex::new(Vec::new()))
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn to_bgr(rgb: u32) -> u32 {
|
||||
let r = (rgb >> 16) & 0xFF;
|
||||
let g = (rgb >> 8) & 0xFF;
|
||||
let b = rgb & 0xFF;
|
||||
(b << 16) | (g << 8) | r
|
||||
}
|
||||
|
||||
unsafe extern "system" fn wndproc(h: HWND, m: u32, w: WPARAM, l: LPARAM) -> LRESULT {
|
||||
if m == WM_PAINT {
|
||||
let mut ps = PAINTSTRUCT::default();
|
||||
let hdc = BeginPaint(h, &mut ps);
|
||||
// Paint the whole client area with the transparent color-key first.
|
||||
let keybrush = CreateSolidBrush(COLORREF(KEY_BGR));
|
||||
FillRect(hdc, &ps.rcPaint, keybrush);
|
||||
let _ = DeleteObject(HGDIOBJ(keybrush.0));
|
||||
let (ox, oy) = *ORIGIN.get().unwrap_or(&(0, 0));
|
||||
if let Ok(list) = strokes().lock() {
|
||||
for st in list.iter() {
|
||||
if st.points.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let pen = CreatePen(PS_SOLID, st.width.max(1), COLORREF(to_bgr(st.color)));
|
||||
let old = SelectObject(hdc, HGDIOBJ(pen.0));
|
||||
let p0 = st.points[0];
|
||||
let _ = MoveToEx(hdc, p0.0 - ox, p0.1 - oy, None);
|
||||
for p in &st.points[1..] {
|
||||
let _ = LineTo(hdc, p.0 - ox, p.1 - oy);
|
||||
}
|
||||
if st.points.len() == 1 {
|
||||
// a single tap: a 1px nub so it is visible
|
||||
let _ = LineTo(hdc, p0.0 - ox + 1, p0.1 - oy + 1);
|
||||
}
|
||||
SelectObject(hdc, old);
|
||||
let _ = DeleteObject(HGDIOBJ(pen.0));
|
||||
}
|
||||
}
|
||||
let _ = EndPaint(h, &ps);
|
||||
return LRESULT(0);
|
||||
}
|
||||
DefWindowProcW(h, m, w, l)
|
||||
}
|
||||
|
||||
/// A live annotation overlay. Drop it to clear + remove the overlay.
|
||||
pub struct Annotations {
|
||||
stop: Arc<AtomicBool>,
|
||||
dirty: Arc<AtomicBool>,
|
||||
hwnd: Arc<AtomicIsize>,
|
||||
thread: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl Annotations {
|
||||
/// Raise the (empty, invisible) overlay; strokes appear as they are added.
|
||||
pub fn start() -> Annotations {
|
||||
if let Ok(mut v) = strokes().lock() {
|
||||
v.clear();
|
||||
}
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let dirty = Arc::new(AtomicBool::new(true));
|
||||
let hwnd = Arc::new(AtomicIsize::new(0));
|
||||
let (s2, d2, h2) = (stop.clone(), dirty.clone(), hwnd.clone());
|
||||
let thread = std::thread::Builder::new()
|
||||
.name("netdesk-annot".into())
|
||||
.spawn(move || unsafe { run(s2, d2, h2) })
|
||||
.ok();
|
||||
Annotations { stop, dirty, hwnd, thread }
|
||||
}
|
||||
|
||||
/// Append a stroke (virtual-desktop coords) and request a repaint.
|
||||
pub fn add_stroke(&self, s: Stroke) {
|
||||
if let Ok(mut v) = strokes().lock() {
|
||||
v.push(s);
|
||||
}
|
||||
self.dirty.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Erase all strokes.
|
||||
pub fn clear(&self) {
|
||||
if let Ok(mut v) = strokes().lock() {
|
||||
v.clear();
|
||||
}
|
||||
self.dirty.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// The overlay window handle (0 until created) — for tests.
|
||||
pub fn hwnd(&self) -> isize {
|
||||
self.hwnd.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Annotations {
|
||||
fn drop(&mut self) {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
if let Some(t) = self.thread.take() {
|
||||
let _ = t.join();
|
||||
}
|
||||
if let Ok(mut v) = strokes().lock() {
|
||||
v.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn run(stop: Arc<AtomicBool>, dirty: Arc<AtomicBool>, hwnd_out: Arc<AtomicIsize>) {
|
||||
let hinst = HINSTANCE::default();
|
||||
let keybrush = CreateSolidBrush(COLORREF(KEY_BGR));
|
||||
let class = w!("NdAnnotOverlay");
|
||||
let wc = WNDCLASSW {
|
||||
lpfnWndProc: Some(wndproc),
|
||||
hInstance: hinst,
|
||||
lpszClassName: class,
|
||||
hbrBackground: keybrush,
|
||||
..Default::default()
|
||||
};
|
||||
let _ = RegisterClassW(&wc);
|
||||
|
||||
let x = GetSystemMetrics(SM_XVIRTUALSCREEN);
|
||||
let y = GetSystemMetrics(SM_YVIRTUALSCREEN);
|
||||
let cx = GetSystemMetrics(SM_CXVIRTUALSCREEN).max(1);
|
||||
let cy = GetSystemMetrics(SM_CYVIRTUALSCREEN).max(1);
|
||||
let _ = ORIGIN.set((x, y));
|
||||
|
||||
let hwnd = match CreateWindowExW(
|
||||
WS_EX_TOPMOST | WS_EX_TOOLWINDOW | WS_EX_LAYERED | WS_EX_TRANSPARENT | WS_EX_NOACTIVATE,
|
||||
class,
|
||||
w!("nd-annot"),
|
||||
WS_POPUP | WS_VISIBLE,
|
||||
x,
|
||||
y,
|
||||
cx,
|
||||
cy,
|
||||
None,
|
||||
None,
|
||||
hinst,
|
||||
None,
|
||||
) {
|
||||
Ok(h) => h,
|
||||
Err(e) => {
|
||||
log::warn!("annot: window create failed: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
// Magenta is the transparent color-key: only the strokes are opaque/visible.
|
||||
let _ = SetLayeredWindowAttributes(hwnd, COLORREF(KEY_BGR), 0, LWA_COLORKEY);
|
||||
// Keep the overlay out of OUR OWN capture (the drawing viewer renders ink
|
||||
// locally, so this prevents a double image on it).
|
||||
if let Err(e) = SetWindowDisplayAffinity(hwnd, WDA_EXCLUDEFROMCAPTURE) {
|
||||
log::warn!("annot: SetWindowDisplayAffinity failed: {e}");
|
||||
}
|
||||
let _ = ShowWindow(hwnd, SW_SHOWNA);
|
||||
hwnd_out.store(hwnd.0 as isize, Ordering::Relaxed);
|
||||
log::info!("annot: overlay up ({cx}x{cy} @ {x},{y}); click-through, excluded from capture");
|
||||
|
||||
let mut msg = MSG::default();
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
if dirty.swap(false, Ordering::Relaxed) {
|
||||
let _ = InvalidateRect(hwnd, None, false);
|
||||
}
|
||||
while PeekMessageW(&mut msg, None, 0, 0, PM_REMOVE).as_bool() {
|
||||
let _ = TranslateMessage(&msg);
|
||||
DispatchMessageW(&msg);
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(16));
|
||||
}
|
||||
let _ = DestroyWindow(hwnd);
|
||||
hwnd_out.store(0, Ordering::Relaxed);
|
||||
log::info!("annot: overlay down");
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Instant;
|
||||
use windows::Win32::UI::WindowsAndMessaging::{GetWindowLongPtrW, GWL_EXSTYLE};
|
||||
|
||||
/// The overlay must be a LAYERED + click-through (TRANSPARENT) + TOPMOST
|
||||
/// window and must accept strokes + clear without panicking. Proves the
|
||||
/// collaborative-overlay mechanism (pixel rendering is verified live after
|
||||
/// the capstone redeploy).
|
||||
#[test]
|
||||
fn overlay_is_layered_clickthrough_topmost_and_takes_strokes() {
|
||||
let a = Annotations::start();
|
||||
let end = Instant::now() + Duration::from_millis(1500);
|
||||
while a.hwnd() == 0 && Instant::now() < end {
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
let h = a.hwnd();
|
||||
assert!(h != 0, "overlay window did not create");
|
||||
let ex = unsafe { GetWindowLongPtrW(HWND(h as *mut core::ffi::c_void), GWL_EXSTYLE) } as u32;
|
||||
assert!(ex & WS_EX_LAYERED.0 != 0, "overlay must be layered");
|
||||
assert!(ex & WS_EX_TRANSPARENT.0 != 0, "overlay must be click-through");
|
||||
assert!(ex & WS_EX_TOPMOST.0 != 0, "overlay must be topmost");
|
||||
a.add_stroke(Stroke {
|
||||
color: 0xFF0000,
|
||||
width: 4,
|
||||
points: vec![(100, 100), (400, 300)],
|
||||
});
|
||||
a.add_stroke(Stroke {
|
||||
color: 0x00AAFF,
|
||||
width: 2,
|
||||
points: vec![(50, 50)],
|
||||
});
|
||||
std::thread::sleep(Duration::from_millis(120));
|
||||
a.clear();
|
||||
std::thread::sleep(Duration::from_millis(60));
|
||||
println!("ANNOT_OVERLAY_OK hwnd={h:#x} exstyle={ex:#x}");
|
||||
drop(a);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,551 @@
|
||||
//! System-audio capture for the native host: WASAPI loopback (capture "what you
|
||||
//! hear" from the default render endpoint) -> downmix to 48 kHz stereo i16 ->
|
||||
//! Opus (20 ms frames) -> `on_packet`. The host engine seals each packet and
|
||||
//! sends it as a `CHAN_AUDIO` frame.
|
||||
//!
|
||||
//! Off by default (NETDESK_AUDIO=1 to enable), so it can never affect a normal
|
||||
//! session. Entirely best-effort: any WASAPI error logs once and ends the
|
||||
//! thread — audio just will not stream, video/input are unaffected.
|
||||
//!
|
||||
//! Stock cpal/soundio cannot do WASAPI loopback (RustDesk pins a cpal fork for
|
||||
//! it), so we drive WASAPI directly through the `windows` crate — no extra
|
||||
//! dependency and full control over the format path.
|
||||
//!
|
||||
//! Developed by Akronlogic.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use netdesk_opus::{Decoder, Encoder, FRAME_SIZE, SAMPLE_RATE};
|
||||
|
||||
use std::collections::VecDeque;
|
||||
|
||||
use windows::core::Interface;
|
||||
use windows::Win32::Media::Audio::{
|
||||
eCapture, eConsole, eRender, IAudioCaptureClient, IAudioClient, IAudioRenderClient,
|
||||
IMMDeviceEnumerator, MMDeviceEnumerator, AUDCLNT_BUFFERFLAGS_SILENT, AUDCLNT_SHAREMODE_SHARED,
|
||||
AUDCLNT_STREAMFLAGS_LOOPBACK, WAVEFORMATEX, WAVEFORMATEXTENSIBLE,
|
||||
};
|
||||
use windows::Win32::System::Com::{
|
||||
CoCreateInstance, CoInitializeEx, CoTaskMemFree, CLSCTX_ALL, COINIT_MULTITHREADED,
|
||||
};
|
||||
|
||||
const WAVE_FORMAT_PCM: u16 = 1;
|
||||
const WAVE_FORMAT_IEEE_FLOAT: u16 = 3;
|
||||
const WAVE_FORMAT_EXTENSIBLE: u16 = 0xFFFE;
|
||||
/// 1 second loopback buffer (100-ns units) — generous so draining can lag while
|
||||
/// idle without overrunning.
|
||||
const BUFFER_HNS: i64 = 10_000_000;
|
||||
|
||||
/// Capture loopback audio until `stop`, encoding Opus only while `active`
|
||||
/// (a viewer is connected). Each Opus packet is handed to `on_packet`.
|
||||
///
|
||||
/// Returns (ending the thread) on any fatal WASAPI/Opus error.
|
||||
pub fn run_loopback_opus(
|
||||
stop: Arc<AtomicBool>,
|
||||
active: impl Fn() -> bool,
|
||||
mut on_packet: impl FnMut(Vec<u8>),
|
||||
) {
|
||||
// Back-compat shim: always-on system loopback.
|
||||
let src = Arc::new(AtomicU8::new(1));
|
||||
run_audio(stop, active, src, move |p| on_packet(p))
|
||||
}
|
||||
|
||||
/// On-demand host audio. `source` selects the live source and may change at
|
||||
/// runtime (the viewer picks it): 0 = off, 1 = SYSTEM (WASAPI render loopback,
|
||||
/// "what you hear"), 2 = MICROPHONE (default input). Opus packets go to
|
||||
/// `on_packet` only while `active` (a viewer is streaming). Changing the source
|
||||
/// re-opens the correct device.
|
||||
pub fn run_audio(
|
||||
stop: Arc<AtomicBool>,
|
||||
active: impl Fn() -> bool,
|
||||
source: Arc<AtomicU8>,
|
||||
mut on_packet: impl FnMut(Vec<u8>),
|
||||
) {
|
||||
unsafe {
|
||||
let _ = CoInitializeEx(None, COINIT_MULTITHREADED);
|
||||
}
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
let src = source.load(Ordering::Relaxed);
|
||||
if src == 0 {
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
continue;
|
||||
}
|
||||
if let Err(e) = capture(&stop, &active, &source, src, &mut on_packet) {
|
||||
// DEVICE_INVALIDATED (0x88890004), a device format change, or a
|
||||
// default-endpoint switch all land here. Re-open the (possibly new)
|
||||
// default endpoint after a short backoff — non-fatal, video/input
|
||||
// are unaffected and the source selection is preserved.
|
||||
log::warn!("audio: capture ended (src={src}): {e} — re-opening endpoint");
|
||||
std::thread::sleep(Duration::from_millis(300));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn format_params(mix: *const WAVEFORMATEX) -> (u32, u16, u16, bool) {
|
||||
// (sample_rate, channels, bits_per_sample, is_float)
|
||||
let base = &*mix;
|
||||
let mut is_float = base.wFormatTag == WAVE_FORMAT_IEEE_FLOAT;
|
||||
if base.wFormatTag == WAVE_FORMAT_EXTENSIBLE {
|
||||
// SubFormat GUID Data1: 3 = IEEE float, 1 = PCM.
|
||||
let ext = &*(mix as *const WAVEFORMATEXTENSIBLE);
|
||||
is_float = ext.SubFormat.data1 == WAVE_FORMAT_IEEE_FLOAT as u32;
|
||||
}
|
||||
(
|
||||
base.nSamplesPerSec,
|
||||
base.nChannels,
|
||||
base.wBitsPerSample,
|
||||
is_float,
|
||||
)
|
||||
}
|
||||
|
||||
fn capture(
|
||||
stop: &AtomicBool,
|
||||
active: &impl Fn() -> bool,
|
||||
source: &AtomicU8,
|
||||
src: u8,
|
||||
on_packet: &mut impl FnMut(Vec<u8>),
|
||||
) -> anyhow::Result<()> {
|
||||
unsafe {
|
||||
// COM: ignore the HRESULT — a non-S_OK (already initialized on this
|
||||
// thread with a different mode) is fine for our dedicated thread.
|
||||
let _ = CoInitializeEx(None, COINIT_MULTITHREADED);
|
||||
|
||||
let enumerator: IMMDeviceEnumerator =
|
||||
CoCreateInstance(&MMDeviceEnumerator, None, CLSCTX_ALL)?;
|
||||
// src 2 = microphone: default CAPTURE device, no loopback flag.
|
||||
// else (system): default RENDER device captured in loopback mode.
|
||||
let (flow, stream_flags) = if src == 2 {
|
||||
(eCapture, 0u32)
|
||||
} else {
|
||||
(eRender, AUDCLNT_STREAMFLAGS_LOOPBACK)
|
||||
};
|
||||
let device = enumerator.GetDefaultAudioEndpoint(flow, eConsole)?;
|
||||
let client: IAudioClient = device.Activate(CLSCTX_ALL, None)?;
|
||||
|
||||
let mix = client.GetMixFormat()?;
|
||||
let (rate, channels, bits, is_float) = format_params(mix);
|
||||
let channels = channels as usize;
|
||||
// Resample to 48 kHz only when the device mix rate differs. On Windows the
|
||||
// shared-mode mix is usually already 48000, so the common path is a no-op.
|
||||
// Previously a 44100 Hz device (many USB headsets / Bluetooth / capture
|
||||
// cards) DISABLED audio entirely via a hard bail; now we capture at the
|
||||
// device's native rate and resample the stereo stream to the 48 kHz the
|
||||
// Opus encoder and the web AudioDecoder both require.
|
||||
let need_resample = rate != SAMPLE_RATE;
|
||||
let mut resampler = LinearResampler::new(rate.max(1), SAMPLE_RATE);
|
||||
if need_resample {
|
||||
log::info!("audio: device mix {rate} Hz, {channels}ch -> resampling to 48000 stereo (src={src})");
|
||||
}
|
||||
|
||||
client.Initialize(
|
||||
AUDCLNT_SHAREMODE_SHARED,
|
||||
stream_flags,
|
||||
BUFFER_HNS,
|
||||
0,
|
||||
mix,
|
||||
None,
|
||||
)?;
|
||||
CoTaskMemFree(Some(mix as *const _));
|
||||
|
||||
let capture: IAudioCaptureClient = client.GetService()?;
|
||||
client.Start()?;
|
||||
|
||||
let mut enc = Encoder::new(SAMPLE_RATE, 2)?;
|
||||
let mut accum: Vec<i16> = Vec::with_capacity(FRAME_SIZE * 2 * 4);
|
||||
let frame_len = FRAME_SIZE * 2; // interleaved stereo samples per 20ms
|
||||
|
||||
// Exit (so run_audio re-opens) when the viewer changes the source.
|
||||
while !stop.load(Ordering::Relaxed) && source.load(Ordering::Relaxed) == src {
|
||||
let mut got_any = false;
|
||||
loop {
|
||||
let packet = capture.GetNextPacketSize()?;
|
||||
if packet == 0 {
|
||||
break;
|
||||
}
|
||||
got_any = true;
|
||||
let mut pdata: *mut u8 = std::ptr::null_mut();
|
||||
let mut nframes: u32 = 0;
|
||||
let mut flags: u32 = 0;
|
||||
capture.GetBuffer(&mut pdata, &mut nframes, &mut flags, None, None)?;
|
||||
let nframes = nframes as usize;
|
||||
|
||||
if active() && nframes > 0 {
|
||||
let silent = (flags & AUDCLNT_BUFFERFLAGS_SILENT.0 as u32) != 0;
|
||||
if silent || pdata.is_null() {
|
||||
if need_resample {
|
||||
for _ in 0..nframes {
|
||||
resampler.push(0.0, 0.0, &mut accum);
|
||||
}
|
||||
} else {
|
||||
for _ in 0..nframes {
|
||||
accum.push(0);
|
||||
accum.push(0);
|
||||
}
|
||||
}
|
||||
} else if need_resample {
|
||||
for_each_stereo(pdata, nframes, channels, bits, is_float, |l, r| {
|
||||
resampler.push(l, r, &mut accum);
|
||||
});
|
||||
} else {
|
||||
for_each_stereo(pdata, nframes, channels, bits, is_float, |l, r| {
|
||||
accum.push(to_i16(l));
|
||||
accum.push(to_i16(r));
|
||||
});
|
||||
}
|
||||
}
|
||||
capture.ReleaseBuffer(nframes as u32)?;
|
||||
}
|
||||
|
||||
// Emit whole 20 ms Opus frames.
|
||||
while accum.len() >= frame_len {
|
||||
match enc.encode(&accum[..frame_len], FRAME_SIZE) {
|
||||
Ok(pkt) => on_packet(pkt),
|
||||
Err(e) => log::warn!("audio: opus encode: {e}"),
|
||||
}
|
||||
accum.drain(..frame_len);
|
||||
}
|
||||
|
||||
if !got_any {
|
||||
// Silence / idle: nothing to drain, back off briefly.
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
|
||||
let _ = client.Stop();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one WASAPI buffer (`nframes` frames of `channels` samples) to stereo
|
||||
/// f32 pairs, calling `emit(l, r)` once per frame. Takes the first two channels
|
||||
/// (duplicating channel 0 for mono); extra channels are ignored. Yielding f32
|
||||
/// pairs lets the caller either quantise straight to i16 (native 48 kHz path) or
|
||||
/// feed them through the resampler (non-48 kHz devices) without a second pass.
|
||||
unsafe fn for_each_stereo(
|
||||
pdata: *mut u8,
|
||||
nframes: usize,
|
||||
channels: usize,
|
||||
bits: u16,
|
||||
is_float: bool,
|
||||
mut emit: impl FnMut(f32, f32),
|
||||
) {
|
||||
if is_float && bits == 32 {
|
||||
let src = std::slice::from_raw_parts(pdata as *const f32, nframes * channels);
|
||||
for f in 0..nframes {
|
||||
let base = f * channels;
|
||||
let l = src[base];
|
||||
let r = if channels >= 2 { src[base + 1] } else { l };
|
||||
emit(l, r);
|
||||
}
|
||||
} else if !is_float && bits == 16 {
|
||||
let src = std::slice::from_raw_parts(pdata as *const i16, nframes * channels);
|
||||
for f in 0..nframes {
|
||||
let base = f * channels;
|
||||
let l = src[base] as f32 / 32768.0;
|
||||
let r = if channels >= 2 { src[base + 1] as f32 / 32768.0 } else { l };
|
||||
emit(l, r);
|
||||
}
|
||||
} else if !is_float && bits == 32 {
|
||||
// 32-bit integer PCM (some pro/USB interfaces present this).
|
||||
let src = std::slice::from_raw_parts(pdata as *const i32, nframes * channels);
|
||||
for f in 0..nframes {
|
||||
let base = f * channels;
|
||||
let l = src[base] as f32 / 2_147_483_648.0;
|
||||
let r = if channels >= 2 { src[base + 1] as f32 / 2_147_483_648.0 } else { l };
|
||||
emit(l, r);
|
||||
}
|
||||
} else if is_float && bits == 64 {
|
||||
let src = std::slice::from_raw_parts(pdata as *const f64, nframes * channels);
|
||||
for f in 0..nframes {
|
||||
let base = f * channels;
|
||||
let l = src[base] as f32;
|
||||
let r = if channels >= 2 { src[base + 1] as f32 } else { l };
|
||||
emit(l, r);
|
||||
}
|
||||
} else {
|
||||
// Unknown bit depth: emit silence to keep timing rather than garbage.
|
||||
for _ in 0..nframes {
|
||||
emit(0.0, 0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal stateful linear resampler: stereo f32 in at `src_rate`, stereo i16 out
|
||||
/// at `dst_rate`, appended to `out`. The phase and previous sample persist across
|
||||
/// calls so consecutive WASAPI buffers resample seamlessly (no per-buffer clicks).
|
||||
/// No anti-alias filter — but transparent for speech/UI/media at the 44100->48000
|
||||
/// upsample that is by far the common non-48 kHz case.
|
||||
struct LinearResampler {
|
||||
step: f64, // source samples to advance per output sample (src/dst)
|
||||
t: f64, // phase in [0,1) between prev and cur
|
||||
prev_l: f32,
|
||||
prev_r: f32,
|
||||
primed: bool,
|
||||
}
|
||||
|
||||
impl LinearResampler {
|
||||
fn new(src_rate: u32, dst_rate: u32) -> Self {
|
||||
Self {
|
||||
step: src_rate as f64 / dst_rate as f64,
|
||||
t: 0.0,
|
||||
prev_l: 0.0,
|
||||
prev_r: 0.0,
|
||||
primed: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn push(&mut self, cur_l: f32, cur_r: f32, out: &mut Vec<i16>) {
|
||||
if !self.primed {
|
||||
self.prev_l = cur_l;
|
||||
self.prev_r = cur_r;
|
||||
self.primed = true;
|
||||
}
|
||||
// Emit every output sample whose source position falls in [prev, cur).
|
||||
while self.t < 1.0 {
|
||||
let ft = self.t as f32;
|
||||
out.push(to_i16(self.prev_l + (cur_l - self.prev_l) * ft));
|
||||
out.push(to_i16(self.prev_r + (cur_r - self.prev_r) * ft));
|
||||
self.t += self.step;
|
||||
}
|
||||
self.t -= 1.0;
|
||||
self.prev_l = cur_l;
|
||||
self.prev_r = cur_r;
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn to_i16(f: f32) -> i16 {
|
||||
(f.clamp(-1.0, 1.0) * 32767.0) as i16
|
||||
}
|
||||
|
||||
/// 200 ms shared-mode render buffer for voice playback.
|
||||
const VOICE_BUFFER_HNS: i64 = 2_000_000;
|
||||
|
||||
/// Play the viewer's microphone (two-way voice call) on the host's default render
|
||||
/// endpoint. Reads Opus packets from `rx`, decodes to 48 kHz stereo, converts to
|
||||
/// the device mix format (resampling when the device isn't 48 kHz), and writes to
|
||||
/// the shared-mode render buffer. Best-effort: any WASAPI/Opus error ends the
|
||||
/// current attempt and re-opens; capture + video are unaffected.
|
||||
pub fn run_voice_playback(stop: Arc<AtomicBool>, rx: std::sync::mpsc::Receiver<Vec<u8>>) {
|
||||
unsafe {
|
||||
let _ = CoInitializeEx(None, COINIT_MULTITHREADED);
|
||||
}
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
match unsafe { voice_playback(&stop, &rx) } {
|
||||
Ok(()) => break, // sender dropped / stop requested
|
||||
Err(e) => {
|
||||
log::warn!("voice: playback ended: {e} — re-opening");
|
||||
std::thread::sleep(Duration::from_millis(400));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn voice_playback(
|
||||
stop: &AtomicBool,
|
||||
rx: &std::sync::mpsc::Receiver<Vec<u8>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let _ = CoInitializeEx(None, COINIT_MULTITHREADED);
|
||||
let enumerator: IMMDeviceEnumerator =
|
||||
CoCreateInstance(&MMDeviceEnumerator, None, CLSCTX_ALL)?;
|
||||
let device = enumerator.GetDefaultAudioEndpoint(eRender, eConsole)?;
|
||||
let client: IAudioClient = device.Activate(CLSCTX_ALL, None)?;
|
||||
let mix = client.GetMixFormat()?;
|
||||
let (rate, channels, bits, is_float) = format_params(mix);
|
||||
let channels = channels as usize;
|
||||
client.Initialize(AUDCLNT_SHAREMODE_SHARED, 0, VOICE_BUFFER_HNS, 0, mix, None)?;
|
||||
CoTaskMemFree(Some(mix as *const _));
|
||||
let buf_frames = client.GetBufferSize()?;
|
||||
let render: IAudioRenderClient = client.GetService()?;
|
||||
let mut dec = Decoder::new(SAMPLE_RATE, 2)?;
|
||||
client.Start()?;
|
||||
|
||||
let need_resample = rate != SAMPLE_RATE;
|
||||
let mut rs = LinearResampler::new(SAMPLE_RATE, rate.max(1));
|
||||
// Decoded, device-rate, interleaved-stereo i16 awaiting the render buffer.
|
||||
let mut pcm: VecDeque<i16> = VecDeque::with_capacity(buf_frames as usize * 2);
|
||||
let mut scratch: Vec<i16> = Vec::new();
|
||||
log::info!("voice: render ready ({rate} Hz, {channels}ch) — awaiting viewer mic");
|
||||
let mut rx_pkts: u64 = 0;
|
||||
let mut rx_peak: i32 = 0;
|
||||
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
// Drain queued Opus packets -> PCM (bounded so a burst can't balloon RAM).
|
||||
loop {
|
||||
if pcm.len() > (rate as usize) * 2 {
|
||||
break; // ~1 s already buffered; let it drain first
|
||||
}
|
||||
match rx.try_recv() {
|
||||
Ok(opus) => {
|
||||
if let Ok(frame) = dec.decode(&opus) {
|
||||
rx_pkts += 1;
|
||||
for &s in frame.iter().step_by(64) {
|
||||
let a = (s as i32).abs();
|
||||
if a > rx_peak { rx_peak = a; }
|
||||
}
|
||||
if rx_pkts % 250 == 0 {
|
||||
log::info!("voice: rx {rx_pkts} pkts from viewer mic, recent peak {rx_peak}");
|
||||
rx_peak = 0;
|
||||
}
|
||||
if need_resample {
|
||||
scratch.clear();
|
||||
let n = frame.len() / 2;
|
||||
for i in 0..n {
|
||||
let l = frame[i * 2] as f32 / 32768.0;
|
||||
let r = frame[i * 2 + 1] as f32 / 32768.0;
|
||||
rs.push(l, r, &mut scratch);
|
||||
}
|
||||
pcm.extend(scratch.iter().copied());
|
||||
} else {
|
||||
pcm.extend(frame.iter().copied());
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(std::sync::mpsc::TryRecvError::Empty) => break,
|
||||
Err(std::sync::mpsc::TryRecvError::Disconnected) => return Ok(()),
|
||||
}
|
||||
}
|
||||
// Write as much PCM as the render buffer currently has room for.
|
||||
let padding = client.GetCurrentPadding()?;
|
||||
let avail = buf_frames.saturating_sub(padding);
|
||||
let have = (pcm.len() / 2) as u32;
|
||||
let to_write = avail.min(have);
|
||||
if to_write > 0 {
|
||||
let data = render.GetBuffer(to_write)?;
|
||||
write_render(data, to_write as usize, channels, bits, is_float, &mut pcm);
|
||||
render.ReleaseBuffer(to_write, 0)?;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(8));
|
||||
}
|
||||
let _ = client.Stop();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Write `nframes` of interleaved stereo i16 (popped from `pcm`) to a WASAPI
|
||||
/// render buffer in the device's mix format. Channel 0=L, 1=R, extra = silence;
|
||||
/// a mono device gets the L/R average.
|
||||
unsafe fn write_render(
|
||||
data: *mut u8,
|
||||
nframes: usize,
|
||||
channels: usize,
|
||||
bits: u16,
|
||||
is_float: bool,
|
||||
pcm: &mut VecDeque<i16>,
|
||||
) {
|
||||
if is_float && bits == 32 {
|
||||
let dst = std::slice::from_raw_parts_mut(data as *mut f32, nframes * channels);
|
||||
for f in 0..nframes {
|
||||
let l = pcm.pop_front().unwrap_or(0) as f32 / 32768.0;
|
||||
let r = pcm.pop_front().unwrap_or(0) as f32 / 32768.0;
|
||||
let base = f * channels;
|
||||
if channels == 1 {
|
||||
dst[base] = (l + r) * 0.5;
|
||||
} else {
|
||||
dst[base] = l;
|
||||
dst[base + 1] = r;
|
||||
for c in 2..channels {
|
||||
dst[base + c] = 0.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if !is_float && bits == 16 {
|
||||
let dst = std::slice::from_raw_parts_mut(data as *mut i16, nframes * channels);
|
||||
for f in 0..nframes {
|
||||
let l = pcm.pop_front().unwrap_or(0);
|
||||
let r = pcm.pop_front().unwrap_or(0);
|
||||
let base = f * channels;
|
||||
if channels == 1 {
|
||||
dst[base] = ((l as i32 + r as i32) / 2) as i16;
|
||||
} else {
|
||||
dst[base] = l;
|
||||
dst[base + 1] = r;
|
||||
for c in 2..channels {
|
||||
dst[base + c] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Unsupported mix format: consume + emit silence to keep the clock.
|
||||
for _ in 0..nframes {
|
||||
pcm.pop_front();
|
||||
pcm.pop_front();
|
||||
}
|
||||
let bytes = nframes * channels * (bits as usize / 8).max(1);
|
||||
let dst = std::slice::from_raw_parts_mut(data, bytes);
|
||||
for b in dst.iter_mut() {
|
||||
*b = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Exercises the real WASAPI loopback init + drain/stop path on this machine
|
||||
// (a Windows box with audio devices). Tolerant: a box with no render
|
||||
// endpoint makes capture() bail early, which is still a clean return. The
|
||||
// point is that init + stop never panic and the thread exits promptly.
|
||||
// The non-48 kHz path that used to be a hard bail. Feed a 44100 Hz stereo
|
||||
// sine through the resampler and prove it (a) yields ~48000/44100x as many
|
||||
// output frames and (b) is not silent — i.e. real audio comes out, so a
|
||||
// 44.1 kHz device now streams instead of being disabled.
|
||||
#[test]
|
||||
fn resampler_44100_to_48000_is_audible() {
|
||||
let mut rs = LinearResampler::new(44100, 48000);
|
||||
let mut out: Vec<i16> = Vec::new();
|
||||
let n_in = 44100usize; // 1 second of source
|
||||
for i in 0..n_in {
|
||||
let t = i as f32 / 44100.0;
|
||||
let s = (2.0 * std::f32::consts::PI * 440.0 * t).sin() * 0.8;
|
||||
rs.push(s, s, &mut out);
|
||||
}
|
||||
let out_frames = out.len() / 2;
|
||||
// Expect ~48000 output frames (±1%).
|
||||
let expected = 48000i64;
|
||||
let diff = (out_frames as i64 - expected).abs();
|
||||
assert!(
|
||||
diff < 480,
|
||||
"resampled frame count {out_frames} not ~{expected} (diff {diff})"
|
||||
);
|
||||
let peak = out.iter().map(|&s| (s as i32).abs()).max().unwrap_or(0);
|
||||
assert!(peak > 10_000, "resampled audio is silent (peak={peak})");
|
||||
// Mono-fed => L and R must stay identical through the resampler.
|
||||
assert!(
|
||||
out.chunks(2).all(|c| c[0] == c[1]),
|
||||
"resampler desynced L/R channels"
|
||||
);
|
||||
println!("RESAMPLER_44K_OK out_frames={out_frames} peak={peak}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loopback_inits_and_stops_cleanly() {
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let got = Arc::new(std::sync::Mutex::new(0usize));
|
||||
let s2 = stop.clone();
|
||||
let g2 = got.clone();
|
||||
let h = std::thread::spawn(move || {
|
||||
run_loopback_opus(
|
||||
s2,
|
||||
|| true,
|
||||
move |_pkt| {
|
||||
*g2.lock().unwrap() += 1;
|
||||
},
|
||||
);
|
||||
});
|
||||
std::thread::sleep(Duration::from_millis(300));
|
||||
stop.store(true, Ordering::Relaxed);
|
||||
let start = std::time::Instant::now();
|
||||
while !h.is_finished() && start.elapsed() < Duration::from_secs(3) {
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
assert!(h.is_finished(), "audio capture thread did not stop within 3s");
|
||||
h.join().unwrap();
|
||||
println!("AUDIO_LOOPBACK_OK packets={}", *got.lock().unwrap());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
//! Encode correctness + honest timing on synthetic MOVING content (independent
|
||||
//! of whether the live desktop is active). Encodes N frames, decodes them back
|
||||
//! with the real viewer decoder, and asserts the roundtrip reproduces the right
|
||||
//! dimensions — proving the Annex-B the host emits is valid for the viewer.
|
||||
//!
|
||||
//! Usage: netdesk-host-encbench [sw|mf] (default sw = openh264 software)
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
fn make_frame(w: usize, h: usize, t: usize) -> Vec<u8> {
|
||||
let mut buf = vec![0u8; w * h * 4];
|
||||
let bx = (t * 13) % (w.saturating_sub(120).max(1));
|
||||
let by = (t * 7) % (h.saturating_sub(120).max(1));
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
let i = (y * w + x) * 4;
|
||||
let inbox = x >= bx && x < bx + 120 && y >= by && y < by + 120;
|
||||
if inbox {
|
||||
buf[i] = 255;
|
||||
buf[i + 1] = 255;
|
||||
buf[i + 2] = 255;
|
||||
buf[i + 3] = 255;
|
||||
} else {
|
||||
buf[i] = ((x + t * 3) & 0xff) as u8;
|
||||
buf[i + 1] = ((y + t * 2) & 0xff) as u8;
|
||||
buf[i + 2] = ((x + y + t) & 0xff) as u8;
|
||||
buf[i + 3] = 255;
|
||||
}
|
||||
}
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
/// Mean (R,G,B) of a decoded row.
|
||||
fn row_means(rgb: &[u8], w: usize, row: usize) -> (f64, f64, f64) {
|
||||
let base = row * w * 3;
|
||||
let (mut r, mut g, mut b) = (0u64, 0u64, 0u64);
|
||||
for col in 0..w {
|
||||
let o = base + col * 3;
|
||||
r += rgb[o] as u64;
|
||||
g += rgb[o + 1] as u64;
|
||||
b += rgb[o + 2] as u64;
|
||||
}
|
||||
let n = w as f64;
|
||||
(r as f64 / n, g as f64 / n, b as f64 / n)
|
||||
}
|
||||
|
||||
/// Solid gray frame (128,128,128) with a moving white box, so a correct decode
|
||||
/// is gray everywhere and any green row stands out immediately.
|
||||
fn gray_frame(w: usize, h: usize, t: usize) -> Vec<u8> {
|
||||
let mut buf = vec![128u8; w * h * 4];
|
||||
for p in buf.chunks_exact_mut(4) {
|
||||
p[3] = 255;
|
||||
}
|
||||
let bx = (t * 11) % (w.saturating_sub(100).max(1));
|
||||
for y in 0..100.min(h) {
|
||||
for x in bx..(bx + 100).min(w) {
|
||||
let i = (y * w + x) * 4;
|
||||
buf[i] = 255;
|
||||
buf[i + 1] = 255;
|
||||
buf[i + 2] = 255;
|
||||
}
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
let mode = std::env::args().nth(1).unwrap_or_else(|| "sw".into());
|
||||
let (w, h) = (1920usize, 1080usize);
|
||||
let n = 120usize;
|
||||
println!("encode roundtrip [{}]: {}x{}, {} frames", mode, w, h, n);
|
||||
|
||||
// Green-line hunt: QSV-encode gray frames, decode through the HARDWARE
|
||||
// decoder (best_decoder), and scan every decoded row for anomalous green.
|
||||
#[cfg(windows)]
|
||||
if mode == "mfdec" {
|
||||
let _ = tracing_subscriber::fmt().with_max_level(tracing::Level::INFO).try_init();
|
||||
let w = std::env::var("ND_W").ok().and_then(|s| s.parse().ok()).unwrap_or(w);
|
||||
let h = std::env::var("ND_H").ok().and_then(|s| s.parse().ok()).unwrap_or(h);
|
||||
println!(" mfdec dims: {}x{}", w, h);
|
||||
let mut enc = netdesk_host::encode_mf::MfEncoder::new(w, h, 60, 8_000_000)?;
|
||||
let mut dec = netdesk_core::decode::best_decoder();
|
||||
println!("encoder: {}", enc.name());
|
||||
let mut aus: Vec<Vec<u8>> = Vec::new();
|
||||
for t in 0..30 {
|
||||
for ef in enc.encode(&gray_frame(w, h, t))? {
|
||||
aus.push(ef.data);
|
||||
}
|
||||
}
|
||||
for ef in enc.drain()? {
|
||||
aus.push(ef.data);
|
||||
}
|
||||
let mut frames_checked = 0usize;
|
||||
let mut green_rows_total = 0usize;
|
||||
let mut first_report = String::new();
|
||||
for au in &aus {
|
||||
if let Some(fr) = dec.decode(au)? {
|
||||
frames_checked += 1;
|
||||
let mut green_here = 0usize;
|
||||
let mut first = None;
|
||||
for row in 0..fr.h {
|
||||
let (r, g, b) = row_means(&fr.rgb, fr.w, row);
|
||||
// "green" = G clearly above both R and B (gray/white have R≈G≈B).
|
||||
if g - r > 40.0 && g - b > 40.0 {
|
||||
green_here += 1;
|
||||
if first.is_none() {
|
||||
first = Some((row, r, g, b));
|
||||
}
|
||||
}
|
||||
}
|
||||
green_rows_total += green_here;
|
||||
if green_here > 0 && first_report.is_empty() {
|
||||
let (row, r, g, b) = first.unwrap();
|
||||
first_report = format!(
|
||||
"frame {}x{}: {} green rows; first row {} rgb=({:.0},{:.0},{:.0})",
|
||||
fr.w, fr.h, green_here, row, r, g, b
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
println!("decoded {} frames; total green rows: {}", frames_checked, green_rows_total);
|
||||
if green_rows_total == 0 {
|
||||
println!("NO GREEN LINES via hardware decode");
|
||||
} else {
|
||||
println!("GREEN DETECTED: {}", first_report);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut dec = netdesk_core::decode::H264Decoder::new()?;
|
||||
let mut keyframes = 0usize;
|
||||
let mut decoded_ok = 0usize;
|
||||
let mut total_bytes = 0usize;
|
||||
let mut enc_ns = 0u128;
|
||||
let mut emitted = 0usize;
|
||||
|
||||
// Collect all AUs (with timing) then decode, so the two encoders share one path.
|
||||
let mut aus: Vec<(Vec<u8>, bool)> = Vec::new();
|
||||
|
||||
#[cfg(windows)]
|
||||
if mode == "mf" {
|
||||
let mut enc = match netdesk_host::encode_mf::MfEncoder::new(w, h, 60, 8_000_000) {
|
||||
Ok(e) => e,
|
||||
Err(e) => {
|
||||
println!("no hardware encoder ({e}); skipping mf test");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
println!("hardware encoder: {} (hw={})", enc.name(), enc.is_hardware());
|
||||
for t in 0..n {
|
||||
let frame = make_frame(w, h, t);
|
||||
let t0 = Instant::now();
|
||||
let outs = enc.encode(&frame)?;
|
||||
enc_ns += t0.elapsed().as_nanos();
|
||||
for ef in outs {
|
||||
aus.push((ef.data, ef.keyframe));
|
||||
}
|
||||
}
|
||||
for ef in enc.drain()? {
|
||||
aus.push((ef.data, ef.keyframe));
|
||||
}
|
||||
}
|
||||
|
||||
if mode != "mf" {
|
||||
let mut enc = netdesk_host::encode::H264Encoder::new(w, h, 60, 8_000_000)?;
|
||||
for t in 0..n {
|
||||
let frame = make_frame(w, h, t);
|
||||
let t0 = Instant::now();
|
||||
let out = enc.encode(&frame)?;
|
||||
enc_ns += t0.elapsed().as_nanos();
|
||||
if let Some(ef) = out {
|
||||
aus.push((ef.data, ef.keyframe));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (data, keyframe) in &aus {
|
||||
emitted += 1;
|
||||
total_bytes += data.len();
|
||||
if *keyframe {
|
||||
keyframes += 1;
|
||||
}
|
||||
if let Some(fr) = dec.decode(data)? {
|
||||
if fr.w == w && fr.h == h {
|
||||
decoded_ok += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let enc_ms = enc_ns as f64 / n as f64 / 1e6;
|
||||
println!(" encode avg: {:.2} ms/frame (~{:.0} fps)", enc_ms, 1000.0 / enc_ms);
|
||||
println!(" AUs emitted: {} keyframes: {} decoded_ok: {}", emitted, keyframes, decoded_ok);
|
||||
println!(" avg AU: {:.1} KB", total_bytes as f64 / emitted.max(1) as f64 / 1024.0);
|
||||
|
||||
if keyframes == 0 {
|
||||
anyhow::bail!("no keyframe emitted");
|
||||
}
|
||||
if decoded_ok < emitted.saturating_sub(2) {
|
||||
anyhow::bail!("roundtrip decode failed: {}/{} decoded", decoded_ok, emitted);
|
||||
}
|
||||
println!("ROUNDTRIP OK");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
//! Standalone NetDesk native host engine. Captures the screen and serves
|
||||
//! viewers over the NetDesk protocol with the best available H.264 encoder and
|
||||
//! SendInput injection. Configured via environment:
|
||||
//! ND_HOST_ID (required) the id viewers connect to
|
||||
//! ND_PASSWORD (required) session password (never leaves this machine)
|
||||
//! ND_WS_URL default wss://netdesk.space/ws
|
||||
//! ND_NAME default "NetDesk (Rust)"
|
||||
//! ND_MONITOR default 0
|
||||
//! ND_FPS default 60
|
||||
//! ND_BITRATE default 8000000
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
|
||||
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let host_id = std::env::var("ND_HOST_ID").map_err(|_| anyhow::anyhow!("set ND_HOST_ID"))?;
|
||||
let password =
|
||||
std::env::var("ND_PASSWORD").map_err(|_| anyhow::anyhow!("set ND_PASSWORD"))?;
|
||||
let cfg = netdesk_host::run::EngineConfig {
|
||||
ws_url: std::env::var("ND_WS_URL").unwrap_or_else(|_| "wss://netdesk.space/ws".into()),
|
||||
host_id,
|
||||
name: std::env::var("ND_NAME").unwrap_or_else(|_| "NetDesk (Rust)".into()),
|
||||
password,
|
||||
monitor: envnum("ND_MONITOR", 0),
|
||||
fps: envnum("ND_FPS", 60),
|
||||
bitrate: envnum("ND_BITRATE", 8_000_000),
|
||||
// Host-enforced permissions + approve-before-connect (see run.rs):
|
||||
// ND_PERM_INPUT / ND_PERM_CLIPBOARD / ND_PERM_FILES / ND_PERM_AUDIO
|
||||
// (default 1) and ND_APPROVE (default 0).
|
||||
perms: netdesk_core::host::HostPerms {
|
||||
input: envflag("ND_PERM_INPUT", true),
|
||||
clipboard: envflag("ND_PERM_CLIPBOARD", true),
|
||||
files: envflag("ND_PERM_FILES", true),
|
||||
audio: envflag("ND_PERM_AUDIO", true),
|
||||
},
|
||||
approve: envflag("ND_APPROVE", false),
|
||||
totp_secret: std::env::var("ND_TOTP").unwrap_or_default(),
|
||||
// Encrypted direct-IP: ND_DIRECT_PORT makes the host listen on that
|
||||
// TCP port for a viewer that dials it directly (no relay);
|
||||
// ND_ALLOW_IPS is an optional CIDR allow-list.
|
||||
direct_port: std::env::var("ND_DIRECT_PORT").ok().and_then(|s| s.parse().ok()),
|
||||
allow_ips: std::env::var("ND_ALLOW_IPS").unwrap_or_default(),
|
||||
// ND_MAX_VIEWERS>1 allows multiple simultaneous direct-IP viewers
|
||||
// (1 controller + N observers of the one capture).
|
||||
max_viewers: std::env::var("ND_MAX_VIEWERS")
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(1)
|
||||
.max(1),
|
||||
};
|
||||
let rt = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?;
|
||||
let handle = netdesk_host::run::EngineHandle::new();
|
||||
rt.block_on(netdesk_host::run::run_engine(cfg, handle))?;
|
||||
}
|
||||
#[cfg(not(windows))]
|
||||
println!("engine is Windows-only");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn envnum<T: std::str::FromStr>(key: &str, default: T) -> T {
|
||||
std::env::var(key).ok().and_then(|s| s.parse().ok()).unwrap_or(default)
|
||||
}
|
||||
|
||||
/// "1"/"true"/"yes" = true, "0"/"false"/"no" = false, unset = `default`.
|
||||
#[cfg(windows)]
|
||||
fn envflag(key: &str, default: bool) -> bool {
|
||||
match std::env::var(key).as_deref().map(str::trim) {
|
||||
Ok("1") | Ok("true") | Ok("yes") => true,
|
||||
Ok("0") | Ok("false") | Ok("no") => false,
|
||||
_ => default,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
//! Privacy-mode feasibility probe. Determines whether a window marked
|
||||
//! `SetWindowDisplayAffinity(WDA_EXCLUDEFROMCAPTURE)` is hidden from THIS host's
|
||||
//! own screen capture — both DXGI Desktop Duplication (`grab`) and the GDI
|
||||
//! BitBlt fallback (`grab_seed`). If it is, a fullscreen black excluded window
|
||||
//! is a real, driver-free privacy mode: the local monitor shows black while the
|
||||
//! remote operator still captures the desktop behind it.
|
||||
//!
|
||||
//! Method: show a small solid-magenta popup at a known spot, capture with the
|
||||
//! affinity OFF (control — magenta must appear) then ON (magenta must vanish),
|
||||
//! and count magenta pixels in each frame. Safe: the window is small and shown
|
||||
//! for ~2s, then destroyed. Developed by Akronlogic.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
use windows::core::w;
|
||||
use windows::Win32::Foundation::{COLORREF, HINSTANCE, HWND, LPARAM, LRESULT, WPARAM};
|
||||
use windows::Win32::Graphics::Gdi::CreateSolidBrush;
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
CreateWindowExW, DefWindowProcW, DestroyWindow, DispatchMessageW, PeekMessageW, RegisterClassW,
|
||||
SetWindowDisplayAffinity, ShowWindow, TranslateMessage, MSG, PM_REMOVE, SW_SHOW,
|
||||
WDA_EXCLUDEFROMCAPTURE, WDA_NONE, WNDCLASSW, WS_EX_TOPMOST, WS_POPUP, WS_VISIBLE,
|
||||
};
|
||||
|
||||
use netdesk_host::capture::DxgiCapturer;
|
||||
|
||||
unsafe extern "system" fn wndproc(h: HWND, m: u32, w: WPARAM, l: LPARAM) -> LRESULT {
|
||||
DefWindowProcW(h, m, w, l)
|
||||
}
|
||||
|
||||
fn pump(ms: u64) {
|
||||
let end = Instant::now() + Duration::from_millis(ms);
|
||||
while Instant::now() < end {
|
||||
unsafe {
|
||||
let mut msg = MSG::default();
|
||||
while PeekMessageW(&mut msg, None, 0, 0, PM_REMOVE).as_bool() {
|
||||
let _ = TranslateMessage(&msg);
|
||||
DispatchMessageW(&msg);
|
||||
}
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(15));
|
||||
}
|
||||
}
|
||||
|
||||
/// Count magenta (B=255,G=0,R=255 with tolerance) pixels in a BGRA buffer.
|
||||
fn count_magenta(buf: &[u8]) -> usize {
|
||||
buf.chunks_exact(4)
|
||||
.filter(|p| p[0] >= 250 && p[1] <= 6 && p[2] >= 250)
|
||||
.count()
|
||||
}
|
||||
|
||||
/// Freshest DXGI frame over `ms` (None if the duplication reported no change).
|
||||
fn freshest_dxgi(cap: &mut DxgiCapturer, ms: u64) -> Option<Vec<u8>> {
|
||||
let end = Instant::now() + Duration::from_millis(ms);
|
||||
let mut last = None;
|
||||
while Instant::now() < end {
|
||||
if let Ok(Some(f)) = cap.grab(80) {
|
||||
last = Some(f);
|
||||
}
|
||||
}
|
||||
last
|
||||
}
|
||||
|
||||
fn main() {
|
||||
unsafe {
|
||||
let hinst = HINSTANCE::default();
|
||||
let brush = CreateSolidBrush(COLORREF(0x00FF_00FF)); // B=FF,G=00,R=FF = magenta
|
||||
let class = w!("NdPrivacyProbe");
|
||||
let wc = WNDCLASSW {
|
||||
lpfnWndProc: Some(wndproc),
|
||||
hInstance: hinst,
|
||||
lpszClassName: class,
|
||||
hbrBackground: brush,
|
||||
..Default::default()
|
||||
};
|
||||
let atom = RegisterClassW(&wc);
|
||||
if atom == 0 {
|
||||
eprintln!("RegisterClassW failed");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
let hwnd = match CreateWindowExW(
|
||||
WS_EX_TOPMOST,
|
||||
class,
|
||||
w!("nd-privacy-probe"),
|
||||
WS_POPUP | WS_VISIBLE,
|
||||
120,
|
||||
120,
|
||||
520,
|
||||
380,
|
||||
None,
|
||||
None,
|
||||
hinst,
|
||||
None,
|
||||
) {
|
||||
Ok(h) => h,
|
||||
Err(e) => {
|
||||
eprintln!("CreateWindowExW failed: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
let _ = ShowWindow(hwnd, SW_SHOW);
|
||||
pump(400);
|
||||
|
||||
let mut cap = match DxgiCapturer::new(0) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
eprintln!("DxgiCapturer::new failed: {e}");
|
||||
let _ = DestroyWindow(hwnd);
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
|
||||
// CONTROL: affinity OFF — the window must be visible to both captures.
|
||||
let _ = SetWindowDisplayAffinity(hwnd, WDA_NONE);
|
||||
pump(300);
|
||||
let ctl_gdi = cap.grab_seed().map(|b| count_magenta(&b)).unwrap_or(usize::MAX);
|
||||
let ctl_dxgi = freshest_dxgi(&mut cap, 700).map(|b| count_magenta(&b));
|
||||
|
||||
// EXCLUDE: affinity ON — the window must vanish from capture.
|
||||
let _ = SetWindowDisplayAffinity(hwnd, WDA_EXCLUDEFROMCAPTURE);
|
||||
pump(400);
|
||||
let exc_gdi = cap.grab_seed().map(|b| count_magenta(&b)).unwrap_or(usize::MAX);
|
||||
let exc_dxgi = freshest_dxgi(&mut cap, 900).map(|b| count_magenta(&b));
|
||||
|
||||
let _ = DestroyWindow(hwnd);
|
||||
|
||||
println!("PRIVACY_PROBE control_gdi_magenta={ctl_gdi} exclude_gdi_magenta={exc_gdi}");
|
||||
println!(
|
||||
"PRIVACY_PROBE control_dxgi_magenta={:?} exclude_dxgi_magenta={:?}",
|
||||
ctl_dxgi, exc_dxgi
|
||||
);
|
||||
|
||||
let gdi_ok = ctl_gdi > 1000 && exc_gdi == 0;
|
||||
// DXGI: control should show magenta; exclude should show none. If the
|
||||
// duplication reported no change in a phase (None), note it.
|
||||
let dxgi_ok = matches!((ctl_dxgi, exc_dxgi), (Some(c), Some(e)) if c > 1000 && e == 0);
|
||||
if gdi_ok && dxgi_ok {
|
||||
println!("VERDICT: PRIVACY_WDA_WORKS — WDA_EXCLUDEFROMCAPTURE hides the window from BOTH DXGI and GDI capture");
|
||||
} else if gdi_ok {
|
||||
println!("VERDICT: PRIVACY_WDA_GDI_ONLY — GDI excludes it; DXGI result {:?}/{:?} (inconclusive or not excluded)", ctl_dxgi, exc_dxgi);
|
||||
} else {
|
||||
println!("VERDICT: PRIVACY_WDA_FAILS — the window was still captured; privacy needs an IDD virtual display");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
//! Capture + encode benchmark: open DXGI duplication on monitor 0, encode each
|
||||
//! frame to H.264, and report per-stage timing + bitstream sizes. Proves the
|
||||
//! native capture->encode pipeline works and shows where we stand on latency.
|
||||
|
||||
#[cfg(windows)]
|
||||
use std::time::Instant;
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
netdesk_host::set_dpi_aware();
|
||||
let mut cap = netdesk_host::capture::DxgiCapturer::new(0)?;
|
||||
let (w, h) = (cap.width() as usize, cap.height() as usize);
|
||||
println!("capture: {}x{}", w, h);
|
||||
|
||||
let mut enc = netdesk_host::encode::H264Encoder::new(w, h, 60, 8_000_000)?;
|
||||
|
||||
let mut last: Option<Vec<u8>> = None;
|
||||
let mut frames = 0usize;
|
||||
let mut keyframes = 0usize;
|
||||
let mut cap_ns = 0u128;
|
||||
let mut enc_ns = 0u128;
|
||||
let mut total_bytes = 0usize;
|
||||
let mut encoded = 0usize;
|
||||
|
||||
// warm up + run ~150 iterations
|
||||
for i in 0..150 {
|
||||
let t0 = Instant::now();
|
||||
let frame = match cap.grab(50)? {
|
||||
Some(buf) => {
|
||||
last = Some(buf);
|
||||
last.as_ref().unwrap()
|
||||
}
|
||||
None => match last.as_ref() {
|
||||
Some(b) => b, // screen unchanged: re-encode last (keeps fps up)
|
||||
None => continue,
|
||||
},
|
||||
};
|
||||
let t1 = Instant::now();
|
||||
let out = enc.encode(frame)?;
|
||||
let t2 = Instant::now();
|
||||
|
||||
if i >= 10 {
|
||||
// skip warmup
|
||||
frames += 1;
|
||||
cap_ns += (t1 - t0).as_nanos();
|
||||
enc_ns += (t2 - t1).as_nanos();
|
||||
if let Some(ef) = &out {
|
||||
encoded += 1;
|
||||
total_bytes += ef.data.len();
|
||||
if ef.keyframe {
|
||||
keyframes += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if frames == 0 {
|
||||
anyhow::bail!("no frames");
|
||||
}
|
||||
let cap_ms = cap_ns as f64 / frames as f64 / 1e6;
|
||||
let enc_ms = enc_ns as f64 / frames as f64 / 1e6;
|
||||
println!("measured over {} frames ({} produced output, {} keyframes)", frames, encoded, keyframes);
|
||||
println!(" capture avg: {:.2} ms", cap_ms);
|
||||
println!(" encode avg: {:.2} ms", enc_ms);
|
||||
println!(" pipeline: {:.2} ms (~{:.0} fps ceiling)", cap_ms + enc_ms, 1000.0 / (cap_ms + enc_ms));
|
||||
if encoded > 0 {
|
||||
let avg_kb = total_bytes as f64 / encoded as f64 / 1024.0;
|
||||
println!(" avg AU size: {:.1} KB (~{:.1} Mbit/s @60fps)", avg_kb, avg_kb * 1024.0 * 8.0 * 60.0 / 1e6);
|
||||
}
|
||||
}
|
||||
#[cfg(not(windows))]
|
||||
println!("probe is Windows-only");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
//! Windows screen capture via DXGI Desktop Duplication — the same GPU path
|
||||
//! dxcam / RustDesk use: the OS hands us the latest desktop frame as a GPU
|
||||
//! texture, we copy it to a CPU-readable staging texture and read BGRA out.
|
||||
//! `grab` returns `None` when the desktop is unchanged (no new frame), so an
|
||||
//! idle screen costs almost nothing.
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use windows::core::Interface;
|
||||
use windows::Win32::Foundation::{HMODULE, E_ACCESSDENIED};
|
||||
use windows::Win32::Graphics::Direct3D::{
|
||||
D3D_DRIVER_TYPE_HARDWARE, D3D_FEATURE_LEVEL_11_0,
|
||||
};
|
||||
use windows::Win32::Graphics::Direct3D11::{
|
||||
D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D,
|
||||
D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_FLAG, D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ,
|
||||
D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
|
||||
};
|
||||
use windows::Win32::Graphics::Dxgi::Common::DXGI_SAMPLE_DESC;
|
||||
use windows::Win32::Graphics::Dxgi::{
|
||||
IDXGIAdapter, IDXGIDevice, IDXGIOutput, IDXGIOutput1, IDXGIOutputDuplication, IDXGIResource,
|
||||
DXGI_ERROR_ACCESS_LOST, DXGI_ERROR_WAIT_TIMEOUT, DXGI_OUTDUPL_FRAME_INFO,
|
||||
};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
GetSystemMetrics, SM_CXVIRTUALSCREEN, SM_CYVIRTUALSCREEN, SM_XVIRTUALSCREEN,
|
||||
SM_YVIRTUALSCREEN,
|
||||
};
|
||||
|
||||
/// Geometry of one display output (flat DXGI output index + its desktop rect).
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct MonitorDesc {
|
||||
pub index: u32,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub left: i32,
|
||||
pub top: i32,
|
||||
}
|
||||
|
||||
/// A live Desktop Duplication session for one monitor.
|
||||
pub struct DxgiCapturer {
|
||||
device: ID3D11Device,
|
||||
context: ID3D11DeviceContext,
|
||||
dupl: IDXGIOutputDuplication,
|
||||
width: u32,
|
||||
height: u32,
|
||||
left: i32,
|
||||
top: i32,
|
||||
staging: Option<ID3D11Texture2D>,
|
||||
}
|
||||
|
||||
impl DxgiCapturer {
|
||||
/// Enumerate all display outputs on the default adapter, in the same flat
|
||||
/// index order `new(index)` uses. Best-effort: returns an empty vec on any
|
||||
/// failure (callers fall back to a single monitor). The index is what the
|
||||
/// host advertises as `monitors[].i` and what a `monitor.switch{mon}` sends.
|
||||
pub fn enumerate() -> Vec<MonitorDesc> {
|
||||
let mut out = Vec::new();
|
||||
unsafe {
|
||||
let mut device: Option<ID3D11Device> = None;
|
||||
let levels = [D3D_FEATURE_LEVEL_11_0];
|
||||
if D3D11CreateDevice(
|
||||
None,
|
||||
D3D_DRIVER_TYPE_HARDWARE,
|
||||
HMODULE::default(),
|
||||
D3D11_CREATE_DEVICE_FLAG(0),
|
||||
Some(&levels),
|
||||
D3D11_SDK_VERSION,
|
||||
Some(&mut device),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err()
|
||||
{
|
||||
return out;
|
||||
}
|
||||
let device = match device {
|
||||
Some(d) => d,
|
||||
None => return out,
|
||||
};
|
||||
let dxgi_device: IDXGIDevice = match device.cast() {
|
||||
Ok(d) => d,
|
||||
Err(_) => return out,
|
||||
};
|
||||
let adapter: IDXGIAdapter = match dxgi_device.GetAdapter() {
|
||||
Ok(a) => a,
|
||||
Err(_) => return out,
|
||||
};
|
||||
let mut i = 0u32;
|
||||
loop {
|
||||
match adapter.EnumOutputs(i) {
|
||||
Ok(output) => {
|
||||
if let Ok(desc) = output.GetDesc() {
|
||||
let c = desc.DesktopCoordinates;
|
||||
out.push(MonitorDesc {
|
||||
index: i,
|
||||
width: (c.right - c.left) as u32,
|
||||
height: (c.bottom - c.top) as u32,
|
||||
left: c.left,
|
||||
top: c.top,
|
||||
});
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
Err(_) => break, // DXGI_ERROR_NOT_FOUND: no more outputs
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Open duplication for `output_index` (0 = first monitor).
|
||||
pub fn new(output_index: u32) -> Result<Self> {
|
||||
unsafe {
|
||||
let mut device: Option<ID3D11Device> = None;
|
||||
let mut context: Option<ID3D11DeviceContext> = None;
|
||||
let levels = [D3D_FEATURE_LEVEL_11_0];
|
||||
D3D11CreateDevice(
|
||||
None,
|
||||
D3D_DRIVER_TYPE_HARDWARE,
|
||||
HMODULE::default(),
|
||||
D3D11_CREATE_DEVICE_FLAG(0),
|
||||
Some(&levels),
|
||||
D3D11_SDK_VERSION,
|
||||
Some(&mut device),
|
||||
None,
|
||||
Some(&mut context),
|
||||
)?;
|
||||
let device = device.ok_or_else(|| anyhow!("D3D11 device is null"))?;
|
||||
let context = context.ok_or_else(|| anyhow!("D3D11 context is null"))?;
|
||||
|
||||
let dxgi_device: IDXGIDevice = device.cast()?;
|
||||
let adapter: IDXGIAdapter = dxgi_device.GetAdapter()?;
|
||||
let output: IDXGIOutput = adapter.EnumOutputs(output_index)?;
|
||||
let desc = output.GetDesc()?;
|
||||
let width = (desc.DesktopCoordinates.right - desc.DesktopCoordinates.left) as u32;
|
||||
let height = (desc.DesktopCoordinates.bottom - desc.DesktopCoordinates.top) as u32;
|
||||
let left = desc.DesktopCoordinates.left;
|
||||
let top = desc.DesktopCoordinates.top;
|
||||
|
||||
let output1: IDXGIOutput1 = output.cast()?;
|
||||
let dupl = output1.DuplicateOutput(&device)?;
|
||||
|
||||
Ok(Self {
|
||||
device,
|
||||
context,
|
||||
dupl,
|
||||
width,
|
||||
height,
|
||||
left,
|
||||
top,
|
||||
staging: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn width(&self) -> u32 {
|
||||
self.width
|
||||
}
|
||||
pub fn height(&self) -> u32 {
|
||||
self.height
|
||||
}
|
||||
/// Top-left of this monitor in virtual-desktop coordinates (for mapping the
|
||||
/// global cursor position into frame-pixel space).
|
||||
pub fn origin(&self) -> (i32, i32) {
|
||||
(self.left, self.top)
|
||||
}
|
||||
|
||||
/// Grab the next frame as tightly-packed BGRA (`width*height*4`). Returns
|
||||
/// `Ok(None)` when nothing changed within `timeout_ms`. On a desktop switch
|
||||
/// the duplication is lost (`access lost`) and the caller re-creates us.
|
||||
pub fn grab(&mut self, timeout_ms: u32) -> Result<Option<Vec<u8>>> {
|
||||
unsafe {
|
||||
let mut info = DXGI_OUTDUPL_FRAME_INFO::default();
|
||||
let mut resource: Option<IDXGIResource> = None;
|
||||
match self.dupl.AcquireNextFrame(timeout_ms, &mut info, &mut resource) {
|
||||
Ok(()) => {}
|
||||
Err(e) if e.code() == DXGI_ERROR_WAIT_TIMEOUT => return Ok(None),
|
||||
Err(e) if e.code() == DXGI_ERROR_ACCESS_LOST || e.code() == E_ACCESSDENIED => {
|
||||
return Err(anyhow!("duplication access lost"));
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
let resource = resource.ok_or_else(|| anyhow!("acquired null resource"))?;
|
||||
let tex: ID3D11Texture2D = resource.cast()?;
|
||||
|
||||
let mut tdesc = D3D11_TEXTURE2D_DESC::default();
|
||||
tex.GetDesc(&mut tdesc);
|
||||
|
||||
if self.staging.is_none() {
|
||||
let sdesc = D3D11_TEXTURE2D_DESC {
|
||||
Width: tdesc.Width,
|
||||
Height: tdesc.Height,
|
||||
MipLevels: 1,
|
||||
ArraySize: 1,
|
||||
Format: tdesc.Format,
|
||||
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 st: Option<ID3D11Texture2D> = None;
|
||||
self.device.CreateTexture2D(&sdesc, None, Some(&mut st))?;
|
||||
self.staging = st;
|
||||
}
|
||||
let staging = self.staging.clone().unwrap();
|
||||
self.context.CopyResource(&staging, &tex);
|
||||
|
||||
let mut mapped = D3D11_MAPPED_SUBRESOURCE::default();
|
||||
self.context
|
||||
.Map(&staging, 0, D3D11_MAP_READ, 0, Some(&mut mapped))?;
|
||||
|
||||
let w = tdesc.Width as usize;
|
||||
let h = tdesc.Height as usize;
|
||||
let pitch = mapped.RowPitch as usize;
|
||||
let src = mapped.pData as *const u8;
|
||||
let mut out = vec![0u8; w * h * 4];
|
||||
for y in 0..h {
|
||||
std::ptr::copy_nonoverlapping(
|
||||
src.add(y * pitch),
|
||||
out.as_mut_ptr().add(y * w * 4),
|
||||
w * 4,
|
||||
);
|
||||
}
|
||||
self.context.Unmap(&staging, 0);
|
||||
let _ = self.dupl.ReleaseFrame();
|
||||
Ok(Some(out))
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the CURRENT desktop image even when nothing has changed. A
|
||||
/// long-running duplication yields `WAIT_TIMEOUT` on a static screen (so
|
||||
/// `grab` returns `Ok(None)` indefinitely), which would leave a freshly
|
||||
/// connected viewer black because the keyframe flush has no pixels to
|
||||
/// encode. This gives DXGI a brief chance to produce the frame and otherwise
|
||||
/// falls back to a GDI BitBlt of the monitor rect, which always returns live
|
||||
/// pixels. Used only to seed a keyframe (viewer join / refresh), never in
|
||||
/// the per-frame hot path.
|
||||
pub fn grab_seed(&mut self) -> Result<Vec<u8>> {
|
||||
for _ in 0..2 {
|
||||
match self.grab(150) {
|
||||
Ok(Some(buf)) => return Ok(buf),
|
||||
Ok(None) => continue,
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
gdi_capture(self.left, self.top, self.width, self.height)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DxgiCapturer {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
let _ = self.dupl.ReleaseFrame();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Capture a monitor rect via GDI BitBlt as top-down BGRA. Unlike DXGI this
|
||||
/// always returns the CURRENT pixels (no "a frame changed" requirement), so it
|
||||
/// can seed the first keyframe for an idle/static screen. The host is
|
||||
/// per-monitor-DPI-aware, so `left/top/width/height` are physical virtual-
|
||||
/// desktop pixels that match the DXGI frame layout (BGRA, 4 bytes/pixel).
|
||||
pub(crate) fn gdi_capture(left: i32, top: i32, width: u32, height: u32) -> Result<Vec<u8>> {
|
||||
use windows::Win32::Graphics::Gdi::{
|
||||
BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, DeleteDC, DeleteObject, GetDC,
|
||||
GetDIBits, ReleaseDC, SelectObject, BITMAPINFO, BITMAPINFOHEADER, BI_RGB, DIB_RGB_COLORS,
|
||||
HGDIOBJ, SRCCOPY,
|
||||
};
|
||||
let (w, h) = (width as i32, height as i32);
|
||||
unsafe {
|
||||
let screen = GetDC(None);
|
||||
if screen.is_invalid() {
|
||||
return Err(anyhow!("GetDC(screen) failed"));
|
||||
}
|
||||
let mem = CreateCompatibleDC(screen);
|
||||
let bmp = CreateCompatibleBitmap(screen, w, h);
|
||||
if mem.is_invalid() || bmp.is_invalid() {
|
||||
ReleaseDC(None, screen);
|
||||
return Err(anyhow!("GDI DC/bitmap alloc failed"));
|
||||
}
|
||||
let old = SelectObject(mem, HGDIOBJ(bmp.0));
|
||||
let blt = BitBlt(mem, 0, 0, w, h, screen, left, top, SRCCOPY);
|
||||
let mut bi = BITMAPINFO {
|
||||
bmiHeader: BITMAPINFOHEADER {
|
||||
biSize: std::mem::size_of::<BITMAPINFOHEADER>() as u32,
|
||||
biWidth: w,
|
||||
biHeight: -h, // negative => top-down, matching the DXGI buffer
|
||||
biPlanes: 1,
|
||||
biBitCount: 32,
|
||||
biCompression: BI_RGB.0 as u32,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let mut out = vec![0u8; (width * height * 4) as usize];
|
||||
let lines = GetDIBits(
|
||||
mem,
|
||||
bmp,
|
||||
0,
|
||||
height,
|
||||
Some(out.as_mut_ptr() as *mut core::ffi::c_void),
|
||||
&mut bi,
|
||||
DIB_RGB_COLORS,
|
||||
);
|
||||
let _ = SelectObject(mem, old);
|
||||
let _ = DeleteObject(HGDIOBJ(bmp.0));
|
||||
let _ = DeleteDC(mem);
|
||||
ReleaseDC(None, screen);
|
||||
if blt.is_err() || lines == 0 {
|
||||
return Err(anyhow!("GDI BitBlt/GetDIBits failed"));
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach the calling (capture) thread to the desktop currently receiving input
|
||||
/// — the normal Default desktop, or the secure Winlogon desktop when the machine
|
||||
/// is locked, at the sign-in screen, or showing a UAC prompt. Running as SYSTEM
|
||||
/// (via the service), this lets GDI capture + SendInput reach the secure desktop
|
||||
/// so a remote viewer can sign in and unlock. Keeps exactly one desktop handle
|
||||
/// open, closing the previous one on each switch.
|
||||
static mut PREV_DESK: isize = 0;
|
||||
pub(crate) fn attach_input_desktop() -> bool {
|
||||
use windows::Win32::System::StationsAndDesktops::{
|
||||
CloseDesktop, OpenInputDesktop, SetThreadDesktop, DESKTOP_ACCESS_FLAGS,
|
||||
DESKTOP_CONTROL_FLAGS, HDESK,
|
||||
};
|
||||
unsafe {
|
||||
// GENERIC_ALL (0x10000000) → the kernel maps it to all desktop rights.
|
||||
match OpenInputDesktop(DESKTOP_CONTROL_FLAGS(0), false, DESKTOP_ACCESS_FLAGS(0x1000_0000)) {
|
||||
Ok(h) => {
|
||||
if SetThreadDesktop(h).is_ok() {
|
||||
if PREV_DESK != 0 {
|
||||
let _ = CloseDesktop(HDESK(PREV_DESK as *mut core::ffi::c_void));
|
||||
}
|
||||
PREV_DESK = h.0 as isize;
|
||||
true
|
||||
} else {
|
||||
let _ = CloseDesktop(h);
|
||||
false
|
||||
}
|
||||
}
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The all-monitors sentinel: `Capturer::open(MON_ALL)` captures the whole
|
||||
/// virtual desktop (every monitor stitched) via GDI. On the wire it is
|
||||
/// advertised + switched as monitor index `-1`.
|
||||
pub const MON_ALL: u32 = u32::MAX;
|
||||
|
||||
/// Bounding rect of the whole virtual desktop (all monitors), physical pixels:
|
||||
/// `(left, top, width, height)`. Width/height are rounded DOWN to even (the
|
||||
/// encoder needs even dimensions).
|
||||
pub fn virtual_desktop_rect() -> (i32, i32, u32, u32) {
|
||||
unsafe {
|
||||
let left = GetSystemMetrics(SM_XVIRTUALSCREEN);
|
||||
let top = GetSystemMetrics(SM_YVIRTUALSCREEN);
|
||||
let w = (GetSystemMetrics(SM_CXVIRTUALSCREEN).max(0) as u32) & !1;
|
||||
let h = (GetSystemMetrics(SM_CYVIRTUALSCREEN).max(0) as u32) & !1;
|
||||
(left, top, w, h)
|
||||
}
|
||||
}
|
||||
|
||||
/// All-monitors capturer: a GDI BitBlt of the whole virtual desktop each grab.
|
||||
/// Slower than DXGI (no GPU path, no "changed-only" wait), so it is an opt-in
|
||||
/// overview source, not the per-monitor default. Always returns live pixels.
|
||||
pub struct VirtualGdiCapturer {
|
||||
left: i32,
|
||||
top: i32,
|
||||
width: u32,
|
||||
height: u32,
|
||||
}
|
||||
|
||||
impl VirtualGdiCapturer {
|
||||
pub fn new() -> Result<Self> {
|
||||
let (left, top, width, height) = virtual_desktop_rect();
|
||||
if width == 0 || height == 0 {
|
||||
return Err(anyhow!("virtual desktop has zero size"));
|
||||
}
|
||||
Ok(Self { left, top, width, height })
|
||||
}
|
||||
pub fn width(&self) -> u32 {
|
||||
self.width
|
||||
}
|
||||
pub fn height(&self) -> u32 {
|
||||
self.height
|
||||
}
|
||||
pub fn origin(&self) -> (i32, i32) {
|
||||
(self.left, self.top)
|
||||
}
|
||||
/// GDI has no "unchanged" signal, so every grab returns the current frame.
|
||||
pub fn grab(&mut self, _timeout_ms: u32) -> Result<Option<Vec<u8>>> {
|
||||
Ok(Some(gdi_capture(self.left, self.top, self.width, self.height)?))
|
||||
}
|
||||
pub fn grab_seed(&mut self) -> Result<Vec<u8>> {
|
||||
gdi_capture(self.left, self.top, self.width, self.height)
|
||||
}
|
||||
}
|
||||
|
||||
/// A capture source: one DXGI-duplicated monitor, or the whole virtual desktop
|
||||
/// (GDI). The capture loop holds this and never branches on the variant — the
|
||||
/// enum delegates grab/seed/geometry, so `MON_ALL` "just works" as a monitor.
|
||||
pub enum Capturer {
|
||||
Dxgi(DxgiCapturer),
|
||||
VirtualAll(VirtualGdiCapturer),
|
||||
}
|
||||
|
||||
impl Capturer {
|
||||
/// Open monitor `mon` (`MON_ALL` => the whole virtual desktop via GDI).
|
||||
pub fn open(mon: u32) -> Result<Self> {
|
||||
if mon == MON_ALL {
|
||||
Ok(Capturer::VirtualAll(VirtualGdiCapturer::new()?))
|
||||
} else {
|
||||
Ok(Capturer::Dxgi(DxgiCapturer::new(mon)?))
|
||||
}
|
||||
}
|
||||
pub fn grab(&mut self, timeout_ms: u32) -> Result<Option<Vec<u8>>> {
|
||||
match self {
|
||||
Capturer::Dxgi(c) => c.grab(timeout_ms),
|
||||
Capturer::VirtualAll(c) => c.grab(timeout_ms),
|
||||
}
|
||||
}
|
||||
pub fn grab_seed(&mut self) -> Result<Vec<u8>> {
|
||||
match self {
|
||||
Capturer::Dxgi(c) => c.grab_seed(),
|
||||
Capturer::VirtualAll(c) => c.grab_seed(),
|
||||
}
|
||||
}
|
||||
pub fn width(&self) -> u32 {
|
||||
match self {
|
||||
Capturer::Dxgi(c) => c.width(),
|
||||
Capturer::VirtualAll(c) => c.width(),
|
||||
}
|
||||
}
|
||||
pub fn height(&self) -> u32 {
|
||||
match self {
|
||||
Capturer::Dxgi(c) => c.height(),
|
||||
Capturer::VirtualAll(c) => c.height(),
|
||||
}
|
||||
}
|
||||
pub fn origin(&self) -> (i32, i32) {
|
||||
match self {
|
||||
Capturer::Dxgi(c) => c.origin(),
|
||||
Capturer::VirtualAll(c) => c.origin(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod gdi_seed_tests {
|
||||
use super::*;
|
||||
|
||||
/// The GDI fallback must return a full-size buffer of live (non-black)
|
||||
/// desktop pixels — this is what rescues an idle screen from a black
|
||||
/// keyframe. Runs only with a real desktop session.
|
||||
#[test]
|
||||
fn gdi_capture_returns_live_pixels() {
|
||||
// A small rect at the virtual-desktop origin; every machine has pixels
|
||||
// here (taskbar / wallpaper), so a correct capture is never all-black.
|
||||
let w = 64u32;
|
||||
let h = 64u32;
|
||||
match gdi_capture(0, 0, w, h) {
|
||||
Ok(buf) => {
|
||||
assert_eq!(buf.len(), (w * h * 4) as usize, "wrong buffer size");
|
||||
let lit = buf.chunks_exact(4).filter(|p| p[0] | p[1] | p[2] != 0).count();
|
||||
assert!(lit > 0, "GDI capture returned an all-black buffer");
|
||||
}
|
||||
Err(e) => panic!("gdi_capture failed: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// The all-monitors capturer reports an even-sized virtual desktop and
|
||||
/// grabs a full-size live frame. Runs only with a real desktop session.
|
||||
#[test]
|
||||
fn virtual_capturer_grabs_full_desktop() {
|
||||
let (_l, _t, vw, vh) = virtual_desktop_rect();
|
||||
assert!(vw > 0 && vh > 0, "virtual desktop size is zero");
|
||||
assert!(vw % 2 == 0 && vh % 2 == 0, "virtual desktop size must be even");
|
||||
let mut cap = Capturer::open(MON_ALL).expect("open all-monitors capturer");
|
||||
assert_eq!((cap.width(), cap.height()), (vw, vh));
|
||||
let buf = cap.grab_seed().expect("grab virtual desktop");
|
||||
assert_eq!(buf.len(), (vw * vh * 4) as usize, "wrong buffer size");
|
||||
let lit = buf.chunks_exact(4).filter(|p| p[0] | p[1] | p[2] != 0).count();
|
||||
assert!(lit > 0, "all-monitors capture was all black");
|
||||
println!("VIRTUAL_CAP_OK {}x{} bytes={}", vw, vh, buf.len());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
//! BGRA -> NV12 color conversion (BT.601 limited range), the input most
|
||||
//! hardware H.264 encoders want. We sample chroma at the top-left of each 2x2
|
||||
//! block (fast; chroma is half-res anyway). Integer coefficients match the
|
||||
//! viewer decoder's NV12->RGB (BT.601 limited) so a capture->encode->decode
|
||||
//! roundtrip is color-stable.
|
||||
//!
|
||||
//! Output is tightly packed NV12: Y plane (w*h) then interleaved UV (w*h/2),
|
||||
//! total w*h*3/2 bytes, stride == width (callers keep width 2-aligned).
|
||||
|
||||
/// Convert a tightly-packed BGRA frame (`w*h*4`) into `dst` as NV12. `dst` must
|
||||
/// be at least `w*h*3/2` bytes. `w` and `h` must be even.
|
||||
pub fn bgra_to_nv12(bgra: &[u8], w: usize, h: usize, dst: &mut [u8]) {
|
||||
debug_assert!(w % 2 == 0 && h % 2 == 0);
|
||||
debug_assert!(bgra.len() >= w * h * 4);
|
||||
debug_assert!(dst.len() >= w * h * 3 / 2);
|
||||
|
||||
let (y_plane, uv_plane) = dst.split_at_mut(w * h);
|
||||
|
||||
// Luma: every pixel.
|
||||
for y in 0..h {
|
||||
let row = y * w;
|
||||
for x in 0..w {
|
||||
let p = (row + x) * 4;
|
||||
let b = bgra[p] as i32;
|
||||
let g = bgra[p + 1] as i32;
|
||||
let r = bgra[p + 2] as i32;
|
||||
// Y = 16 + (66R + 129G + 25B)/256
|
||||
let yy = (66 * r + 129 * g + 25 * b + 128) >> 8;
|
||||
y_plane[row + x] = (yy + 16).clamp(0, 255) as u8;
|
||||
}
|
||||
}
|
||||
|
||||
// Chroma: one U,V per 2x2 block, sampled at the block's top-left pixel.
|
||||
let cw = w / 2;
|
||||
for by in 0..h / 2 {
|
||||
let srow = (by * 2) * w;
|
||||
let drow = by * cw * 2;
|
||||
for bx in 0..cw {
|
||||
let p = (srow + bx * 2) * 4;
|
||||
let b = bgra[p] as i32;
|
||||
let g = bgra[p + 1] as i32;
|
||||
let r = bgra[p + 2] as i32;
|
||||
// U = 128 + (-38R - 74G + 112B)/256 ; V = 128 + (112R - 94G - 18B)/256
|
||||
let u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
|
||||
let v = ((112 * r - 94 * g - 18 * b + 128) >> 8) + 128;
|
||||
uv_plane[drow + bx * 2] = u.clamp(0, 255) as u8;
|
||||
uv_plane[drow + bx * 2 + 1] = v.clamp(0, 255) as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Bilinear downscale of a tightly-packed BGRA image `sw`x`sh` into `dst`
|
||||
/// (`dw`x`dh`, also tightly packed BGRA). Used for the quality "scale" lever so
|
||||
/// a constrained link can trade resolution for bandwidth. Only called when
|
||||
/// dw<sw (downscale), where bilinear is a good quality/speed trade-off.
|
||||
pub fn resize_bgra(src: &[u8], sw: usize, sh: usize, dst: &mut [u8], dw: usize, dh: usize) {
|
||||
if dw == 0 || dh == 0 || sw == 0 || sh == 0 {
|
||||
return;
|
||||
}
|
||||
let x_ratio = sw as f32 / dw as f32;
|
||||
let y_ratio = sh as f32 / dh as f32;
|
||||
for dy in 0..dh {
|
||||
let sy = (dy as f32 + 0.5) * y_ratio - 0.5;
|
||||
let sy = sy.clamp(0.0, (sh - 1) as f32);
|
||||
let y0 = sy.floor() as usize;
|
||||
let y1 = (y0 + 1).min(sh - 1);
|
||||
let fy = sy - y0 as f32;
|
||||
let r0 = y0 * sw * 4;
|
||||
let r1 = y1 * sw * 4;
|
||||
let drow = dy * dw * 4;
|
||||
for dx in 0..dw {
|
||||
let sx = (dx as f32 + 0.5) * x_ratio - 0.5;
|
||||
let sx = sx.clamp(0.0, (sw - 1) as f32);
|
||||
let x0 = sx.floor() as usize;
|
||||
let x1 = (x0 + 1).min(sw - 1);
|
||||
let fx = sx - x0 as f32;
|
||||
let o = drow + dx * 4;
|
||||
let (a, b, c, d) = (r0 + x0 * 4, r0 + x1 * 4, r1 + x0 * 4, r1 + x1 * 4);
|
||||
for ch in 0..4 {
|
||||
let p00 = src[a + ch] as f32;
|
||||
let p01 = src[b + ch] as f32;
|
||||
let p10 = src[c + ch] as f32;
|
||||
let p11 = src[d + ch] as f32;
|
||||
let top = p00 + (p01 - p00) * fx;
|
||||
let bot = p10 + (p11 - p10) * fx;
|
||||
dst[o + ch] = (top + (bot - top) * fy).round().clamp(0.0, 255.0) as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Target encode dimensions for a capture `w`x`h` at `scale` (0..1], rounded to
|
||||
/// even and clamped to a sane minimum.
|
||||
pub fn scaled_even(w: u32, h: u32, scale: f32) -> (u32, u32) {
|
||||
let s = scale.clamp(0.1, 1.0);
|
||||
let sw = (((w as f32 * s) as u32) & !1).max(64).min(w);
|
||||
let sh = (((h as f32 * s) as u32) & !1).max(64).min(h);
|
||||
(sw, sh)
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
//! Remote cursor tracking (RustDesk-style overlay): report the host cursor's
|
||||
//! position and shape to the viewer, sent on the control channel separately
|
||||
//! from video so an idle screen still shows a live, correctly-shaped cursor.
|
||||
//!
|
||||
//! Shapes are reported as CSS names for the standard system cursors (arrow,
|
||||
//! text, hand, resize, …) — which the viewer renders as crisp OS-native
|
||||
//! pointers — so there is no fragile per-frame bitmap extraction. Non-standard
|
||||
//! (app-custom) cursors fall back to the viewer's default pointer.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
GetCursorInfo, LoadCursorW, CURSORINFO, CURSOR_SHOWING, IDC_APPSTARTING, IDC_ARROW, IDC_CROSS,
|
||||
IDC_HAND, IDC_HELP, IDC_IBEAM, IDC_NO, IDC_SIZEALL, IDC_SIZENESW, IDC_SIZENS, IDC_SIZENWSE,
|
||||
IDC_SIZEWE, IDC_WAIT,
|
||||
};
|
||||
|
||||
/// `(hcursor_or_0, global_x, global_y)`. `0` means the cursor is hidden.
|
||||
pub fn cursor_state() -> Option<(u64, i32, i32)> {
|
||||
unsafe {
|
||||
let mut ci = CURSORINFO {
|
||||
cbSize: std::mem::size_of::<CURSORINFO>() as u32,
|
||||
..Default::default()
|
||||
};
|
||||
if GetCursorInfo(&mut ci).is_err() {
|
||||
return None;
|
||||
}
|
||||
let showing = (ci.flags.0 & CURSOR_SHOWING.0) != 0 && !ci.hCursor.is_invalid();
|
||||
let h = if showing { ci.hCursor.0 as u64 } else { 0 };
|
||||
Some((h, ci.ptScreenPos.x, ci.ptScreenPos.y))
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a cursor handle to a CSS cursor name by comparing it against the loaded
|
||||
/// standard system cursors. Returns `None` for app-custom cursors.
|
||||
pub fn cursor_css(hcursor: u64) -> Option<&'static str> {
|
||||
if hcursor == 0 {
|
||||
return None;
|
||||
}
|
||||
const STD: &[(windows::core::PCWSTR, &str)] = &[
|
||||
(IDC_ARROW, "default"),
|
||||
(IDC_IBEAM, "text"),
|
||||
(IDC_HAND, "pointer"),
|
||||
(IDC_CROSS, "crosshair"),
|
||||
(IDC_WAIT, "wait"),
|
||||
(IDC_APPSTARTING, "progress"),
|
||||
(IDC_SIZEALL, "move"),
|
||||
(IDC_SIZENS, "ns-resize"),
|
||||
(IDC_SIZEWE, "ew-resize"),
|
||||
(IDC_SIZENWSE, "nwse-resize"),
|
||||
(IDC_SIZENESW, "nesw-resize"),
|
||||
(IDC_NO, "not-allowed"),
|
||||
(IDC_HELP, "help"),
|
||||
];
|
||||
unsafe {
|
||||
for (idc, css) in STD {
|
||||
if let Ok(h) = LoadCursorW(None, *idc) {
|
||||
if h.0 as u64 == hcursor {
|
||||
return Some(css);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// A decoded custom-cursor bitmap ready for `cursor.shape` (PNG + hotspot), so
|
||||
/// app-specific cursors render on the viewer instead of falling back to the
|
||||
/// default pointer. RustDesk sends the raw cursor image the same way.
|
||||
pub struct CursorBitmap {
|
||||
pub png: Vec<u8>,
|
||||
pub w: u32,
|
||||
pub h: u32,
|
||||
pub hx: i32,
|
||||
pub hy: i32,
|
||||
}
|
||||
|
||||
/// Extract a (color) cursor's RGBA bitmap + hotspot and PNG-encode it. Returns
|
||||
/// None for monochrome cursors (rare for app-custom shapes; the standard mono
|
||||
/// ones are already mapped to CSS) and on any GDI failure, so the caller simply
|
||||
/// keeps the default pointer. Mirrors capture.rs's GetDIBits usage.
|
||||
pub fn cursor_bitmap(hcursor: u64) -> Option<CursorBitmap> {
|
||||
use windows::Win32::Foundation::HANDLE;
|
||||
use windows::Win32::Graphics::Gdi::{
|
||||
DeleteObject, GetDC, GetDIBits, GetObjectW, ReleaseDC, BITMAP, BITMAPINFO,
|
||||
BITMAPINFOHEADER, BI_RGB, DIB_RGB_COLORS, HBITMAP, HGDIOBJ,
|
||||
};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{GetIconInfo, HICON, ICONINFO};
|
||||
|
||||
if hcursor == 0 {
|
||||
return None;
|
||||
}
|
||||
unsafe {
|
||||
let mut ii = ICONINFO::default();
|
||||
if GetIconInfo(HICON(hcursor as *mut std::ffi::c_void), &mut ii).is_err() {
|
||||
return None;
|
||||
}
|
||||
// Always release the bitmaps GetIconInfo created for us.
|
||||
let cleanup = |ii: &ICONINFO| {
|
||||
if !ii.hbmColor.is_invalid() {
|
||||
let _ = DeleteObject(HGDIOBJ(ii.hbmColor.0));
|
||||
}
|
||||
if !ii.hbmMask.is_invalid() {
|
||||
let _ = DeleteObject(HGDIOBJ(ii.hbmMask.0));
|
||||
}
|
||||
};
|
||||
if ii.hbmColor.is_invalid() {
|
||||
// Monochrome cursor — fall back to the default pointer.
|
||||
cleanup(&ii);
|
||||
return None;
|
||||
}
|
||||
// Dimensions from the color bitmap.
|
||||
let mut bm = BITMAP::default();
|
||||
let got = GetObjectW(
|
||||
HGDIOBJ(ii.hbmColor.0),
|
||||
std::mem::size_of::<BITMAP>() as i32,
|
||||
Some(&mut bm as *mut _ as *mut std::ffi::c_void),
|
||||
);
|
||||
if got == 0 || bm.bmWidth <= 0 || bm.bmHeight <= 0 || bm.bmWidth > 512 || bm.bmHeight > 512 {
|
||||
cleanup(&ii);
|
||||
return None;
|
||||
}
|
||||
let (w, h) = (bm.bmWidth as u32, bm.bmHeight as u32);
|
||||
let read_bgra = |hbm: HBITMAP| -> Option<Vec<u8>> {
|
||||
let screen = GetDC(None);
|
||||
let mut bi = BITMAPINFO {
|
||||
bmiHeader: BITMAPINFOHEADER {
|
||||
biSize: std::mem::size_of::<BITMAPINFOHEADER>() as u32,
|
||||
biWidth: w as i32,
|
||||
biHeight: -(h as i32), // top-down
|
||||
biPlanes: 1,
|
||||
biBitCount: 32,
|
||||
biCompression: BI_RGB.0,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let mut buf = vec![0u8; (w * h * 4) as usize];
|
||||
let lines = GetDIBits(
|
||||
screen,
|
||||
hbm,
|
||||
0,
|
||||
h,
|
||||
Some(buf.as_mut_ptr() as *mut std::ffi::c_void),
|
||||
&mut bi,
|
||||
DIB_RGB_COLORS,
|
||||
);
|
||||
ReleaseDC(None, screen);
|
||||
if lines == 0 {
|
||||
None
|
||||
} else {
|
||||
Some(buf)
|
||||
}
|
||||
};
|
||||
let color = read_bgra(ii.hbmColor);
|
||||
// The AND mask (same 32bpp read): nonzero means "transparent" where the
|
||||
// color bitmap carries no alpha of its own.
|
||||
let mask = read_bgra(ii.hbmMask);
|
||||
let hotspot = (ii.xHotspot as i32, ii.yHotspot as i32);
|
||||
cleanup(&ii);
|
||||
let color = color?;
|
||||
|
||||
// BGRA -> RGBA. Use the color bitmap's own alpha if it has any; otherwise
|
||||
// derive it from the AND mask (mask pixel != 0 => transparent).
|
||||
let has_alpha = color.chunks_exact(4).any(|px| px[3] != 0);
|
||||
let mut rgba = vec![0u8; (w * h * 4) as usize];
|
||||
for (i, px) in color.chunks_exact(4).enumerate() {
|
||||
let (b, g, r) = (px[0], px[1], px[2]);
|
||||
let a = if has_alpha {
|
||||
px[3]
|
||||
} else if let Some(m) = &mask {
|
||||
// AND mask: a "set" (white) pixel is transparent.
|
||||
let mp = &m[i * 4..i * 4 + 4];
|
||||
if mp[0] | mp[1] | mp[2] != 0 {
|
||||
0
|
||||
} else {
|
||||
255
|
||||
}
|
||||
} else {
|
||||
255
|
||||
};
|
||||
let o = i * 4;
|
||||
rgba[o] = r;
|
||||
rgba[o + 1] = g;
|
||||
rgba[o + 2] = b;
|
||||
rgba[o + 3] = a;
|
||||
}
|
||||
let _ = HANDLE::default(); // keep the HANDLE import used across windows versions
|
||||
let png = encode_png(&rgba, w, h)?;
|
||||
Some(CursorBitmap { png, w, h, hx: hotspot.0, hy: hotspot.1 })
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode a top-down RGBA buffer to PNG bytes.
|
||||
pub(crate) fn encode_png(rgba: &[u8], w: u32, h: u32) -> Option<Vec<u8>> {
|
||||
let mut out = Vec::new();
|
||||
{
|
||||
let mut enc = png::Encoder::new(&mut out, w, h);
|
||||
enc.set_color(png::ColorType::Rgba);
|
||||
enc.set_depth(png::BitDepth::Eight);
|
||||
let mut writer = enc.write_header().ok()?;
|
||||
writer.write_image_data(rgba).ok()?;
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// Decode PNG bytes to top-down RGBA8 + (w, h). Used for clipboard images that
|
||||
/// arrive from a viewer, to hand arboard an `ImageData`.
|
||||
pub(crate) fn decode_png(bytes: &[u8]) -> Option<(Vec<u8>, u32, u32)> {
|
||||
let dec = png::Decoder::new(std::io::Cursor::new(bytes));
|
||||
let mut reader = dec.read_info().ok()?;
|
||||
let mut buf = vec![0u8; reader.output_buffer_size()];
|
||||
let info = reader.next_frame(&mut buf).ok()?;
|
||||
buf.truncate(info.buffer_size());
|
||||
// Normalize to RGBA8.
|
||||
let rgba = match (info.color_type, info.bit_depth) {
|
||||
(png::ColorType::Rgba, png::BitDepth::Eight) => buf,
|
||||
(png::ColorType::Rgb, png::BitDepth::Eight) => {
|
||||
let mut out = Vec::with_capacity((info.width * info.height * 4) as usize);
|
||||
for px in buf.chunks_exact(3) {
|
||||
out.extend_from_slice(&[px[0], px[1], px[2], 255]);
|
||||
}
|
||||
out
|
||||
}
|
||||
_ => return None, // uncommon formats -> skip (clipboard images are RGB/RGBA8)
|
||||
};
|
||||
Some((rgba, info.width, info.height))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Extract a real system cursor's bitmap (IDC_HAND is a color cursor on
|
||||
// modern Windows). In production cursor_css() catches the standard ones
|
||||
// first; here we call cursor_bitmap directly to prove the GDI extraction +
|
||||
// PNG encode produce a valid, sensibly-sized image.
|
||||
#[test]
|
||||
fn extracts_a_real_cursor_to_png() {
|
||||
let h = unsafe { LoadCursorW(None, IDC_HAND) };
|
||||
let h = match h {
|
||||
Ok(h) => h.0 as u64,
|
||||
Err(_) => return, // no interactive cursor resource (headless CI) — skip
|
||||
};
|
||||
if let Some(bmp) = cursor_bitmap(h) {
|
||||
assert!(bmp.w >= 8 && bmp.w <= 512, "width sane: {}", bmp.w);
|
||||
assert!(bmp.h >= 8 && bmp.h <= 512, "height sane: {}", bmp.h);
|
||||
// Valid PNG signature.
|
||||
assert_eq!(&bmp.png[..8], &[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]);
|
||||
// Round-trips through a decoder back to w*h*4 RGBA.
|
||||
let dec = png::Decoder::new(std::io::Cursor::new(&bmp.png));
|
||||
let mut reader = dec.read_info().expect("png header");
|
||||
let mut buf = vec![0u8; reader.output_buffer_size()];
|
||||
let info = reader.next_frame(&mut buf).expect("png frame");
|
||||
assert_eq!(info.width, bmp.w);
|
||||
assert_eq!(info.height, bmp.h);
|
||||
}
|
||||
// If extraction returned None (e.g. a monochrome IDC_HAND on some
|
||||
// builds), that's the documented safe fallback — not a failure.
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod clip_png_tests {
|
||||
use super::{decode_png, encode_png};
|
||||
|
||||
#[test]
|
||||
fn png_rgba_roundtrip_for_clipboard_images() {
|
||||
// A 2x2 RGBA image round-trips encode -> decode byte-for-byte.
|
||||
let w = 2u32;
|
||||
let h = 2u32;
|
||||
let rgba: Vec<u8> = vec![
|
||||
255, 0, 0, 255, 0, 255, 0, 255,
|
||||
0, 0, 255, 128, 10, 20, 30, 40,
|
||||
];
|
||||
let png = encode_png(&rgba, w, h).expect("encode");
|
||||
assert_eq!(&png[..8], &[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]);
|
||||
let (got, gw, gh) = decode_png(&png).expect("decode");
|
||||
assert_eq!((gw, gh), (w, h));
|
||||
assert_eq!(got, rgba, "clipboard image pixels survive the round-trip");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
//! H.264 encode via openh264 (pure-Rust-friendly, statically linked C core) —
|
||||
//! the portable software path that works on every Windows PC with no driver or
|
||||
//! GPU dependency. BGRA (from DXGI capture) is converted to I420 and encoded as
|
||||
//! Annex-B access units: keyframes carry SPS+PPS+IDR, deltas carry a P slice,
|
||||
//! exactly what the viewer's openh264/Media Foundation decoder expects.
|
||||
//!
|
||||
//! Hardware encoders (NVENC / QSV / AMF) layer in later behind this same
|
||||
//! `EncodedFrame` output; software-first guarantees the host always works.
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use openh264::encoder::{Encoder, EncoderConfig, FrameType, RateControlMode, UsageType};
|
||||
use openh264::formats::{BgraSliceU8, YUVBuffer};
|
||||
use openh264::OpenH264API;
|
||||
|
||||
/// One encoded access unit ready to hand to `core::host`, which prefixes the
|
||||
/// video channel + keyframe byte and encrypts it.
|
||||
pub struct EncodedFrame {
|
||||
pub data: Vec<u8>,
|
||||
pub keyframe: bool,
|
||||
}
|
||||
|
||||
/// A live H.264 encoder for one resolution. Resolution changes re-init the
|
||||
/// encoder automatically (openh264 does this on dimension change).
|
||||
pub struct H264Encoder {
|
||||
enc: Encoder,
|
||||
width: usize,
|
||||
height: usize,
|
||||
}
|
||||
|
||||
impl H264Encoder {
|
||||
pub fn new(width: usize, height: usize, fps: u32, bitrate_bps: u32) -> Result<Self> {
|
||||
let config = EncoderConfig::new()
|
||||
.set_bitrate_bps(bitrate_bps)
|
||||
.max_frame_rate(fps as f32)
|
||||
.rate_control_mode(RateControlMode::Bitrate)
|
||||
.usage_type(UsageType::ScreenContentRealTime)
|
||||
.enable_skip_frame(true)
|
||||
.set_multiple_thread_idc(0); // auto: use all cores
|
||||
let enc = Encoder::with_api_config(OpenH264API::from_source(), config)
|
||||
.map_err(|e| anyhow!("openh264 encoder init: {e}"))?;
|
||||
Ok(Self { enc, width, height })
|
||||
}
|
||||
|
||||
/// Encode one tightly-packed BGRA frame (`width*height*4`). Returns `None`
|
||||
/// when rate control skipped this frame (nothing to send this tick).
|
||||
pub fn encode(&mut self, bgra: &[u8]) -> Result<Option<EncodedFrame>> {
|
||||
let src = BgraSliceU8::new(bgra, (self.width, self.height));
|
||||
let yuv = YUVBuffer::from_rgb_source(src);
|
||||
let bs = self.enc.encode(&yuv).map_err(|e| anyhow!("encode: {e}"))?;
|
||||
let ft = bs.frame_type();
|
||||
if ft == FrameType::Skip || ft == FrameType::Invalid {
|
||||
return Ok(None);
|
||||
}
|
||||
let data = bs.to_vec();
|
||||
if data.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let keyframe = matches!(ft, FrameType::IDR | FrameType::I);
|
||||
Ok(Some(EncodedFrame { data, keyframe }))
|
||||
}
|
||||
|
||||
/// Force the next encoded frame to be a keyframe (viewer asked for a refresh
|
||||
/// after a drop, or a new viewer joined).
|
||||
pub fn request_keyframe(&mut self) {
|
||||
self.enc.force_intra_frame();
|
||||
}
|
||||
|
||||
pub fn width(&self) -> usize {
|
||||
self.width
|
||||
}
|
||||
pub fn height(&self) -> usize {
|
||||
self.height
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,486 @@
|
||||
//! Windows hardware H.264 encode via Media Foundation.
|
||||
//!
|
||||
//! Enumerates the PC's hardware H.264 encoder MFTs (NVIDIA NVENC, Intel Quick
|
||||
//! Sync, AMD VCE/AMF — whatever is present) and drives the async MFT with a
|
||||
//! deadlock-free event pump. BGRA from DXGI capture is converted to NV12 and
|
||||
//! fed in; Annex-B access units come out, with SPS/PPS guaranteed on keyframes.
|
||||
//!
|
||||
//! This is the "use any hardware encoder and any GPU accelerator on any Windows
|
||||
//! PC" path. If no hardware encoder is present or init fails, the caller falls
|
||||
//! back to the portable openh264 software encoder.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use crate::color::bgra_to_nv12;
|
||||
use crate::encode::EncodedFrame;
|
||||
use anyhow::{anyhow, Result};
|
||||
use std::ffi::c_void;
|
||||
|
||||
use windows::core::{Interface, VARIANT, GUID, PWSTR};
|
||||
use windows::Win32::Media::MediaFoundation::*;
|
||||
use windows::Win32::System::Com::{CoInitializeEx, CoTaskMemFree, COINIT_MULTITHREADED};
|
||||
|
||||
/// MF_MT_FRAME_SIZE / FRAME_RATE / PIXEL_ASPECT_RATIO are packed UINT64s
|
||||
/// ((hi << 32) | lo). windows-rs does not expose the inline MFSetAttribute*
|
||||
/// helpers, so we pack them ourselves.
|
||||
unsafe fn set_size(mt: &IMFMediaType, key: &GUID, w: u32, h: u32) -> Result<()> {
|
||||
mt.SetUINT64(key, ((w as u64) << 32) | (h as u64))?;
|
||||
Ok(())
|
||||
}
|
||||
unsafe fn set_ratio(mt: &IMFMediaType, key: &GUID, num: u32, den: u32) -> Result<()> {
|
||||
mt.SetUINT64(key, ((num as u64) << 32) | (den as u64))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// ~60 fps in 100-ns units; the stream is low-latency so stamps only need to be
|
||||
/// monotonic.
|
||||
const HNS_PER_FRAME: i64 = 166_667;
|
||||
|
||||
pub struct MfEncoder {
|
||||
transform: IMFTransform,
|
||||
events: IMFMediaEventGenerator,
|
||||
width: usize,
|
||||
height: usize,
|
||||
nv12: Vec<u8>,
|
||||
seq_header: Vec<u8>,
|
||||
provides_samples: bool,
|
||||
output_size: usize,
|
||||
need_credits: u32,
|
||||
pts: i64,
|
||||
hardware: bool,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl MfEncoder {
|
||||
pub fn new(width: usize, height: usize, fps: u32, bitrate_bps: u32) -> Result<Self> {
|
||||
unsafe {
|
||||
let _ = CoInitializeEx(None, COINIT_MULTITHREADED);
|
||||
MFStartup(MF_VERSION, MFSTARTUP_FULL).map_err(|e| anyhow!("MFStartup: {e}"))?;
|
||||
|
||||
// Enumerate hardware H.264 encoders: NV12 in, H264 out.
|
||||
let input_info = MFT_REGISTER_TYPE_INFO {
|
||||
guidMajorType: MFMediaType_Video,
|
||||
guidSubtype: MFVideoFormat_NV12,
|
||||
};
|
||||
let output_info = MFT_REGISTER_TYPE_INFO {
|
||||
guidMajorType: MFMediaType_Video,
|
||||
guidSubtype: MFVideoFormat_H264,
|
||||
};
|
||||
let flags =
|
||||
MFT_ENUM_FLAG_HARDWARE | MFT_ENUM_FLAG_ASYNCMFT | MFT_ENUM_FLAG_SORTANDFILTER;
|
||||
|
||||
let mut pp: *mut Option<IMFActivate> = std::ptr::null_mut();
|
||||
let mut count: u32 = 0;
|
||||
MFTEnumEx(
|
||||
MFT_CATEGORY_VIDEO_ENCODER,
|
||||
flags,
|
||||
Some(&input_info),
|
||||
Some(&output_info),
|
||||
&mut pp,
|
||||
&mut count,
|
||||
)
|
||||
.map_err(|e| anyhow!("MFTEnumEx: {e}"))?;
|
||||
|
||||
if count == 0 || pp.is_null() {
|
||||
return Err(anyhow!("no hardware H.264 encoder MFT found"));
|
||||
}
|
||||
let activations = std::slice::from_raw_parts(pp, count as usize);
|
||||
|
||||
let mut chosen: Option<(IMFTransform, String)> = None;
|
||||
for act in activations.iter().flatten() {
|
||||
let mut pw = PWSTR::null();
|
||||
let mut len = 0u32;
|
||||
let name = if act
|
||||
.GetAllocatedString(&MFT_FRIENDLY_NAME_Attribute, &mut pw, &mut len)
|
||||
.is_ok()
|
||||
&& !pw.is_null()
|
||||
{
|
||||
let s =
|
||||
String::from_utf16_lossy(std::slice::from_raw_parts(pw.0, len as usize));
|
||||
CoTaskMemFree(Some(pw.0 as *const c_void));
|
||||
s
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
if let Ok(t) = act.ActivateObject::<IMFTransform>() {
|
||||
chosen = Some((t, name));
|
||||
break;
|
||||
}
|
||||
}
|
||||
CoTaskMemFree(Some(pp as *const c_void));
|
||||
|
||||
let (transform, name) =
|
||||
chosen.ok_or_else(|| anyhow!("failed to activate any encoder MFT"))?;
|
||||
|
||||
// Async MFT: unlock before use, request low latency (1-in-1-out).
|
||||
if let Ok(attrs) = transform.GetAttributes() {
|
||||
let _ = attrs.SetUINT32(&MF_TRANSFORM_ASYNC_UNLOCK, 1);
|
||||
let _ = attrs.SetUINT32(&MF_LOW_LATENCY, 1);
|
||||
}
|
||||
|
||||
// OUTPUT type first (required for encoders): H.264.
|
||||
let out_t = MFCreateMediaType().map_err(|e| anyhow!("MFCreateMediaType(out): {e}"))?;
|
||||
out_t.SetGUID(&MF_MT_MAJOR_TYPE, &MFMediaType_Video)?;
|
||||
out_t.SetGUID(&MF_MT_SUBTYPE, &MFVideoFormat_H264)?;
|
||||
out_t.SetUINT32(&MF_MT_AVG_BITRATE, bitrate_bps)?;
|
||||
out_t.SetUINT32(&MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive.0 as u32)?;
|
||||
out_t.SetUINT32(&MF_MT_MPEG2_PROFILE, eAVEncH264VProfile_Main.0 as u32)?;
|
||||
set_size(&out_t, &MF_MT_FRAME_SIZE, width as u32, height as u32)?;
|
||||
set_ratio(&out_t, &MF_MT_FRAME_RATE, fps, 1)?;
|
||||
set_ratio(&out_t, &MF_MT_PIXEL_ASPECT_RATIO, 1, 1)?;
|
||||
transform
|
||||
.SetOutputType(0, &out_t, 0)
|
||||
.map_err(|e| anyhow!("SetOutputType(H264): {e}"))?;
|
||||
|
||||
// INPUT type: NV12.
|
||||
let in_t = MFCreateMediaType().map_err(|e| anyhow!("MFCreateMediaType(in): {e}"))?;
|
||||
in_t.SetGUID(&MF_MT_MAJOR_TYPE, &MFMediaType_Video)?;
|
||||
in_t.SetGUID(&MF_MT_SUBTYPE, &MFVideoFormat_NV12)?;
|
||||
in_t.SetUINT32(&MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive.0 as u32)?;
|
||||
set_size(&in_t, &MF_MT_FRAME_SIZE, width as u32, height as u32)?;
|
||||
set_ratio(&in_t, &MF_MT_FRAME_RATE, fps, 1)?;
|
||||
set_ratio(&in_t, &MF_MT_PIXEL_ASPECT_RATIO, 1, 1)?;
|
||||
transform
|
||||
.SetInputType(0, &in_t, 0)
|
||||
.map_err(|e| anyhow!("SetInputType(NV12): {e}"))?;
|
||||
|
||||
// Prefer ON-DEMAND keyframes (viewer join / refresh / QoS step) over a
|
||||
// tight periodic GOP: set a long GOP as a safety net and drive IDRs via
|
||||
// ICodecAPI force-keyframe (see request_keyframe), so the bitstream stays
|
||||
// lean on a static screen and a join gets a crisp keyframe without the
|
||||
// old whole-encoder re-init hitch. Best-effort (ignored if unsupported).
|
||||
if let Ok(api) = transform.cast::<ICodecAPI>() {
|
||||
let _ = api.SetValue(&CODECAPI_AVEncMPVGOPSize, &VARIANT::from(120u32));
|
||||
}
|
||||
|
||||
// SPS/PPS sequence header, to guarantee on keyframes.
|
||||
let mut seq_header = Vec::new();
|
||||
if let Ok(cur) = transform.GetOutputCurrentType(0) {
|
||||
if let Ok(sz) = cur.GetBlobSize(&MF_MT_MPEG_SEQUENCE_HEADER) {
|
||||
if sz > 0 {
|
||||
let mut buf = vec![0u8; sz as usize];
|
||||
if cur
|
||||
.GetBlob(&MF_MT_MPEG_SEQUENCE_HEADER, &mut buf, None)
|
||||
.is_ok()
|
||||
{
|
||||
seq_header = buf;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let info = transform
|
||||
.GetOutputStreamInfo(0)
|
||||
.map_err(|e| anyhow!("GetOutputStreamInfo: {e}"))?;
|
||||
let provides = (MFT_OUTPUT_STREAM_PROVIDES_SAMPLES.0
|
||||
| MFT_OUTPUT_STREAM_CAN_PROVIDE_SAMPLES.0) as u32;
|
||||
let provides_samples = (info.dwFlags & provides) != 0;
|
||||
let output_size = info.cbSize as usize;
|
||||
|
||||
let events: IMFMediaEventGenerator = transform
|
||||
.cast()
|
||||
.map_err(|e| anyhow!("MFT is not an event generator (not async?): {e}"))?;
|
||||
|
||||
transform.ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, 0)?;
|
||||
transform.ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, 0)?;
|
||||
|
||||
Ok(Self {
|
||||
transform,
|
||||
events,
|
||||
width,
|
||||
height,
|
||||
nv12: vec![0u8; width * height * 3 / 2],
|
||||
seq_header,
|
||||
provides_samples,
|
||||
output_size,
|
||||
need_credits: 0,
|
||||
pts: 0,
|
||||
hardware: true,
|
||||
name,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
pub fn is_hardware(&self) -> bool {
|
||||
self.hardware
|
||||
}
|
||||
pub fn width(&self) -> usize {
|
||||
self.width
|
||||
}
|
||||
pub fn height(&self) -> usize {
|
||||
self.height
|
||||
}
|
||||
|
||||
/// Encode one BGRA frame. Returns 0+ encoded AUs (an async encoder may run a
|
||||
/// frame or two of latency; outputs arrive in order).
|
||||
pub fn encode(&mut self, bgra: &[u8]) -> Result<Vec<EncodedFrame>> {
|
||||
unsafe {
|
||||
bgra_to_nv12(bgra, self.width, self.height, &mut self.nv12);
|
||||
let sample = self.make_input_sample()?;
|
||||
let mut out = Vec::new();
|
||||
|
||||
// Wait (blocking) until the MFT has asked for input at least once.
|
||||
while self.need_credits == 0 {
|
||||
let ev = self
|
||||
.events
|
||||
.GetEvent(MF_EVENT_FLAG_NONE)
|
||||
.map_err(|e| anyhow!("GetEvent(block): {e}"))?;
|
||||
self.handle_event(&ev, &mut out)?;
|
||||
}
|
||||
self.transform
|
||||
.ProcessInput(0, &sample, 0)
|
||||
.map_err(|e| anyhow!("ProcessInput: {e}"))?;
|
||||
self.need_credits -= 1;
|
||||
|
||||
// Collect whatever output is ready right now, without blocking.
|
||||
self.drain_ready(&mut out)?;
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush the pipeline tail (used by the bench / on shutdown).
|
||||
pub fn drain(&mut self) -> Result<Vec<EncodedFrame>> {
|
||||
unsafe {
|
||||
let mut out = Vec::new();
|
||||
self.transform
|
||||
.ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, 0)?;
|
||||
self.transform.ProcessMessage(MFT_MESSAGE_COMMAND_DRAIN, 0)?;
|
||||
loop {
|
||||
let ev = self
|
||||
.events
|
||||
.GetEvent(MF_EVENT_FLAG_NONE)
|
||||
.map_err(|e| anyhow!("GetEvent(drain): {e}"))?;
|
||||
let t = ev.GetType()?;
|
||||
if t == METransformDrainComplete.0 as u32 {
|
||||
break;
|
||||
}
|
||||
self.handle_event(&ev, &mut out)?;
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn drain_ready(&mut self, out: &mut Vec<EncodedFrame>) -> Result<()> {
|
||||
loop {
|
||||
match self.events.GetEvent(MF_EVENT_FLAG_NO_WAIT) {
|
||||
Ok(ev) => self.handle_event(&ev, out)?,
|
||||
Err(e) if e.code() == MF_E_NO_EVENTS_AVAILABLE => return Ok(()),
|
||||
Err(e) => return Err(anyhow!("GetEvent(nowait): {e}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn handle_event(
|
||||
&mut self,
|
||||
ev: &IMFMediaEvent,
|
||||
out: &mut Vec<EncodedFrame>,
|
||||
) -> Result<()> {
|
||||
let t = ev.GetType()?;
|
||||
if t == METransformNeedInput.0 as u32 {
|
||||
self.need_credits += 1;
|
||||
} else if t == METransformHaveOutput.0 as u32 {
|
||||
if let Some(f) = self.pull_output()? {
|
||||
out.push(f);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
unsafe fn make_input_sample(&mut self) -> Result<IMFSample> {
|
||||
let sample = MFCreateSample()?;
|
||||
let buffer = MFCreateMemoryBuffer(self.nv12.len() as u32)?;
|
||||
let mut ptr: *mut u8 = std::ptr::null_mut();
|
||||
buffer.Lock(&mut ptr, None, None)?;
|
||||
std::ptr::copy_nonoverlapping(self.nv12.as_ptr(), ptr, self.nv12.len());
|
||||
buffer.Unlock()?;
|
||||
buffer.SetCurrentLength(self.nv12.len() as u32)?;
|
||||
sample.AddBuffer(&buffer)?;
|
||||
let _ = sample.SetSampleTime(self.pts);
|
||||
let _ = sample.SetSampleDuration(HNS_PER_FRAME);
|
||||
self.pts += HNS_PER_FRAME;
|
||||
Ok(sample)
|
||||
}
|
||||
|
||||
/// Pick the H.264 output type (re)offered by the MFT and refresh the stored
|
||||
/// SPS/PPS. Hardware encoders finalize their output type after the first
|
||||
/// input and signal MF_E_TRANSFORM_STREAM_CHANGE; this renegotiates it.
|
||||
unsafe fn renegotiate_output(&mut self) -> Result<()> {
|
||||
let t = self
|
||||
.transform
|
||||
.GetOutputAvailableType(0, 0)
|
||||
.map_err(|e| anyhow!("GetOutputAvailableType(renegotiate): {e}"))?;
|
||||
self.transform
|
||||
.SetOutputType(0, &t, 0)
|
||||
.map_err(|e| anyhow!("SetOutputType(renegotiate): {e}"))?;
|
||||
if let Ok(sz) = t.GetBlobSize(&MF_MT_MPEG_SEQUENCE_HEADER) {
|
||||
if sz > 0 {
|
||||
let mut buf = vec![0u8; sz as usize];
|
||||
if t.GetBlob(&MF_MT_MPEG_SEQUENCE_HEADER, &mut buf, None).is_ok() {
|
||||
self.seq_header = buf;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Ok(info) = self.transform.GetOutputStreamInfo(0) {
|
||||
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(())
|
||||
}
|
||||
|
||||
unsafe fn pull_output(&mut self) -> Result<Option<EncodedFrame>> {
|
||||
let sample_slot: Option<IMFSample> = if self.provides_samples {
|
||||
None
|
||||
} else {
|
||||
let s = MFCreateSample()?;
|
||||
let b = MFCreateMemoryBuffer(self.output_size.max(self.width * self.height) as u32)?;
|
||||
s.AddBuffer(&b)?;
|
||||
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);
|
||||
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(|| anyhow!("ProcessOutput gave no sample"))?;
|
||||
let keyframe = s.GetUINT32(&MFSampleExtension_CleanPoint).unwrap_or(0) == 1;
|
||||
let buf = s.GetBufferByIndex(0)?;
|
||||
let mut ptr: *mut u8 = std::ptr::null_mut();
|
||||
let mut len = 0u32;
|
||||
buf.Lock(&mut ptr, None, Some(&mut len))?;
|
||||
let mut data = std::slice::from_raw_parts(ptr, len as usize).to_vec();
|
||||
buf.Unlock()?;
|
||||
|
||||
if keyframe && !self.seq_header.is_empty() && !starts_with_sps(&data) {
|
||||
let mut with_hdr = self.seq_header.clone();
|
||||
with_hdr.extend_from_slice(&data);
|
||||
data = with_hdr;
|
||||
}
|
||||
Ok(Some(EncodedFrame { data, keyframe }))
|
||||
}
|
||||
Err(e) if e.code() == MF_E_TRANSFORM_NEED_MORE_INPUT => Ok(None),
|
||||
Err(e) if e.code() == MF_E_TRANSFORM_STREAM_CHANGE => {
|
||||
// Encoder finalized its output format. Renegotiate; the pending
|
||||
// output arrives on the next METransformHaveOutput event (QSV
|
||||
// catastrophes on a synchronous retry here).
|
||||
self.renegotiate_output()?;
|
||||
Ok(None)
|
||||
}
|
||||
Err(e) => Err(anyhow!("ProcessOutput: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Force the next encoded frame to be a keyframe (IDR) via ICodecAPI
|
||||
/// CODECAPI_AVEncVideoForceKeyFrame — NVENC/QSV/AMF all honor it — instead of
|
||||
/// tearing down + rebuilding the whole MFT (the old approach, which stalled
|
||||
/// the stream on every viewer join / refresh / QoS bitrate step).
|
||||
pub fn request_keyframe(&mut self) {
|
||||
unsafe {
|
||||
if let Ok(api) = self.transform.cast::<ICodecAPI>() {
|
||||
let _ = api.SetValue(&CODECAPI_AVEncVideoForceKeyFrame, &VARIANT::from(1u32));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Prove the ICodecAPI force-keyframe path emits an IDR ON DEMAND without a
|
||||
/// rebuild. With the long (120) GOP set at init, the only natural keyframe in
|
||||
/// a short run is the initial IDR at frame 0; a keyframe appearing shortly
|
||||
/// after request_keyframe() (and well before frame 120) is therefore the
|
||||
/// forced one. Skips cleanly if the box has no hardware H.264 MFT.
|
||||
#[test]
|
||||
fn forced_keyframe_emits_idr_on_demand() {
|
||||
let (w, h) = (320usize, 240usize);
|
||||
let mut enc = match MfEncoder::new(w, h, 30, 1_000_000) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return, // no hardware encoder on this box — skip
|
||||
};
|
||||
let frame = vec![0x70u8; w * h * 4]; // solid BGRA
|
||||
let mut kf_idx = Vec::new();
|
||||
let mut idx = 0usize;
|
||||
// Prime ~40 frames: collect keyframe positions. With GOP=120 only frame 0
|
||||
// should be a keyframe.
|
||||
for _ in 0..40 {
|
||||
if let Ok(frames) = enc.encode(&frame) {
|
||||
for f in frames {
|
||||
if f.keyframe {
|
||||
kf_idx.push(idx);
|
||||
}
|
||||
idx += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
let before = idx;
|
||||
enc.request_keyframe();
|
||||
for _ in 0..12 {
|
||||
if let Ok(frames) = enc.encode(&frame) {
|
||||
for f in frames {
|
||||
if f.keyframe {
|
||||
kf_idx.push(idx);
|
||||
}
|
||||
idx += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Ok(frames) = enc.drain() {
|
||||
for f in frames {
|
||||
if f.keyframe {
|
||||
kf_idx.push(idx);
|
||||
}
|
||||
idx += 1;
|
||||
}
|
||||
}
|
||||
// If the encoder emitted essentially no frames (some CI GPUs stall under
|
||||
// a headless session), don't assert on timing — just ensure no panic.
|
||||
if before < 5 {
|
||||
return;
|
||||
}
|
||||
assert!(
|
||||
kf_idx.iter().any(|&p| p >= before),
|
||||
"forced keyframe should emit an IDR after request_keyframe(); \
|
||||
keyframes at {kf_idx:?}, forced after frame {before} (GOP=120)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Does this Annex-B buffer begin (after a start code) with an SPS NAL (type 7)?
|
||||
fn starts_with_sps(d: &[u8]) -> bool {
|
||||
let mut i = 0;
|
||||
if d.len() >= 4 && d[0] == 0 && d[1] == 0 && d[2] == 0 && d[3] == 1 {
|
||||
i = 4;
|
||||
} else if d.len() >= 3 && d[0] == 0 && d[1] == 0 && d[2] == 1 {
|
||||
i = 3;
|
||||
}
|
||||
d.get(i).map(|b| (b & 0x1f) == 7).unwrap_or(false)
|
||||
}
|
||||
|
||||
impl Drop for MfEncoder {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
let _ = self
|
||||
.transform
|
||||
.ProcessMessage(MFT_MESSAGE_NOTIFY_END_STREAMING, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The MFT is only ever touched from the one thread that owns this encoder; MF
|
||||
// objects are agile under MTA.
|
||||
unsafe impl Send for MfEncoder {}
|
||||
@@ -0,0 +1,218 @@
|
||||
//! Unified video encoder: prefer the platform hardware encoder (Media
|
||||
//! Foundation NVENC/QSV/AMF on Windows), fall back to the portable openh264
|
||||
//! software encoder, or — on constrained links — VP9 software (libvpx), which
|
||||
//! is far better per-bitrate than H.264. One `encode` call returns 0+ coded
|
||||
//! frames (the async hardware path can run a frame of latency; software is 1:1).
|
||||
|
||||
use crate::encode::{EncodedFrame, H264Encoder};
|
||||
use anyhow::Result;
|
||||
|
||||
pub enum VideoEncoder {
|
||||
#[cfg(windows)]
|
||||
Hardware(crate::encode_mf::MfEncoder),
|
||||
Software(H264Encoder),
|
||||
#[cfg(windows)]
|
||||
Vp9(Vp9State),
|
||||
}
|
||||
|
||||
/// VP9 encoder state. libvpx is stateless about wall-clock, so we track a start
|
||||
/// instant (ms PTS) and a pending-keyframe flag (the encode loop asks for a
|
||||
/// keyframe on viewer join / refresh / re-init).
|
||||
#[cfg(windows)]
|
||||
pub struct Vp9State {
|
||||
enc: netdesk_vpx::Encoder,
|
||||
want_key: bool,
|
||||
start: std::time::Instant,
|
||||
}
|
||||
|
||||
impl VideoEncoder {
|
||||
/// Build the best available H.264 encoder for this machine (hardware MF,
|
||||
/// else openh264 software). This is the default fast-link path.
|
||||
pub fn best(width: usize, height: usize, fps: u32, bitrate_bps: u32) -> Result<Self> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
match crate::encode_mf::MfEncoder::new(width, height, fps, bitrate_bps) {
|
||||
Ok(e) => {
|
||||
log::info!("encode: hardware \"{}\"", e.name());
|
||||
return Ok(VideoEncoder::Hardware(e));
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!("encode: no hardware encoder ({e}); using openh264 software");
|
||||
}
|
||||
}
|
||||
}
|
||||
let e = H264Encoder::new(width, height, fps, bitrate_bps)?;
|
||||
log::info!("encode: openh264 software");
|
||||
Ok(VideoEncoder::Software(e))
|
||||
}
|
||||
|
||||
/// Build an encoder for a specific codec (`proto::VideoCodec`): VP9 (1) uses
|
||||
/// libvpx software; anything else falls back to [`best`] (H.264). VP9 needs
|
||||
/// even dimensions, so callers must pass even `width`/`height` (the capture
|
||||
/// loop uses `color::scaled_even`).
|
||||
pub fn with_codec(
|
||||
codec_id: u8,
|
||||
i444: bool,
|
||||
width: usize,
|
||||
height: usize,
|
||||
fps: u32,
|
||||
bitrate_bps: u32,
|
||||
) -> Result<Self> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
if codec_id == 1 {
|
||||
let kbps = (bitrate_bps / 1000).max(50);
|
||||
match netdesk_vpx::Encoder::new_i444(
|
||||
netdesk_vpx::VpxCodec::Vp9,
|
||||
width as u32,
|
||||
height as u32,
|
||||
fps,
|
||||
kbps,
|
||||
i444,
|
||||
) {
|
||||
Ok(enc) => {
|
||||
log::info!(
|
||||
"encode: VP9 software (libvpx) {width}x{height} @ {kbps} kbps{}",
|
||||
if i444 { " 4:4:4" } else { "" }
|
||||
);
|
||||
return Ok(VideoEncoder::Vp9(Vp9State {
|
||||
enc,
|
||||
want_key: true, // first frame is a keyframe anyway
|
||||
start: std::time::Instant::now(),
|
||||
}));
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!("encode: VP9 init failed ({e}); using H.264");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = codec_id;
|
||||
Self::best(width, height, fps, bitrate_bps)
|
||||
}
|
||||
|
||||
pub fn encode(&mut self, bgra: &[u8]) -> Result<Vec<EncodedFrame>> {
|
||||
match self {
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Hardware(e) => e.encode(bgra),
|
||||
VideoEncoder::Software(e) => Ok(e.encode(bgra)?.into_iter().collect()),
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Vp9(s) => {
|
||||
let pts = s.start.elapsed().as_millis() as i64;
|
||||
let force_key = std::mem::take(&mut s.want_key);
|
||||
let frames = s.enc.encode_bgra(bgra, pts, force_key)?;
|
||||
Ok(frames
|
||||
.into_iter()
|
||||
.map(|f| EncodedFrame {
|
||||
data: f.data,
|
||||
keyframe: f.key,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn request_keyframe(&mut self) {
|
||||
match self {
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Hardware(e) => e.request_keyframe(),
|
||||
VideoEncoder::Software(e) => e.request_keyframe(),
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Vp9(s) => s.want_key = true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Codec id for the frames this encoder emits (`proto::VideoCodec`):
|
||||
/// 0 = H.264 (hardware MF / openh264 software), 1 = VP9.
|
||||
pub fn codec_id(&self) -> u8 {
|
||||
match self {
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Hardware(_) => 0,
|
||||
VideoEncoder::Software(_) => 0,
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Vp9(_) => 1,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_hardware(&self) -> bool {
|
||||
match self {
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Hardware(_) => true,
|
||||
VideoEncoder::Software(_) => false,
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Vp9(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn label(&self) -> String {
|
||||
match self {
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Hardware(e) => format!("hardware: {}", e.name()),
|
||||
VideoEncoder::Software(_) => "software: openh264".to_string(),
|
||||
#[cfg(windows)]
|
||||
VideoEncoder::Vp9(_) => "software: VP9 (libvpx)".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, windows))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use netdesk_core::decode::best_decoder_for;
|
||||
|
||||
// End-to-end through the real host encoder wrapper + the core decoder
|
||||
// selector: BGRA -> VideoEncoder(VP9) -> best_decoder_for(1) -> RGB, proving
|
||||
// the codec_id tag, pts/keyframe handling, and the BGRA<->RGB conversions
|
||||
// all line up (not just the raw netdesk-vpx roundtrip).
|
||||
#[test]
|
||||
fn vp9_encode_decode_via_video_encoder() {
|
||||
let (w, h) = (128usize, 96usize);
|
||||
let mut bgra = vec![0u8; w * h * 4];
|
||||
for j in 0..h {
|
||||
for i in 0..w {
|
||||
let p = (j * w + i) * 4;
|
||||
let (b, g, r) = if j < h / 2 && i < w / 2 {
|
||||
(30u8, 30, 200) // red
|
||||
} else if j < h / 2 {
|
||||
(30, 200, 30) // green
|
||||
} else if i < w / 2 {
|
||||
(200, 30, 30) // blue
|
||||
} else {
|
||||
(200, 200, 200) // gray
|
||||
};
|
||||
bgra[p] = b;
|
||||
bgra[p + 1] = g;
|
||||
bgra[p + 2] = r;
|
||||
bgra[p + 3] = 255;
|
||||
}
|
||||
}
|
||||
|
||||
let mut enc = VideoEncoder::with_codec(1, false, w, h, 30, 2_000_000).expect("vp9 encoder");
|
||||
assert_eq!(enc.codec_id(), 1, "must be VP9");
|
||||
let frames = enc.encode(&bgra).expect("encode");
|
||||
assert!(!frames.is_empty(), "no VP9 packets");
|
||||
assert!(frames[0].keyframe, "first VP9 frame must be a keyframe");
|
||||
|
||||
let mut dec = best_decoder_for(1);
|
||||
let mut got = None;
|
||||
for f in &frames {
|
||||
if let Ok(Some(fr)) = dec.decode(&f.data) {
|
||||
got = Some(fr);
|
||||
}
|
||||
}
|
||||
let fr = got.expect("decoded a frame");
|
||||
assert_eq!((fr.w, fr.h), (w, h), "decoded dims");
|
||||
|
||||
let sample = |x: usize, y: usize| {
|
||||
let p = (y * w + x) * 3;
|
||||
(fr.rgb[p], fr.rgb[p + 1], fr.rgb[p + 2]) // (R, G, B)
|
||||
};
|
||||
let approx = |a: u8, b: u8| (a as i32 - b as i32).abs() <= 48;
|
||||
let (r, g, b) = sample(w / 4, h / 4); // red quadrant
|
||||
assert!(
|
||||
approx(r, 200) && approx(g, 30) && approx(b, 30),
|
||||
"red quad decoded as r{r} g{g} b{b}"
|
||||
);
|
||||
println!("VP9_E2E_OK {w}x{h} key_bytes={}", frames[0].data.len());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
//! Host-side file manager: enumerate drives/roots and directory contents, and
|
||||
//! the chunked read/write used by download/upload. Items are reported in the
|
||||
//! viewer's wire shape `{n:name, t:"d"|"f", s:size, m:mtime}`.
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
/// Top-level roots: every present drive on Windows (`C:\`, `D:\`, …), plus the
|
||||
/// user's home, so the viewer's file panel has somewhere to start.
|
||||
pub fn roots() -> Vec<Value> {
|
||||
let mut items = Vec::new();
|
||||
#[cfg(windows)]
|
||||
{
|
||||
for c in b'A'..=b'Z' {
|
||||
let p = format!("{}:\\", c as char);
|
||||
if Path::new(&p).exists() {
|
||||
items.push(json!({ "n": p, "t": "d", "s": 0, "m": 0 }));
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(not(windows))]
|
||||
{
|
||||
items.push(json!({ "n": "/", "t": "d", "s": 0, "m": 0 }));
|
||||
}
|
||||
items
|
||||
}
|
||||
|
||||
/// List a directory. Entries are sorted folders-first then by name; hidden
|
||||
/// handling is left to the OS (we surface everything `read_dir` returns).
|
||||
pub fn list_dir(path: &str) -> std::io::Result<Vec<Value>> {
|
||||
let mut items = Vec::new();
|
||||
for entry in std::fs::read_dir(path)? {
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
let meta = match entry.metadata() {
|
||||
Ok(m) => m,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let is_dir = meta.is_dir();
|
||||
let size = if is_dir { 0 } else { meta.len() };
|
||||
let mtime = meta
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|t| t.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
items.push(json!({
|
||||
"n": name,
|
||||
"t": if is_dir { "d" } else { "f" },
|
||||
"s": size,
|
||||
"m": mtime,
|
||||
}));
|
||||
}
|
||||
items.sort_by(|a, b| {
|
||||
let ad = a.get("t").and_then(|x| x.as_str()) == Some("d");
|
||||
let bd = b.get("t").and_then(|x| x.as_str()) == Some("d");
|
||||
bd.cmp(&ad).then_with(|| {
|
||||
a.get("n")
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.to_ascii_lowercase()
|
||||
.cmp(&b.get("n").and_then(|x| x.as_str()).unwrap_or("").to_ascii_lowercase())
|
||||
})
|
||||
});
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
/// Recursively walk `root` and return every FILE under it as
|
||||
/// `{ "rel": "<forward-slash relative path>", "s": <size> }`, so a viewer can
|
||||
/// recreate the tree locally and download each entry. Directories are implied
|
||||
/// by the relative paths (the viewer mkdir -p's each parent). Bounded to avoid
|
||||
/// a runaway walk: at most `MAX_ENTRIES` files and `MAX_DEPTH` deep; symlinks
|
||||
/// are not followed. Returns (entries, truncated).
|
||||
pub fn walk_dir(root: &str) -> std::io::Result<(Vec<Value>, bool)> {
|
||||
const MAX_ENTRIES: usize = 50_000;
|
||||
const MAX_DEPTH: usize = 40;
|
||||
let root_path = PathBuf::from(root);
|
||||
let mut out = Vec::new();
|
||||
let mut truncated = false;
|
||||
// Iterative DFS: (dir, depth). rel paths are built from `root`.
|
||||
let mut stack = vec![(root_path.clone(), 0usize)];
|
||||
while let Some((dir, depth)) = stack.pop() {
|
||||
if depth > MAX_DEPTH {
|
||||
truncated = true;
|
||||
continue;
|
||||
}
|
||||
let rd = match std::fs::read_dir(&dir) {
|
||||
Ok(r) => r,
|
||||
Err(_) => continue, // unreadable dir -> skip, keep going
|
||||
};
|
||||
for entry in rd {
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
// symlink_metadata: do NOT follow symlinks (avoid cycles / escapes).
|
||||
let meta = match entry.metadata() {
|
||||
Ok(m) => m,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let path = entry.path();
|
||||
if meta.file_type().is_symlink() {
|
||||
continue;
|
||||
}
|
||||
if meta.is_dir() {
|
||||
stack.push((path, depth + 1));
|
||||
} else if meta.is_file() {
|
||||
if out.len() >= MAX_ENTRIES {
|
||||
truncated = true;
|
||||
break;
|
||||
}
|
||||
// Relative path with forward slashes (the viewer joins it under
|
||||
// the chosen local folder).
|
||||
let rel = path
|
||||
.strip_prefix(&root_path)
|
||||
.unwrap_or(&path)
|
||||
.to_string_lossy()
|
||||
.replace('\\', "/");
|
||||
out.push(json!({ "rel": rel, "s": meta.len() }));
|
||||
}
|
||||
}
|
||||
if out.len() >= MAX_ENTRIES {
|
||||
truncated = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
out.sort_by(|a, b| {
|
||||
a.get("rel").and_then(|x| x.as_str()).unwrap_or("")
|
||||
.cmp(b.get("rel").and_then(|x| x.as_str()).unwrap_or(""))
|
||||
});
|
||||
Ok((out, truncated))
|
||||
}
|
||||
|
||||
/// Where an in-progress upload is written before the atomic rename into place.
|
||||
pub fn part_path(final_path: &str) -> PathBuf {
|
||||
PathBuf::from(format!("{final_path}.netdesk.part"))
|
||||
}
|
||||
|
||||
/// Resolve an upload destination: a relative/empty dir drops onto the Desktop
|
||||
/// (TeamViewer-style default), mirroring the existing host.
|
||||
pub fn resolve_upload_path(path: &str) -> PathBuf {
|
||||
let p = Path::new(path);
|
||||
if p.is_absolute() {
|
||||
return p.to_path_buf();
|
||||
}
|
||||
// Relative: drop on the Desktop.
|
||||
let home = std::env::var("USERPROFILE")
|
||||
.or_else(|_| std::env::var("HOME"))
|
||||
.unwrap_or_else(|_| ".".into());
|
||||
Path::new(&home).join("Desktop").join(path)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn walk_dir_recurses_and_relativizes() {
|
||||
// Build a small temp tree: root/{a.txt, sub/b.txt, sub/deep/c.txt}.
|
||||
let base = std::env::temp_dir().join(format!("nd_walk_{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&base);
|
||||
std::fs::create_dir_all(base.join("sub").join("deep")).unwrap();
|
||||
std::fs::write(base.join("a.txt"), b"AA").unwrap();
|
||||
std::fs::write(base.join("sub").join("b.txt"), b"BBBB").unwrap();
|
||||
std::fs::write(base.join("sub").join("deep").join("c.txt"), b"CCCCCC").unwrap();
|
||||
|
||||
let (entries, truncated) = walk_dir(&base.to_string_lossy()).unwrap();
|
||||
assert!(!truncated);
|
||||
let rels: Vec<String> = entries
|
||||
.iter()
|
||||
.map(|e| e.get("rel").and_then(|x| x.as_str()).unwrap().to_string())
|
||||
.collect();
|
||||
assert_eq!(rels, vec!["a.txt", "sub/b.txt", "sub/deep/c.txt"]); // sorted, forward slashes
|
||||
// sizes carried through
|
||||
let size = |rel: &str| {
|
||||
entries.iter().find(|e| e.get("rel").and_then(|x| x.as_str()) == Some(rel))
|
||||
.and_then(|e| e.get("s").and_then(|x| x.as_u64())).unwrap()
|
||||
};
|
||||
assert_eq!(size("a.txt"), 2);
|
||||
assert_eq!(size("sub/b.txt"), 4);
|
||||
assert_eq!(size("sub/deep/c.txt"), 6);
|
||||
let _ = std::fs::remove_dir_all(&base);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,521 @@
|
||||
//! Input injection via SendInput — mouse move/buttons/wheel and keyboard, wire
|
||||
//! compatible with the existing NetDesk viewers (same convention the Python host
|
||||
//! uses): pointer coordinates arrive in virtual-desktop pixel space and are
|
||||
//! mapped to the 0..65535 absolute range over the whole virtual desktop; keys
|
||||
//! arrive as DOM-style codes ("KeyA", "Enter", "ArrowLeft") and are injected as
|
||||
//! scan codes so games and low-level apps see them correctly.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use windows::Win32::UI::Input::KeyboardAndMouse::{
|
||||
MapVirtualKeyW, SendInput, INPUT, INPUT_0, INPUT_KEYBOARD, INPUT_MOUSE, KEYBDINPUT,
|
||||
KEYBD_EVENT_FLAGS, KEYEVENTF_EXTENDEDKEY, KEYEVENTF_KEYUP, KEYEVENTF_SCANCODE,
|
||||
KEYEVENTF_UNICODE, MAPVK_VK_TO_VSC,
|
||||
MOUSEINPUT, MOUSE_EVENT_FLAGS, MOUSEEVENTF_ABSOLUTE, MOUSEEVENTF_HWHEEL, MOUSEEVENTF_LEFTDOWN,
|
||||
MOUSEEVENTF_LEFTUP, MOUSEEVENTF_MIDDLEDOWN, MOUSEEVENTF_MIDDLEUP, MOUSEEVENTF_MOVE,
|
||||
MOUSEEVENTF_RIGHTDOWN, MOUSEEVENTF_RIGHTUP, MOUSEEVENTF_VIRTUALDESK, MOUSEEVENTF_WHEEL,
|
||||
MOUSEEVENTF_XDOWN, MOUSEEVENTF_XUP, VIRTUAL_KEY,
|
||||
};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
GetSystemMetrics, SM_CXVIRTUALSCREEN, SM_CYVIRTUALSCREEN, SM_XVIRTUALSCREEN, SM_YVIRTUALSCREEN,
|
||||
};
|
||||
|
||||
/// Stateful input injector. Tracks the current mouse-button bitmask so it can
|
||||
/// emit the correct press/release transitions.
|
||||
pub struct Injector {
|
||||
buttons: u32,
|
||||
}
|
||||
|
||||
impl Default for Injector {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Injector {
|
||||
pub fn new() -> Self {
|
||||
Self { buttons: 0 }
|
||||
}
|
||||
|
||||
/// Inject one pointer update: move, then any button transitions, then wheel.
|
||||
pub fn mouse(&mut self, x: i32, y: i32, buttons: u32, wheel: f32) {
|
||||
self.mouse_move(x, y);
|
||||
if buttons != self.buttons {
|
||||
self.mouse_buttons(buttons, self.buttons);
|
||||
self.buttons = buttons;
|
||||
}
|
||||
if wheel != 0.0 {
|
||||
self.wheel(wheel);
|
||||
}
|
||||
}
|
||||
|
||||
/// Inject a RELATIVE pointer update (game / pointer-lock mode): move by
|
||||
/// (dx, dy) pixels using MOUSEEVENTF_MOVE WITHOUT ABSOLUTE, then any
|
||||
/// button transitions, then wheel. Apps that read raw mouse motion (FPS
|
||||
/// games, 3D/CAD) receive the deltas directly. Subject to the system
|
||||
/// pointer-speed/accel settings, like any relative motion.
|
||||
pub fn mouse_rel(&mut self, dx: i32, dy: i32, buttons: u32, wheel: f32) {
|
||||
if dx != 0 || dy != 0 {
|
||||
unsafe { send_mouse(dx, dy, 0, MOUSEEVENTF_MOVE) };
|
||||
}
|
||||
if buttons != self.buttons {
|
||||
self.mouse_buttons(buttons, self.buttons);
|
||||
self.buttons = buttons;
|
||||
}
|
||||
if wheel != 0.0 {
|
||||
self.wheel(wheel);
|
||||
}
|
||||
}
|
||||
|
||||
/// Inject a string as literal Unicode (translate / international-keyboard
|
||||
/// mode): each UTF-16 code unit is sent via KEYEVENTF_UNICODE, so the
|
||||
/// exact character is typed regardless of the host's keyboard layout —
|
||||
/// this is how AltGr combos, dead-key compositions and non-US layouts
|
||||
/// come through correctly (the viewer's OS already resolved them to a
|
||||
/// character). No VK/scan is used, so it is layout-independent.
|
||||
pub fn text(&self, s: &str) {
|
||||
for u in s.encode_utf16() {
|
||||
unsafe {
|
||||
send_key(VIRTUAL_KEY(0), u, KEYEVENTF_UNICODE);
|
||||
send_key(VIRTUAL_KEY(0), u, KEYEVENTF_UNICODE | KEYEVENTF_KEYUP);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Release any held buttons (viewer left / session end), so nothing stays
|
||||
/// stuck down.
|
||||
pub fn release_buttons(&mut self) {
|
||||
if self.buttons != 0 {
|
||||
self.mouse_buttons(0, self.buttons);
|
||||
self.buttons = 0;
|
||||
}
|
||||
}
|
||||
|
||||
fn mouse_move(&self, x: i32, y: i32) {
|
||||
unsafe {
|
||||
let vl = GetSystemMetrics(SM_XVIRTUALSCREEN);
|
||||
let vt = GetSystemMetrics(SM_YVIRTUALSCREEN);
|
||||
let vw = GetSystemMetrics(SM_CXVIRTUALSCREEN).max(1);
|
||||
let vh = GetSystemMetrics(SM_CYVIRTUALSCREEN).max(1);
|
||||
let nx = ((x - vl) as i64 * 65535 / (vw - 1).max(1) as i64) as i32;
|
||||
let ny = ((y - vt) as i64 * 65535 / (vh - 1).max(1) as i64) as i32;
|
||||
send_mouse(
|
||||
nx,
|
||||
ny,
|
||||
0,
|
||||
MOUSEEVENTF_MOVE | MOUSEEVENTF_ABSOLUTE | MOUSEEVENTF_VIRTUALDESK,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn mouse_buttons(&self, new_b: u32, old_b: u32) {
|
||||
let mut evts: Vec<(MOUSE_EVENT_FLAGS, i32)> = Vec::new();
|
||||
for (mask, down, up) in [
|
||||
(1u32, MOUSEEVENTF_LEFTDOWN, MOUSEEVENTF_LEFTUP),
|
||||
(2, MOUSEEVENTF_RIGHTDOWN, MOUSEEVENTF_RIGHTUP),
|
||||
(4, MOUSEEVENTF_MIDDLEDOWN, MOUSEEVENTF_MIDDLEUP),
|
||||
] {
|
||||
if new_b & mask != 0 && old_b & mask == 0 {
|
||||
evts.push((down, 0));
|
||||
}
|
||||
if new_b & mask == 0 && old_b & mask != 0 {
|
||||
evts.push((up, 0));
|
||||
}
|
||||
}
|
||||
// Side buttons: bit 8 = X1 (back), bit 16 = X2 (forward).
|
||||
for (mask, xb) in [(8u32, 0x0001i32), (16u32, 0x0002i32)] {
|
||||
if new_b & mask != 0 && old_b & mask == 0 {
|
||||
evts.push((MOUSEEVENTF_XDOWN, xb));
|
||||
}
|
||||
if new_b & mask == 0 && old_b & mask != 0 {
|
||||
evts.push((MOUSEEVENTF_XUP, xb));
|
||||
}
|
||||
}
|
||||
for (flags, data) in evts {
|
||||
unsafe { send_mouse(0, 0, (data as u32) << 16, flags) };
|
||||
}
|
||||
}
|
||||
|
||||
fn wheel(&self, rows: f32) {
|
||||
unsafe {
|
||||
send_mouse(0, 0, (rows * 120.0) as i32 as u32, MOUSEEVENTF_WHEEL);
|
||||
}
|
||||
}
|
||||
|
||||
/// Inject a key by DOM code + down/up. Unknown codes are ignored.
|
||||
pub fn key(&self, code: &str, down: bool) {
|
||||
let Some(vk) = code_to_vk(code) else {
|
||||
return;
|
||||
};
|
||||
unsafe {
|
||||
let scan = MapVirtualKeyW(vk as u32, MAPVK_VK_TO_VSC) as u16;
|
||||
let mut flags = KEYEVENTF_SCANCODE;
|
||||
if is_extended(code) {
|
||||
flags |= KEYEVENTF_EXTENDEDKEY;
|
||||
}
|
||||
if !down {
|
||||
flags |= KEYEVENTF_KEYUP;
|
||||
}
|
||||
if scan != 0 {
|
||||
send_key(VIRTUAL_KEY(0), scan, flags);
|
||||
} else {
|
||||
// No scan code (rare) — fall back to the virtual key.
|
||||
let mut f = KEYBD_EVENT_FLAGS(0);
|
||||
if !down {
|
||||
f |= KEYEVENTF_KEYUP;
|
||||
}
|
||||
send_key(VIRTUAL_KEY(vk), 0, f);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Horizontal wheel (not surfaced by the core event yet, but available).
|
||||
#[allow(dead_code)]
|
||||
fn hwheel(&self, rows: f32) {
|
||||
unsafe {
|
||||
send_mouse(0, 0, (rows * 120.0) as i32 as u32, MOUSEEVENTF_HWHEEL);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn send_mouse(dx: i32, dy: i32, mouse_data: u32, flags: MOUSE_EVENT_FLAGS) {
|
||||
let input = INPUT {
|
||||
r#type: INPUT_MOUSE,
|
||||
Anonymous: INPUT_0 {
|
||||
mi: MOUSEINPUT {
|
||||
dx,
|
||||
dy,
|
||||
mouseData: mouse_data,
|
||||
dwFlags: flags,
|
||||
time: 0,
|
||||
dwExtraInfo: 0,
|
||||
},
|
||||
},
|
||||
};
|
||||
SendInput(&[input], std::mem::size_of::<INPUT>() as i32);
|
||||
}
|
||||
|
||||
unsafe fn send_key(vk: VIRTUAL_KEY, scan: u16, flags: KEYBD_EVENT_FLAGS) {
|
||||
let input = INPUT {
|
||||
r#type: INPUT_KEYBOARD,
|
||||
Anonymous: INPUT_0 {
|
||||
ki: KEYBDINPUT {
|
||||
wVk: vk,
|
||||
wScan: scan,
|
||||
dwFlags: flags,
|
||||
time: 0,
|
||||
dwExtraInfo: 0,
|
||||
},
|
||||
},
|
||||
};
|
||||
SendInput(&[input], std::mem::size_of::<INPUT>() as i32);
|
||||
}
|
||||
|
||||
/// DOM `KeyboardEvent.code` -> Windows virtual-key. Mirrors the Python host so
|
||||
/// every existing viewer keeps working unchanged.
|
||||
fn code_to_vk(code: &str) -> Option<u16> {
|
||||
if let Some(c) = code.strip_prefix("Key") {
|
||||
let b = c.as_bytes();
|
||||
if b.len() == 1 && b[0].is_ascii_uppercase() {
|
||||
return Some(0x41 + (b[0] - b'A') as u16);
|
||||
}
|
||||
}
|
||||
if let Some(d) = code.strip_prefix("Digit") {
|
||||
let b = d.as_bytes();
|
||||
if b.len() == 1 && b[0].is_ascii_digit() {
|
||||
return Some(0x30 + (b[0] - b'0') as u16);
|
||||
}
|
||||
}
|
||||
if let Some(d) = code.strip_prefix("Numpad") {
|
||||
let b = d.as_bytes();
|
||||
if b.len() == 1 && b[0].is_ascii_digit() {
|
||||
return Some(0x60 + (b[0] - b'0') as u16);
|
||||
}
|
||||
}
|
||||
if let Some(n) = code.strip_prefix('F') {
|
||||
if let Ok(num) = n.parse::<u16>() {
|
||||
if (1..=24).contains(&num) {
|
||||
return Some(0x70 + num - 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(match code {
|
||||
"Backquote" => 0xC0,
|
||||
"Minus" => 0xBD,
|
||||
"Equal" => 0xBB,
|
||||
"BracketLeft" => 0xDB,
|
||||
"BracketRight" => 0xDD,
|
||||
"Backslash" => 0xDC,
|
||||
"Semicolon" => 0xBA,
|
||||
"Quote" => 0xDE,
|
||||
"Comma" => 0xBC,
|
||||
"Period" => 0xBE,
|
||||
"Slash" => 0xBF,
|
||||
"IntlBackslash" => 0xE2,
|
||||
"CapsLock" => 0x14,
|
||||
"Space" => 0x20,
|
||||
"Enter" | "NumpadEnter" => 0x0D,
|
||||
"Backspace" => 0x08,
|
||||
"Tab" => 0x09,
|
||||
"Escape" => 0x1B,
|
||||
"Delete" => 0x2E,
|
||||
"Insert" => 0x2D,
|
||||
"Home" => 0x24,
|
||||
"End" => 0x23,
|
||||
"PageUp" => 0x21,
|
||||
"PageDown" => 0x22,
|
||||
"ArrowLeft" => 0x25,
|
||||
"ArrowUp" => 0x26,
|
||||
"ArrowRight" => 0x27,
|
||||
"ArrowDown" => 0x28,
|
||||
"ControlLeft" => 0xA2,
|
||||
"ControlRight" => 0xA3,
|
||||
"AltLeft" => 0xA4,
|
||||
"AltRight" => 0xA5,
|
||||
"ShiftLeft" => 0xA0,
|
||||
"ShiftRight" => 0xA1,
|
||||
"MetaLeft" => 0x5B,
|
||||
"MetaRight" => 0x5C,
|
||||
"ContextMenu" => 0x5D,
|
||||
"PrintScreen" => 0x2C,
|
||||
"ScrollLock" => 0x91,
|
||||
"Pause" => 0x13,
|
||||
"NumLock" => 0x90,
|
||||
"NumpadMultiply" => 0x6A,
|
||||
"NumpadAdd" => 0x6B,
|
||||
"NumpadSubtract" => 0x6D,
|
||||
"NumpadDecimal" => 0x6E,
|
||||
"NumpadDivide" => 0x6F,
|
||||
// Media / volume keys (DOM UI Events code values) -> VK_* media codes, so
|
||||
// a viewer that grabs them (system-key hook) drives the host's volume,
|
||||
// playback and browser the same as sitting at the machine.
|
||||
"AudioVolumeMute" => 0xAD,
|
||||
"AudioVolumeDown" => 0xAE,
|
||||
"AudioVolumeUp" => 0xAF,
|
||||
"MediaTrackNext" => 0xB0,
|
||||
"MediaTrackPrevious" => 0xB1,
|
||||
"MediaStop" => 0xB2,
|
||||
"MediaPlayPause" => 0xB3,
|
||||
"LaunchMail" => 0xB4,
|
||||
"MediaSelect" | "LaunchMediaPlayer" => 0xB5,
|
||||
"LaunchApp1" => 0xB6,
|
||||
"LaunchApp2" => 0xB7,
|
||||
"BrowserBack" => 0xA6,
|
||||
"BrowserForward" => 0xA7,
|
||||
"BrowserRefresh" => 0xA8,
|
||||
"BrowserStop" => 0xA9,
|
||||
"BrowserSearch" => 0xAA,
|
||||
"BrowserFavorites" => 0xAB,
|
||||
"BrowserHome" => 0xAC,
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_extended(code: &str) -> bool {
|
||||
matches!(
|
||||
code,
|
||||
"ArrowRight"
|
||||
| "ArrowLeft"
|
||||
| "ArrowUp"
|
||||
| "ArrowDown"
|
||||
| "PageUp"
|
||||
| "PageDown"
|
||||
| "Home"
|
||||
| "End"
|
||||
| "Insert"
|
||||
| "Delete"
|
||||
| "NumLock"
|
||||
| "PrintScreen"
|
||||
| "Pause"
|
||||
| "ControlRight"
|
||||
| "AltRight"
|
||||
| "MetaLeft"
|
||||
| "MetaRight"
|
||||
| "ContextMenu"
|
||||
| "NumpadEnter"
|
||||
| "NumpadDivide"
|
||||
| "ScrollLock"
|
||||
// media / volume / browser keys are extended
|
||||
| "AudioVolumeMute"
|
||||
| "AudioVolumeDown"
|
||||
| "AudioVolumeUp"
|
||||
| "MediaTrackNext"
|
||||
| "MediaTrackPrevious"
|
||||
| "MediaStop"
|
||||
| "MediaPlayPause"
|
||||
| "LaunchMail"
|
||||
| "MediaSelect"
|
||||
| "LaunchMediaPlayer"
|
||||
| "LaunchApp1"
|
||||
| "LaunchApp2"
|
||||
| "BrowserBack"
|
||||
| "BrowserForward"
|
||||
| "BrowserRefresh"
|
||||
| "BrowserStop"
|
||||
| "BrowserSearch"
|
||||
| "BrowserFavorites"
|
||||
| "BrowserHome"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::code_to_vk;
|
||||
|
||||
#[test]
|
||||
fn maps_media_and_volume_keys() {
|
||||
// Volume + media-transport keys the viewer's system-key grab forwards.
|
||||
assert_eq!(code_to_vk("AudioVolumeMute"), Some(0xAD));
|
||||
assert_eq!(code_to_vk("AudioVolumeDown"), Some(0xAE));
|
||||
assert_eq!(code_to_vk("AudioVolumeUp"), Some(0xAF));
|
||||
assert_eq!(code_to_vk("MediaTrackNext"), Some(0xB0));
|
||||
assert_eq!(code_to_vk("MediaTrackPrevious"), Some(0xB1));
|
||||
assert_eq!(code_to_vk("MediaStop"), Some(0xB2));
|
||||
assert_eq!(code_to_vk("MediaPlayPause"), Some(0xB3));
|
||||
assert_eq!(code_to_vk("BrowserBack"), Some(0xA6));
|
||||
assert_eq!(code_to_vk("BrowserHome"), Some(0xAC));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_system_keys_the_hook_grabs() {
|
||||
// The WH_KEYBOARD_LL grab forwards these; the host must inject them.
|
||||
assert_eq!(code_to_vk("MetaLeft"), Some(0x5B)); // LWin
|
||||
assert_eq!(code_to_vk("MetaRight"), Some(0x5C)); // RWin
|
||||
assert_eq!(code_to_vk("PrintScreen"), Some(0x2C));
|
||||
assert_eq!(code_to_vk("Tab"), Some(0x09));
|
||||
assert_eq!(code_to_vk("Escape"), Some(0x1B));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn still_maps_ordinary_keys() {
|
||||
assert_eq!(code_to_vk("KeyA"), Some(0x41));
|
||||
assert_eq!(code_to_vk("Digit0"), Some(0x30));
|
||||
assert_eq!(code_to_vk("F5"), Some(0x74));
|
||||
assert_eq!(code_to_vk("Nonsense"), None);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#[cfg(test)]
|
||||
mod rel_mouse_tests {
|
||||
use super::*;
|
||||
use windows::Win32::Foundation::POINT;
|
||||
use windows::Win32::UI::WindowsAndMessaging::{GetCursorPos, SetCursorPos};
|
||||
|
||||
/// Real OS effect: a relative move must shift the ACTUAL cursor in the
|
||||
/// commanded direction by a sane magnitude. Mirrors the e2e GetCursorPos
|
||||
/// proof but as a focused, networkless unit test. The exact magnitude is
|
||||
/// subject to the system pointer-speed/accel curve, so we assert direction
|
||||
/// + a generous band around the sent (40, 30) rather than an exact match.
|
||||
#[test]
|
||||
fn mouse_rel_moves_cursor_relatively() {
|
||||
unsafe {
|
||||
let vw = GetSystemMetrics(SM_CXVIRTUALSCREEN).max(2);
|
||||
let vh = GetSystemMetrics(SM_CYVIRTUALSCREEN).max(2);
|
||||
// Park at center so the delta is not clamped at a screen edge.
|
||||
let _ = SetCursorPos(vw / 2, vh / 2);
|
||||
std::thread::sleep(std::time::Duration::from_millis(30));
|
||||
let mut before = POINT::default();
|
||||
let _ = GetCursorPos(&mut before);
|
||||
|
||||
let mut inj = Injector::new();
|
||||
inj.mouse_rel(40, 30, 0, 0.0);
|
||||
std::thread::sleep(std::time::Duration::from_millis(40));
|
||||
|
||||
let mut after = POINT::default();
|
||||
let _ = GetCursorPos(&mut after);
|
||||
let (ddx, ddy) = (after.x - before.x, after.y - before.y);
|
||||
println!(
|
||||
"REL_MOUSE_OK before=({},{}) after=({},{}) delta=({},{})",
|
||||
before.x, before.y, after.x, after.y, ddx, ddy
|
||||
);
|
||||
// "Enhance pointer precision" (mouse acceleration) scales the whole
|
||||
// relative vector by a system-dependent factor (seen 1.5x–3.2x on
|
||||
// this box run-to-run), so assert the DIRECTION + the 40:30 ≈ 4:3
|
||||
// ratio (preserved under uniform scaling) + sane bounds, never an
|
||||
// absolute magnitude. 3*ddx ≈ 4*ddy proves the sent vector landed.
|
||||
let ratio_ok = (3 * ddx - 4 * ddy).abs() <= (3 * ddx).max(1) / 4;
|
||||
assert!(
|
||||
ddx > 3 && ddy > 2 && ddx > ddy && ddx < 1000 && ratio_ok,
|
||||
"relative move delta ({ddx},{ddy}) must track the sent (40,30) vector \
|
||||
(right+down, ratio ~4:3), not stuck/absolute"
|
||||
);
|
||||
let _ = SetCursorPos(vw / 2, vh / 2);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#[cfg(test)]
|
||||
mod translate_kbd_tests {
|
||||
use super::*;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
use windows::Win32::Foundation::{LPARAM, LRESULT, WPARAM};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
CallNextHookEx, DispatchMessageW, PeekMessageW, SetWindowsHookExW, TranslateMessage,
|
||||
UnhookWindowsHookEx, HHOOK, KBDLLHOOKSTRUCT, MSG, PM_REMOVE, WH_KEYBOARD_LL, WM_KEYDOWN,
|
||||
};
|
||||
|
||||
// (vkCode, scanCode) captured on each keydown the LL hook sees.
|
||||
static CAP: OnceLock<Mutex<Vec<(u32, u32)>>> = OnceLock::new();
|
||||
|
||||
unsafe extern "system" fn cap_proc(code: i32, wparam: WPARAM, lparam: LPARAM) -> LRESULT {
|
||||
if code == 0 && wparam.0 as u32 == WM_KEYDOWN {
|
||||
let kb = &*(lparam.0 as *const KBDLLHOOKSTRUCT);
|
||||
if let Some(m) = CAP.get() {
|
||||
if let Ok(mut v) = m.lock() {
|
||||
v.push((kb.vkCode, kb.scanCode));
|
||||
}
|
||||
}
|
||||
}
|
||||
CallNextHookEx(HHOOK::default(), code, wparam, lparam)
|
||||
}
|
||||
|
||||
fn pump(ms: u64) {
|
||||
let end = Instant::now() + Duration::from_millis(ms);
|
||||
while Instant::now() < end {
|
||||
unsafe {
|
||||
let mut msg = MSG::default();
|
||||
while PeekMessageW(&mut msg, None, 0, 0, PM_REMOVE).as_bool() {
|
||||
let _ = TranslateMessage(&msg);
|
||||
DispatchMessageW(&msg);
|
||||
}
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
}
|
||||
|
||||
/// Real OS effect, focus-INDEPENDENT: Injector::text injects via
|
||||
/// KEYEVENTF_UNICODE, which a WH_KEYBOARD_LL hook sees as VK_PACKET (0xE7)
|
||||
/// with the UTF-16 code unit in `scanCode`. So "Ae-acute" ("A" + U+00E9)
|
||||
/// must produce keydown packets 0x41 then 0xE9 — proving the exact character
|
||||
/// is injected regardless of the host keyboard layout (the point of
|
||||
/// translate mode).
|
||||
#[test]
|
||||
fn text_injects_unicode_as_vk_packet() {
|
||||
const VK_PACKET: u32 = 0xE7;
|
||||
let _ = CAP.set(Mutex::new(Vec::new()));
|
||||
let hook = unsafe { SetWindowsHookExW(WH_KEYBOARD_LL, Some(cap_proc), None, 0) }
|
||||
.expect("install LL hook");
|
||||
|
||||
Injector::new().text("A\u{00e9}");
|
||||
pump(400);
|
||||
|
||||
unsafe {
|
||||
let _ = UnhookWindowsHookEx(hook);
|
||||
}
|
||||
let got = CAP.get().unwrap().lock().unwrap().clone();
|
||||
let packets: Vec<u32> = got
|
||||
.iter()
|
||||
.filter(|(vk, _)| *vk == VK_PACKET)
|
||||
.map(|(_, sc)| *sc)
|
||||
.collect();
|
||||
println!("TRANSLATE_KBD_OK keydown unicode packets = {:04X?}", packets);
|
||||
assert_eq!(
|
||||
packets,
|
||||
vec![0x41u32, 0xE9],
|
||||
"expected Unicode packets for 'A','e-acute'"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
//! System-key grab for the viewer — a low-level keyboard hook (WH_KEYBOARD_LL)
|
||||
//! that catches the keys Windows normally eats before an app can see them: the
|
||||
//! Windows keys, PrintScreen, and the Alt+Tab / Alt+Esc / Ctrl+Esc task-switch
|
||||
//! combos. While [`set_active`]`(true)` (the viewer calls this when a session is
|
||||
//! live AND its window is focused), each grabbed key is forwarded to the host
|
||||
//! (as a DOM `code`, through the install() channel) and the LOCAL action is
|
||||
//! SUPPRESSED — so pressing Alt+Tab or the Windows key switches/opens on the
|
||||
//! REMOTE machine instead of this one (RustDesk's "transmit system keys").
|
||||
//!
|
||||
//! The hook is process-global but only acts while active, so it is inert when no
|
||||
//! session is focused. It lives in this shared crate (next to input.rs's
|
||||
//! injection) so it is unit-testable; only the viewer installs it.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
|
||||
use std::sync::mpsc::{Receiver, Sender};
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
||||
use windows::Win32::Foundation::{LPARAM, LRESULT, WPARAM};
|
||||
use windows::Win32::UI::Input::KeyboardAndMouse::{GetAsyncKeyState, VK_CONTROL, VK_MENU};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
CallNextHookEx, SetWindowsHookExW, UnhookWindowsHookEx, HHOOK, KBDLLHOOKSTRUCT, WH_KEYBOARD_LL,
|
||||
WM_KEYDOWN, WM_SYSKEYDOWN,
|
||||
};
|
||||
|
||||
static ACTIVE: AtomicBool = AtomicBool::new(false);
|
||||
static SINK: OnceLock<Sender<(&'static str, bool)>> = OnceLock::new();
|
||||
static RX: OnceLock<Mutex<Receiver<(&'static str, bool)>>> = OnceLock::new();
|
||||
static HOOK: AtomicIsize = AtomicIsize::new(0);
|
||||
|
||||
/// Turn grabbing on/off. The viewer sets this to `session_active && focused`
|
||||
/// every frame; when false the hook passes everything through untouched.
|
||||
pub fn set_active(on: bool) {
|
||||
ACTIVE.store(on, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Whether the grab is currently installed + active.
|
||||
pub fn is_active() -> bool {
|
||||
ACTIVE.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// The DOM `code` the host's `code_to_vk` expects for a grabbed virtual key.
|
||||
fn system_code(vk: u16) -> Option<&'static str> {
|
||||
Some(match vk {
|
||||
0x5B => "MetaLeft", // Left Windows
|
||||
0x5C => "MetaRight", // Right Windows
|
||||
0x2C => "PrintScreen", // PrtScn
|
||||
0x09 => "Tab", // only grabbed with Alt down (Alt+Tab)
|
||||
0x1B => "Escape", // only grabbed with Alt/Ctrl down (Alt/Ctrl+Esc)
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
fn alt_down() -> bool {
|
||||
unsafe { (GetAsyncKeyState(VK_MENU.0 as i32) as u16 & 0x8000) != 0 }
|
||||
}
|
||||
fn ctrl_down() -> bool {
|
||||
unsafe { (GetAsyncKeyState(VK_CONTROL.0 as i32) as u16 & 0x8000) != 0 }
|
||||
}
|
||||
|
||||
/// Grab the Windows keys + PrtScn unconditionally; Tab/Esc only as a task-switch
|
||||
/// combo (so ordinary Tab/Esc still go through egui normally).
|
||||
fn should_grab(vk: u16) -> bool {
|
||||
match vk {
|
||||
0x5B | 0x5C | 0x2C => true,
|
||||
0x09 => alt_down(), // Alt+Tab
|
||||
0x1B => alt_down() || ctrl_down(), // Alt+Esc / Ctrl+Esc
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
unsafe extern "system" fn proc(code: i32, wparam: WPARAM, lparam: LPARAM) -> LRESULT {
|
||||
// HC_ACTION == 0: the struct is valid and we may inspect/act on it.
|
||||
if code == 0 && ACTIVE.load(Ordering::Relaxed) {
|
||||
let kb = &*(lparam.0 as *const KBDLLHOOKSTRUCT);
|
||||
let vk = kb.vkCode as u16;
|
||||
if should_grab(vk) {
|
||||
if let Some(code_str) = system_code(vk) {
|
||||
let down = matches!(wparam.0 as u32, WM_KEYDOWN | WM_SYSKEYDOWN);
|
||||
if let Some(tx) = SINK.get() {
|
||||
let _ = tx.send((code_str, down));
|
||||
}
|
||||
return LRESULT(1); // swallow it locally — it goes to the remote
|
||||
}
|
||||
}
|
||||
}
|
||||
CallNextHookEx(HHOOK::default(), code, wparam, lparam)
|
||||
}
|
||||
|
||||
/// Install the global low-level keyboard hook once (owns its own channel; drain
|
||||
/// grabbed keys with [`poll`]). Returns false if the hook could not be
|
||||
/// installed. Must be called from a thread that pumps messages (the viewer's UI
|
||||
/// thread does); idempotent.
|
||||
pub fn install() -> bool {
|
||||
if HOOK.load(Ordering::Relaxed) != 0 {
|
||||
return true; // already installed
|
||||
}
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
let _ = SINK.set(tx);
|
||||
let _ = RX.set(Mutex::new(rx));
|
||||
match unsafe { SetWindowsHookExW(WH_KEYBOARD_LL, Some(proc), None, 0) } {
|
||||
Ok(h) => {
|
||||
HOOK.store(h.0 as isize, Ordering::Relaxed);
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!("keyhook: SetWindowsHookExW failed: {e}");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drain the keys grabbed since the last call, as `(dom_code, down)`. The viewer
|
||||
/// calls this each frame and forwards each to the host via `Session::key`.
|
||||
pub fn poll() -> Vec<(&'static str, bool)> {
|
||||
let mut out = Vec::new();
|
||||
if let Some(m) = RX.get() {
|
||||
if let Ok(rx) = m.lock() {
|
||||
while let Ok(ev) = rx.try_recv() {
|
||||
out.push(ev);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Remove the hook (process exit also removes it).
|
||||
pub fn uninstall() {
|
||||
let raw = HOOK.swap(0, Ordering::Relaxed);
|
||||
if raw != 0 {
|
||||
unsafe {
|
||||
let _ = UnhookWindowsHookEx(HHOOK(raw as *mut core::ffi::c_void));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::{Duration, Instant};
|
||||
use windows::Win32::UI::Input::KeyboardAndMouse::{
|
||||
SendInput, INPUT, INPUT_0, INPUT_KEYBOARD, KEYBDINPUT, KEYBD_EVENT_FLAGS, KEYEVENTF_KEYUP,
|
||||
VIRTUAL_KEY,
|
||||
};
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
DispatchMessageW, PeekMessageW, TranslateMessage, MSG, PM_REMOVE,
|
||||
};
|
||||
|
||||
unsafe fn send_vk(vk: u16, down: bool) {
|
||||
let mut flags = KEYBD_EVENT_FLAGS(0);
|
||||
if !down {
|
||||
flags |= KEYEVENTF_KEYUP;
|
||||
}
|
||||
let input = INPUT {
|
||||
r#type: INPUT_KEYBOARD,
|
||||
Anonymous: INPUT_0 {
|
||||
ki: KEYBDINPUT {
|
||||
wVk: VIRTUAL_KEY(vk),
|
||||
wScan: 0,
|
||||
dwFlags: flags,
|
||||
time: 0,
|
||||
dwExtraInfo: 0,
|
||||
},
|
||||
},
|
||||
};
|
||||
SendInput(&[input], std::mem::size_of::<INPUT>() as i32);
|
||||
}
|
||||
|
||||
fn pump(ms: u64) {
|
||||
let end = Instant::now() + Duration::from_millis(ms);
|
||||
while Instant::now() < end {
|
||||
unsafe {
|
||||
let mut msg = MSG::default();
|
||||
while PeekMessageW(&mut msg, None, 0, 0, PM_REMOVE).as_bool() {
|
||||
let _ = TranslateMessage(&msg);
|
||||
DispatchMessageW(&msg);
|
||||
}
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
|
||||
// One test (the hook + ACTIVE are process-global, so we must not run two in
|
||||
// parallel). Injected keys ARE delivered to WH_KEYBOARD_LL hooks, so we prove
|
||||
// the grab end to end: while ACTIVE, injecting the Left Windows key is
|
||||
// forwarded as "MetaLeft" down+up (and suppressed locally, returning 1);
|
||||
// while INACTIVE, nothing is grabbed (keys pass through to the app).
|
||||
#[test]
|
||||
fn grabs_system_keys_only_while_active() {
|
||||
assert!(install(), "hook should install");
|
||||
|
||||
// ACTIVE: the Windows key is grabbed + forwarded.
|
||||
set_active(true);
|
||||
unsafe {
|
||||
send_vk(0x5B, true);
|
||||
send_vk(0x5B, false);
|
||||
}
|
||||
pump(400);
|
||||
let got = poll();
|
||||
assert!(
|
||||
got.contains(&("MetaLeft", true)) && got.contains(&("MetaLeft", false)),
|
||||
"LWin should be grabbed as MetaLeft down+up while active; got {got:?}"
|
||||
);
|
||||
|
||||
// INACTIVE: the same key is passed through, grabbed nothing.
|
||||
set_active(false);
|
||||
unsafe {
|
||||
send_vk(0x5B, true);
|
||||
send_vk(0x5B, false);
|
||||
}
|
||||
pump(200);
|
||||
uninstall();
|
||||
assert!(poll().is_empty(), "nothing should be grabbed while inactive");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
//! netdesk-host — the native Rust host engine: DXGI screen capture + H.264
|
||||
//! encode + SendInput injection, wired to `netdesk_core::host::run_host`. This
|
||||
//! replaces the Python host so the Windows app is one unified Rust binary
|
||||
//! (host + viewer). Developed by Akronlogic.
|
||||
|
||||
#[cfg(windows)]
|
||||
pub mod capture;
|
||||
|
||||
pub mod encode;
|
||||
pub mod color;
|
||||
#[cfg(windows)]
|
||||
pub mod encode_mf;
|
||||
|
||||
#[cfg(windows)]
|
||||
pub mod input;
|
||||
|
||||
#[cfg(windows)]
|
||||
pub mod sys;
|
||||
|
||||
#[cfg(windows)]
|
||||
pub mod cursor;
|
||||
pub mod encoder;
|
||||
#[cfg(windows)]
|
||||
pub mod audio;
|
||||
pub mod files;
|
||||
#[cfg(windows)]
|
||||
pub mod qos;
|
||||
pub mod run;
|
||||
#[cfg(windows)]
|
||||
pub mod telemetry;
|
||||
#[cfg(windows)]
|
||||
pub mod service;
|
||||
pub mod wol;
|
||||
#[cfg(windows)]
|
||||
pub mod term;
|
||||
|
||||
#[cfg(windows)]
|
||||
pub mod annot;
|
||||
pub mod privacy;
|
||||
|
||||
#[cfg(windows)]
|
||||
pub mod keyhook;
|
||||
|
||||
/// Make this process per-monitor DPI aware so DXGI capture returns real
|
||||
/// physical pixels (not a virtualized, scaled desktop). Safe to call once at
|
||||
/// startup; a no-op on non-Windows.
|
||||
#[cfg(windows)]
|
||||
pub fn set_dpi_aware() {
|
||||
use windows::Win32::UI::HiDpi::{
|
||||
SetProcessDpiAwarenessContext, DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2,
|
||||
};
|
||||
unsafe {
|
||||
let _ = SetProcessDpiAwarenessContext(DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
pub fn set_dpi_aware() {}
|
||||
@@ -0,0 +1,131 @@
|
||||
//! Driver-free privacy mode. A fullscreen, topmost, opaque-black overlay marked
|
||||
//! `WDA_EXCLUDEFROMCAPTURE`, so:
|
||||
//! - the LOCAL monitor(s) show solid black — a bystander at the host sees
|
||||
//! nothing of what the remote operator is doing;
|
||||
//! - the host's OWN screen capture (DXGI Desktop Duplication + the GDI BitBlt
|
||||
//! fallback) EXCLUDES the overlay, so the remote operator still sees and
|
||||
//! controls the real desktop behind it.
|
||||
//! The overlay is also click-through + non-activating (WS_EX_TRANSPARENT |
|
||||
//! WS_EX_NOACTIVATE), so the operator's injected SendInput passes straight
|
||||
//! through to the real windows; blocking the LOCAL keyboard/mouse is the
|
||||
//! separate `block_input` toggle. No virtual-display driver is needed — verified
|
||||
//! by the `netdesk-privacy-probe` bin (WDA_EXCLUDEFROMCAPTURE hid a window from
|
||||
//! both DXGI and GDI capture, 197600 -> 0 magenta pixels).
|
||||
//!
|
||||
//! The overlay lives on its own thread with a message pump (a window must be
|
||||
//! serviced on its creating thread); dropping [`Privacy`] tears it down and
|
||||
//! restores the screen.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::thread::JoinHandle;
|
||||
use std::time::Duration;
|
||||
|
||||
use windows::core::w;
|
||||
use windows::Win32::Foundation::{COLORREF, HINSTANCE, HWND, LPARAM, LRESULT, WPARAM};
|
||||
use windows::Win32::Graphics::Gdi::CreateSolidBrush;
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
CreateWindowExW, DefWindowProcW, DestroyWindow, DispatchMessageW, GetSystemMetrics,
|
||||
PeekMessageW, RegisterClassW, SetLayeredWindowAttributes, SetWindowDisplayAffinity, ShowWindow,
|
||||
TranslateMessage, LWA_ALPHA, MSG, PM_REMOVE, SM_CXVIRTUALSCREEN, SM_CYVIRTUALSCREEN,
|
||||
SM_XVIRTUALSCREEN, SM_YVIRTUALSCREEN, SW_SHOWNA, WDA_EXCLUDEFROMCAPTURE, WNDCLASSW,
|
||||
WS_EX_LAYERED, WS_EX_NOACTIVATE, WS_EX_TOOLWINDOW, WS_EX_TOPMOST, WS_EX_TRANSPARENT, WS_POPUP,
|
||||
WS_VISIBLE,
|
||||
};
|
||||
|
||||
unsafe extern "system" fn wndproc(h: HWND, m: u32, w: WPARAM, l: LPARAM) -> LRESULT {
|
||||
DefWindowProcW(h, m, w, l)
|
||||
}
|
||||
|
||||
/// A live privacy-mode blackout. Drop it to restore the screen.
|
||||
pub struct Privacy {
|
||||
stop: Arc<AtomicBool>,
|
||||
thread: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl Privacy {
|
||||
/// Raise the blackout overlay (returns once the window thread has started).
|
||||
pub fn start() -> Privacy {
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let s2 = stop.clone();
|
||||
let thread = std::thread::Builder::new()
|
||||
.name("netdesk-privacy".into())
|
||||
.spawn(move || unsafe { run(s2) })
|
||||
.ok();
|
||||
Privacy { stop, thread }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Privacy {
|
||||
fn drop(&mut self) {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
if let Some(t) = self.thread.take() {
|
||||
let _ = t.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn run(stop: Arc<AtomicBool>) {
|
||||
let hinst = HINSTANCE::default();
|
||||
let brush = CreateSolidBrush(COLORREF(0x0000_0000)); // solid black
|
||||
let class = w!("NdPrivacyBlackout");
|
||||
let wc = WNDCLASSW {
|
||||
lpfnWndProc: Some(wndproc),
|
||||
hInstance: hinst,
|
||||
lpszClassName: class,
|
||||
hbrBackground: brush,
|
||||
..Default::default()
|
||||
};
|
||||
// 0 if the class is already registered (a prior toggle this process) — fine.
|
||||
let _ = RegisterClassW(&wc);
|
||||
|
||||
// Cover the whole virtual desktop (all monitors).
|
||||
let x = GetSystemMetrics(SM_XVIRTUALSCREEN);
|
||||
let y = GetSystemMetrics(SM_YVIRTUALSCREEN);
|
||||
let cx = GetSystemMetrics(SM_CXVIRTUALSCREEN).max(1);
|
||||
let cy = GetSystemMetrics(SM_CYVIRTUALSCREEN).max(1);
|
||||
|
||||
let hwnd = match CreateWindowExW(
|
||||
WS_EX_TOPMOST | WS_EX_TOOLWINDOW | WS_EX_LAYERED | WS_EX_TRANSPARENT | WS_EX_NOACTIVATE,
|
||||
class,
|
||||
w!("nd-privacy"),
|
||||
WS_POPUP | WS_VISIBLE,
|
||||
x,
|
||||
y,
|
||||
cx,
|
||||
cy,
|
||||
None,
|
||||
None,
|
||||
hinst,
|
||||
None,
|
||||
) {
|
||||
Ok(h) => h,
|
||||
Err(e) => {
|
||||
log::warn!("privacy: window create failed: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
// Opaque (alpha 255) so the overlay is solid black on the physical screen.
|
||||
let _ = SetLayeredWindowAttributes(hwnd, COLORREF(0), 255, LWA_ALPHA);
|
||||
// Exclude from THIS host's own capture so the remote still sees the desktop.
|
||||
if let Err(e) = SetWindowDisplayAffinity(hwnd, WDA_EXCLUDEFROMCAPTURE) {
|
||||
log::warn!("privacy: SetWindowDisplayAffinity failed: {e}");
|
||||
}
|
||||
// Show without activating (keyboard focus stays on the real foreground app,
|
||||
// so injected keystrokes still land there).
|
||||
let _ = ShowWindow(hwnd, SW_SHOWNA);
|
||||
log::info!("privacy: blackout up ({cx}x{cy} @ {x},{y}); excluded from capture, click-through");
|
||||
|
||||
let mut msg = MSG::default();
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
while PeekMessageW(&mut msg, None, 0, 0, PM_REMOVE).as_bool() {
|
||||
let _ = TranslateMessage(&msg);
|
||||
DispatchMessageW(&msg);
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(30));
|
||||
}
|
||||
let _ = DestroyWindow(hwnd);
|
||||
log::info!("privacy: blackout down; screen restored");
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
//! Adaptive quality controller — the smoothness engine. A faithful adaptation
|
||||
//! of RustDesk's `src/server/video_qos.rs` (delay-threshold ABR + FPS/scale
|
||||
//! backoff with congestion confirmation and gradual recovery), distilled to a
|
||||
//! pure, deterministic state machine the host feeds with viewer telemetry.
|
||||
//!
|
||||
//! The viewer reports, ~once a second, how far behind it is — a `delay` (how
|
||||
//! long frames queue/decode in front of it) plus the transport RTT and packet
|
||||
//! loss. [`Qos::on_sample`] turns that into the next `(fps, bitrate, scale)` for
|
||||
//! the live encoder (the same knobs the manual `set.quality` lever drives). The
|
||||
//! user's manual quality is the CEILING; the controller only ever trims BELOW it
|
||||
//! under congestion and restores back toward it when the link is healthy.
|
||||
//!
|
||||
//! Priority mirrors RustDesk: **reduce bitrate first** (keep motion smooth),
|
||||
//! then frame-rate (down to a floor), then resolution scale as the last resort;
|
||||
//! recover in the reverse-ish order (frame-rate, then bitrate, then full
|
||||
//! resolution). Congestion must be CONFIRMED (several bad samples in a row)
|
||||
//! before cutting, so a single jittery reply does not thrash the stream.
|
||||
|
||||
/// A good-vs-bad network boundary: frames more than this far behind mean the
|
||||
/// link cannot keep up (RustDesk's `DELAY_THRESHOLD_150MS`).
|
||||
pub const DELAY_THRESHOLD_MS: u32 = 150;
|
||||
/// Packet loss above this also counts as congestion.
|
||||
pub const LOSS_THRESHOLD: f32 = 0.05;
|
||||
/// Consecutive bad samples required before a reduction (confirm, don't thrash).
|
||||
const BAD_CONFIRM: u8 = 3;
|
||||
/// Consecutive good samples required before a step of recovery.
|
||||
const GOOD_CONFIRM: u8 = 3;
|
||||
/// Multiplicative bitrate cut on MILD congestion (keep 80%). Severe congestion
|
||||
/// cuts harder (see [`severity_cut`]) so the stream sheds load fast when it is
|
||||
/// badly behind, then trims gently near the threshold — RustDesk's graduated
|
||||
/// backoff rather than one fixed step.
|
||||
const BITRATE_CUT: f32 = 0.8;
|
||||
|
||||
/// Bitrate retention factor for how far past the delay/loss threshold a sample
|
||||
/// is: gentle (0.8) just over the line, aggressive (0.6) when badly congested.
|
||||
fn severity_cut(delay_ms: u32, loss: f32, threshold_ms: u32) -> f32 {
|
||||
let t = threshold_ms.max(1);
|
||||
if delay_ms >= t * 3 || loss >= 0.15 {
|
||||
0.6
|
||||
} else if delay_ms >= t * 2 || loss >= 0.10 {
|
||||
0.7
|
||||
} else {
|
||||
BITRATE_CUT
|
||||
}
|
||||
}
|
||||
/// FPS step down/up.
|
||||
const FPS_STEP: u32 = 5;
|
||||
/// Scale step down/up (resolution), last resort.
|
||||
const SCALE_STEP: f32 = 0.1;
|
||||
|
||||
/// Startup frame-rate the ramp begins at (RustDesk-style slow start).
|
||||
pub const INIT_FPS: u32 = 15;
|
||||
/// How long the startup ramp takes to climb from [`INIT_FPS`] to the target.
|
||||
pub const RAMP_MS: u64 = 3000;
|
||||
|
||||
/// Startup FPS ramp: a fresh session begins at [`INIT_FPS`] and climbs linearly
|
||||
/// to the target over the first [`RAMP_MS`], so the initial keyframe + buffer
|
||||
/// fill don't arrive as a full-rate burst that stutters before the decode /
|
||||
/// jitter pipeline has warmed up. After [`RAMP_MS`] it returns the target
|
||||
/// unchanged (no limiting), and a target already at/below [`INIT_FPS`] is never
|
||||
/// slowed. Pure + deterministic; the capture loop caps its pacing fps with this.
|
||||
pub fn startup_fps(elapsed_ms: u64, target_fps: u32) -> u32 {
|
||||
let target = target_fps.max(1);
|
||||
if elapsed_ms >= RAMP_MS || target <= INIT_FPS {
|
||||
return target;
|
||||
}
|
||||
// Linear climb INIT_FPS -> target across RAMP_MS.
|
||||
let span = (target - INIT_FPS) as u64;
|
||||
let add = (span * elapsed_ms / RAMP_MS) as u32;
|
||||
(INIT_FPS + add).min(target)
|
||||
}
|
||||
|
||||
/// Static bounds + the user's manual ceiling. `ceil_*` are the manual quality
|
||||
/// (the maximum the controller will restore to).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct QosConfig {
|
||||
pub ceil_fps: u32,
|
||||
pub ceil_bitrate: u32, // kbps
|
||||
pub min_fps: u32,
|
||||
pub min_bitrate: u32, // kbps
|
||||
pub min_scale: f32,
|
||||
pub delay_threshold_ms: u32,
|
||||
}
|
||||
|
||||
impl QosConfig {
|
||||
/// Defaults around a 1080p/30fps/6Mbps session, floors that stay usable on a
|
||||
/// bad link (5fps, 400kbps, half resolution).
|
||||
pub fn new(ceil_fps: u32, ceil_bitrate: u32) -> Self {
|
||||
Self {
|
||||
ceil_fps: ceil_fps.max(1),
|
||||
ceil_bitrate: ceil_bitrate.max(1),
|
||||
min_fps: 5,
|
||||
min_bitrate: 400,
|
||||
min_scale: 0.5,
|
||||
delay_threshold_ms: DELAY_THRESHOLD_MS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the controller wants the encoder set to right now.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct QosTarget {
|
||||
pub fps: u32,
|
||||
pub bitrate: u32, // kbps
|
||||
pub scale: f32,
|
||||
}
|
||||
|
||||
/// The adaptive controller. Pure + deterministic: feed samples, read targets.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Qos {
|
||||
cfg: QosConfig,
|
||||
fps: u32,
|
||||
bitrate: u32,
|
||||
scale: f32,
|
||||
bad_run: u8,
|
||||
good_run: u8,
|
||||
}
|
||||
|
||||
impl Qos {
|
||||
/// Start at the user's ceiling (full quality) — the controller only trims
|
||||
/// below it under measured congestion.
|
||||
pub fn new(cfg: QosConfig) -> Self {
|
||||
Self {
|
||||
fps: cfg.ceil_fps,
|
||||
bitrate: cfg.ceil_bitrate,
|
||||
scale: 1.0,
|
||||
cfg,
|
||||
bad_run: 0,
|
||||
good_run: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// The current target.
|
||||
pub fn target(&self) -> QosTarget {
|
||||
QosTarget {
|
||||
fps: self.fps,
|
||||
bitrate: self.bitrate,
|
||||
scale: self.scale,
|
||||
}
|
||||
}
|
||||
|
||||
/// Raise the ceiling (the user moved the manual quality lever). Current
|
||||
/// values are re-clamped under the new ceiling; a raised ceiling lets the
|
||||
/// controller recover higher.
|
||||
pub fn set_ceiling(&mut self, ceil_fps: u32, ceil_bitrate: u32) {
|
||||
self.cfg.ceil_fps = ceil_fps.max(1);
|
||||
self.cfg.ceil_bitrate = ceil_bitrate.max(1);
|
||||
self.fps = self.fps.min(self.cfg.ceil_fps);
|
||||
self.bitrate = self.bitrate.min(self.cfg.ceil_bitrate);
|
||||
}
|
||||
|
||||
/// Feed one telemetry sample (delay in ms, packet loss 0..1). Returns the
|
||||
/// new target. Congestion is confirmed over [`BAD_CONFIRM`] samples before a
|
||||
/// cut; recovery takes [`GOOD_CONFIRM`] good samples per step.
|
||||
pub fn on_sample(&mut self, delay_ms: u32, loss: f32) -> QosTarget {
|
||||
let congested = delay_ms >= self.cfg.delay_threshold_ms || loss >= LOSS_THRESHOLD;
|
||||
if congested {
|
||||
self.good_run = 0;
|
||||
self.bad_run = self.bad_run.saturating_add(1);
|
||||
if self.bad_run >= BAD_CONFIRM {
|
||||
let cut = severity_cut(delay_ms, loss, self.cfg.delay_threshold_ms);
|
||||
self.reduce(cut);
|
||||
self.bad_run = 0; // require fresh confirmation before the next cut
|
||||
}
|
||||
} else {
|
||||
self.bad_run = 0;
|
||||
self.good_run = self.good_run.saturating_add(1);
|
||||
if self.good_run >= GOOD_CONFIRM {
|
||||
self.grow();
|
||||
self.good_run = 0;
|
||||
}
|
||||
}
|
||||
self.target()
|
||||
}
|
||||
|
||||
/// Trim one step: bitrate first (keep motion), then fps, then resolution.
|
||||
/// `cut` is the bitrate retention factor for this sample's severity.
|
||||
fn reduce(&mut self, cut: f32) {
|
||||
if self.bitrate > self.cfg.min_bitrate {
|
||||
let next = (self.bitrate as f32 * cut) as u32;
|
||||
self.bitrate = next.max(self.cfg.min_bitrate);
|
||||
} else if self.fps > self.cfg.min_fps {
|
||||
self.fps = self.fps.saturating_sub(FPS_STEP).max(self.cfg.min_fps);
|
||||
} else if self.scale > self.cfg.min_scale + 1e-3 {
|
||||
self.scale = (self.scale - SCALE_STEP).max(self.cfg.min_scale);
|
||||
}
|
||||
}
|
||||
|
||||
/// Restore one step toward the ceiling: full resolution, then frame-rate,
|
||||
/// then bitrate.
|
||||
fn grow(&mut self) {
|
||||
if self.scale < 1.0 - 1e-3 {
|
||||
self.scale = (self.scale + SCALE_STEP).min(1.0);
|
||||
} else if self.fps < self.cfg.ceil_fps {
|
||||
self.fps = (self.fps + FPS_STEP).min(self.cfg.ceil_fps);
|
||||
} else if self.bitrate < self.cfg.ceil_bitrate {
|
||||
// +≤ a quarter of the ceiling per step, like RustDesk's capped ramp.
|
||||
let step = (self.cfg.ceil_bitrate / 4).max(200);
|
||||
self.bitrate = (self.bitrate + step).min(self.cfg.ceil_bitrate);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn qos() -> Qos {
|
||||
Qos::new(QosConfig::new(30, 6000))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starts_at_ceiling() {
|
||||
let t = qos().target();
|
||||
assert_eq!((t.fps, t.bitrate, t.scale), (30, 6000, 1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn healthy_link_stays_at_ceiling() {
|
||||
let mut q = qos();
|
||||
for _ in 0..50 {
|
||||
q.on_sample(20, 0.0);
|
||||
}
|
||||
let t = q.target();
|
||||
assert_eq!((t.fps, t.bitrate, t.scale), (30, 6000, 1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_bad_sample_does_not_cut() {
|
||||
let mut q = qos();
|
||||
q.on_sample(300, 0.0); // one bad reply
|
||||
assert_eq!(q.target().bitrate, 6000, "must confirm congestion before cutting");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confirmed_congestion_cuts_bitrate_first() {
|
||||
let mut q = qos();
|
||||
for _ in 0..BAD_CONFIRM {
|
||||
q.on_sample(400, 0.0);
|
||||
}
|
||||
let t = q.target();
|
||||
assert!(t.bitrate < 6000, "bitrate cut first");
|
||||
assert_eq!(t.fps, 30, "fps untouched while bitrate has room");
|
||||
assert_eq!(t.scale, 1.0, "scale untouched while bitrate has room");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn heavy_congestion_walks_bitrate_then_fps_then_scale() {
|
||||
let mut q = qos();
|
||||
for _ in 0..400 {
|
||||
q.on_sample(500, 0.1); // relentless congestion
|
||||
}
|
||||
let t = q.target();
|
||||
assert_eq!(t.bitrate, 400, "bitrate at floor");
|
||||
assert_eq!(t.fps, 5, "fps at floor");
|
||||
assert!((t.scale - 0.5).abs() < 1e-3, "scale at floor");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn severe_congestion_cuts_harder_than_mild() {
|
||||
// One confirmed cut at severe delay (>=3x threshold) must drop bitrate
|
||||
// further than one confirmed cut just over the threshold.
|
||||
let mut mild = qos();
|
||||
for _ in 0..BAD_CONFIRM {
|
||||
mild.on_sample(DELAY_THRESHOLD_MS + 5, 0.0); // ~threshold -> 0.8
|
||||
}
|
||||
let mut severe = qos();
|
||||
for _ in 0..BAD_CONFIRM {
|
||||
severe.on_sample(DELAY_THRESHOLD_MS * 4, 0.0); // deep -> 0.6
|
||||
}
|
||||
assert_eq!(mild.target().bitrate, (6000.0 * 0.8) as u32);
|
||||
assert_eq!(severe.target().bitrate, (6000.0 * 0.6) as u32);
|
||||
assert!(
|
||||
severe.target().bitrate < mild.target().bitrate,
|
||||
"severe congestion should shed more bitrate per cut"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loss_alone_triggers_congestion() {
|
||||
let mut q = qos();
|
||||
for _ in 0..BAD_CONFIRM {
|
||||
q.on_sample(10, 0.2); // low delay but 20% loss
|
||||
}
|
||||
assert!(q.target().bitrate < 6000, "packet loss counts as congestion");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recovers_toward_ceiling_after_congestion() {
|
||||
let mut q = qos();
|
||||
for _ in 0..400 {
|
||||
q.on_sample(500, 0.1);
|
||||
}
|
||||
assert!(q.target().bitrate <= 400);
|
||||
// link heals
|
||||
for _ in 0..1000 {
|
||||
q.on_sample(20, 0.0);
|
||||
}
|
||||
let t = q.target();
|
||||
assert_eq!(t.fps, 30, "fps recovered to ceiling");
|
||||
assert_eq!(t.bitrate, 6000, "bitrate recovered to ceiling");
|
||||
assert!((t.scale - 1.0).abs() < 1e-3, "scale recovered to full");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn never_exceeds_ceiling_on_recovery() {
|
||||
let mut q = qos();
|
||||
for _ in 0..5000 {
|
||||
q.on_sample(5, 0.0);
|
||||
}
|
||||
let t = q.target();
|
||||
assert!(t.fps <= 30 && t.bitrate <= 6000 && t.scale <= 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raising_ceiling_lets_it_climb() {
|
||||
let mut q = Qos::new(QosConfig::new(15, 3000));
|
||||
for _ in 0..200 {
|
||||
q.on_sample(10, 0.0);
|
||||
}
|
||||
assert_eq!(q.target().fps, 15);
|
||||
q.set_ceiling(60, 12000);
|
||||
for _ in 0..1000 {
|
||||
q.on_sample(10, 0.0);
|
||||
}
|
||||
assert_eq!(q.target().fps, 60, "climbs to the raised ceiling");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_ramp_begins_at_init_fps() {
|
||||
assert_eq!(startup_fps(0, 60), INIT_FPS, "a fresh session starts at the ramp floor");
|
||||
assert_eq!(startup_fps(0, 30), INIT_FPS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_ramp_reaches_target_by_ramp_ms() {
|
||||
assert_eq!(startup_fps(RAMP_MS, 60), 60, "full target once the ramp window elapses");
|
||||
assert_eq!(startup_fps(RAMP_MS + 1000, 60), 60, "stays at target after the ramp");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_ramp_is_monotonic_nondecreasing() {
|
||||
let mut prev = 0;
|
||||
for ms in (0..=RAMP_MS).step_by(100) {
|
||||
let f = startup_fps(ms, 60);
|
||||
assert!(f >= prev, "ramp must never step backward (ms={ms})");
|
||||
assert!(f >= INIT_FPS && f <= 60, "ramp stays within [INIT_FPS, target]");
|
||||
prev = f;
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_ramp_midpoint_is_between_floor_and_target() {
|
||||
let mid = startup_fps(RAMP_MS / 2, 60);
|
||||
assert!(mid > INIT_FPS && mid < 60, "half-way through, fps is between floor and target (got {mid})");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_ramp_never_slows_a_low_target() {
|
||||
// A target already at/below the floor is served at full rate immediately.
|
||||
assert_eq!(startup_fps(0, 15), 15);
|
||||
assert_eq!(startup_fps(0, 10), 10);
|
||||
assert_eq!(startup_fps(500, 5), 5);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,289 @@
|
||||
//! Windows SYSTEM service for UNATTENDED access + lock/login-screen capture.
|
||||
//!
|
||||
//! The service runs as LocalSystem (session 0) and spawns the NetDesk host INTO
|
||||
//! the active console session **as SYSTEM**, with its thread desktop following
|
||||
//! the current input desktop. That lets the host capture and inject on the
|
||||
//! secure desktop (the lock screen, the Windows sign-in screen, and UAC) — so a
|
||||
//! remote viewer can sign in and unlock an unattended machine, TeamViewer-style.
|
||||
//!
|
||||
//! Opt-in and isolated: nothing here runs unless `--install-service` is executed
|
||||
//! elevated. The ordinary user-session host (Desktop autostart / admin task) is
|
||||
//! untouched, so installing or removing the service never risks the path that
|
||||
//! already works.
|
||||
//!
|
||||
//! Entry points (wired in the viewer binary's `main`):
|
||||
//! --install-service register + start the service (self-elevates via UAC)
|
||||
//! --uninstall-service stop + delete the service
|
||||
//! --service-run the SCM service entry (run by Windows, not by hand)
|
||||
//! --session-agent the per-session host the service spawns (follows the
|
||||
//! input desktop so the lock screen is captured)
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use std::ffi::c_void;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use windows::core::{PCWSTR, PWSTR};
|
||||
use windows::Win32::Foundation::{CloseHandle, HANDLE, LUID};
|
||||
use windows::Win32::Security::{
|
||||
AdjustTokenPrivileges, DuplicateTokenEx, LookupPrivilegeValueW, SetTokenInformation,
|
||||
SecurityIdentification, TokenPrimary, TokenSessionId, LUID_AND_ATTRIBUTES,
|
||||
SE_PRIVILEGE_ENABLED, TOKEN_ACCESS_MASK, TOKEN_ADJUST_DEFAULT, TOKEN_ADJUST_PRIVILEGES,
|
||||
TOKEN_ADJUST_SESSIONID, TOKEN_ASSIGN_PRIMARY, TOKEN_DUPLICATE, TOKEN_PRIVILEGES, TOKEN_QUERY,
|
||||
};
|
||||
use windows::Win32::System::RemoteDesktop::WTSGetActiveConsoleSessionId;
|
||||
use windows::Win32::System::Services::{
|
||||
RegisterServiceCtrlHandlerW, SetServiceStatus, StartServiceCtrlDispatcherW,
|
||||
SERVICE_STATUS, SERVICE_STATUS_HANDLE, SERVICE_TABLE_ENTRYW, SERVICE_ACCEPT_STOP,
|
||||
SERVICE_CONTROL_SHUTDOWN, SERVICE_CONTROL_STOP, SERVICE_RUNNING, SERVICE_START_PENDING,
|
||||
SERVICE_STOP_PENDING, SERVICE_STOPPED, SERVICE_WIN32_OWN_PROCESS,
|
||||
};
|
||||
use windows::Win32::System::Threading::{
|
||||
CreateProcessAsUserW, GetCurrentProcess, OpenProcessToken, TerminateProcess,
|
||||
CREATE_NO_WINDOW, CREATE_UNICODE_ENVIRONMENT, PROCESS_INFORMATION, STARTUPINFOW,
|
||||
};
|
||||
|
||||
const SERVICE_NAME: &str = "NetDeskHost";
|
||||
|
||||
static STOP: AtomicBool = AtomicBool::new(false);
|
||||
static mut STATUS_HANDLE: isize = 0;
|
||||
|
||||
fn wide(s: &str) -> Vec<u16> {
|
||||
s.encode_utf16().chain(std::iter::once(0)).collect()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// install / uninstall (shell out to sc.exe; requires elevation)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Register the service (LocalSystem, auto-start, auto-restart) and start it.
|
||||
pub fn install_service(exe: &str) -> Result<()> {
|
||||
let bin = format!("\"{exe}\" --service-run");
|
||||
// Create (quirk: sc wants "key=" and the value as separate tokens).
|
||||
run_sc(&["create", SERVICE_NAME, "binPath=", &bin, "start=", "auto",
|
||||
"obj=", "LocalSystem", "DisplayName=", "NetDesk Host (unattended)"])?;
|
||||
let _ = run_sc(&["description", SERVICE_NAME,
|
||||
"NetDesk unattended host — captures the lock/sign-in screen so you can access this PC remotely."]);
|
||||
// Restart on crash so unattended access survives faults.
|
||||
let _ = run_sc(&["failure", SERVICE_NAME, "reset=", "0", "actions=",
|
||||
"restart/5000/restart/5000/restart/5000"]);
|
||||
run_sc(&["start", SERVICE_NAME])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stop and delete the service.
|
||||
pub fn uninstall_service() -> Result<()> {
|
||||
let _ = run_sc(&["stop", SERVICE_NAME]);
|
||||
run_sc(&["delete", SERVICE_NAME])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run_sc(args: &[&str]) -> Result<()> {
|
||||
let status = std::process::Command::new("sc")
|
||||
.args(args)
|
||||
.status()
|
||||
.map_err(|e| anyhow!("sc {args:?}: {e}"))?;
|
||||
if status.success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(anyhow!("sc {:?} exited with {:?}", args, status.code()))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// service entry (run by the SCM via `--service-run`)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Hand control to the Service Control Manager. Blocks until the service stops.
|
||||
pub fn run_service() -> Result<()> {
|
||||
let name = wide(SERVICE_NAME);
|
||||
let table = [
|
||||
SERVICE_TABLE_ENTRYW {
|
||||
lpServiceName: PWSTR(name.as_ptr() as *mut u16),
|
||||
lpServiceProc: Some(service_main),
|
||||
},
|
||||
SERVICE_TABLE_ENTRYW {
|
||||
lpServiceName: PWSTR::null(),
|
||||
lpServiceProc: None,
|
||||
},
|
||||
];
|
||||
unsafe {
|
||||
StartServiceCtrlDispatcherW(table.as_ptr())
|
||||
.map_err(|e| anyhow!("StartServiceCtrlDispatcher: {e}"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
extern "system" fn service_main(_argc: u32, _argv: *mut PWSTR) {
|
||||
unsafe {
|
||||
let name = wide(SERVICE_NAME);
|
||||
match RegisterServiceCtrlHandlerW(PCWSTR(name.as_ptr()), Some(handler)) {
|
||||
Ok(h) => STATUS_HANDLE = h.0 as isize,
|
||||
Err(_) => return,
|
||||
}
|
||||
set_status(SERVICE_START_PENDING, 0);
|
||||
set_status(SERVICE_RUNNING, SERVICE_ACCEPT_STOP);
|
||||
// Main loop: keep a host alive in the active console session.
|
||||
let exe = std::env::current_exe()
|
||||
.map(|p| p.display().to_string())
|
||||
.unwrap_or_default();
|
||||
let mut child: Option<(HANDLE, u32)> = None;
|
||||
while !STOP.load(Ordering::Relaxed) {
|
||||
let session = WTSGetActiveConsoleSessionId();
|
||||
let need_spawn = match child {
|
||||
// 0xFFFFFFFF => no one attached to the console right now.
|
||||
_ if session == 0xFFFF_FFFF => false,
|
||||
None => true,
|
||||
Some((_, sess)) => sess != session || !process_alive(child),
|
||||
};
|
||||
if need_spawn && session != 0xFFFF_FFFF {
|
||||
if let Some((h, _)) = child.take() {
|
||||
let _ = TerminateProcess(h, 1);
|
||||
let _ = CloseHandle(h);
|
||||
}
|
||||
match spawn_host_in_session(session, &exe) {
|
||||
Ok(h) => child = Some((h, session)),
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(1500));
|
||||
}
|
||||
if let Some((h, _)) = child.take() {
|
||||
let _ = TerminateProcess(h, 0);
|
||||
let _ = CloseHandle(h);
|
||||
}
|
||||
set_status(SERVICE_STOPPED, 0);
|
||||
}
|
||||
}
|
||||
|
||||
extern "system" fn handler(control: u32) {
|
||||
if control == SERVICE_CONTROL_STOP || control == SERVICE_CONTROL_SHUTDOWN {
|
||||
STOP.store(true, Ordering::Relaxed);
|
||||
unsafe { set_status(SERVICE_STOP_PENDING, 0) };
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn set_status(state: windows::Win32::System::Services::SERVICE_STATUS_CURRENT_STATE, accept: u32) {
|
||||
if STATUS_HANDLE == 0 {
|
||||
return;
|
||||
}
|
||||
let mut st = SERVICE_STATUS {
|
||||
dwServiceType: SERVICE_WIN32_OWN_PROCESS,
|
||||
dwCurrentState: state,
|
||||
dwControlsAccepted: accept,
|
||||
dwWin32ExitCode: 0,
|
||||
dwServiceSpecificExitCode: 0,
|
||||
dwCheckPoint: 0,
|
||||
dwWaitHint: 3000,
|
||||
};
|
||||
let _ = SetServiceStatus(SERVICE_STATUS_HANDLE(STATUS_HANDLE as *mut c_void), &mut st);
|
||||
}
|
||||
|
||||
fn process_alive(child: Option<(HANDLE, u32)>) -> bool {
|
||||
use windows::Win32::Foundation::{STILL_ACTIVE, WAIT_OBJECT_0};
|
||||
use windows::Win32::System::Threading::{GetExitCodeProcess, WaitForSingleObject};
|
||||
match child {
|
||||
None => false,
|
||||
Some((h, _)) => unsafe {
|
||||
// Fast path: a 0ms wait signals only if the process has exited.
|
||||
if WaitForSingleObject(h, 0) == WAIT_OBJECT_0 {
|
||||
return false;
|
||||
}
|
||||
let mut code = 0u32;
|
||||
if GetExitCodeProcess(h, &mut code).is_ok() {
|
||||
code == STILL_ACTIVE.0 as u32
|
||||
} else {
|
||||
false
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// spawn the host as SYSTEM in the active session (so it can reach the secure
|
||||
// desktop). Duplicates the service's SYSTEM token, re-targets it at the active
|
||||
// session, and CreateProcessAsUser's the host on winsta0\default.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn spawn_host_in_session(session: u32, exe: &str) -> Result<HANDLE> {
|
||||
unsafe {
|
||||
let proc = GetCurrentProcess();
|
||||
let access = TOKEN_ACCESS_MASK(
|
||||
TOKEN_DUPLICATE.0
|
||||
| TOKEN_QUERY.0
|
||||
| TOKEN_ASSIGN_PRIMARY.0
|
||||
| TOKEN_ADJUST_DEFAULT.0
|
||||
| TOKEN_ADJUST_SESSIONID.0
|
||||
| TOKEN_ADJUST_PRIVILEGES.0,
|
||||
);
|
||||
let mut tok = HANDLE::default();
|
||||
OpenProcessToken(proc, access, &mut tok).map_err(|e| anyhow!("OpenProcessToken: {e}"))?;
|
||||
enable_privilege(tok, "SeTcbPrivilege");
|
||||
|
||||
let mut dup = HANDLE::default();
|
||||
DuplicateTokenEx(
|
||||
tok,
|
||||
TOKEN_ACCESS_MASK(0x02000000), // MAXIMUM_ALLOWED
|
||||
None,
|
||||
SecurityIdentification,
|
||||
TokenPrimary,
|
||||
&mut dup,
|
||||
)
|
||||
.map_err(|e| anyhow!("DuplicateTokenEx: {e}"))?;
|
||||
|
||||
// Re-target the primary token at the active console session.
|
||||
let sess = session;
|
||||
SetTokenInformation(
|
||||
dup,
|
||||
TokenSessionId,
|
||||
&sess as *const u32 as *const c_void,
|
||||
std::mem::size_of::<u32>() as u32,
|
||||
)
|
||||
.map_err(|e| anyhow!("SetTokenInformation(session): {e}"))?;
|
||||
|
||||
let mut cmd = wide(&format!("\"{exe}\" --session-agent"));
|
||||
let mut desktop = wide("winsta0\\default");
|
||||
let mut si = STARTUPINFOW::default();
|
||||
si.cb = std::mem::size_of::<STARTUPINFOW>() as u32;
|
||||
si.lpDesktop = PWSTR(desktop.as_mut_ptr());
|
||||
let mut pi = PROCESS_INFORMATION::default();
|
||||
|
||||
CreateProcessAsUserW(
|
||||
dup,
|
||||
PCWSTR::null(),
|
||||
PWSTR(cmd.as_mut_ptr()),
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
CREATE_NO_WINDOW | CREATE_UNICODE_ENVIRONMENT,
|
||||
None,
|
||||
PCWSTR::null(),
|
||||
&si,
|
||||
&mut pi,
|
||||
)
|
||||
.map_err(|e| anyhow!("CreateProcessAsUser: {e}"))?;
|
||||
|
||||
let _ = CloseHandle(pi.hThread);
|
||||
let _ = CloseHandle(dup);
|
||||
let _ = CloseHandle(tok);
|
||||
Ok(pi.hProcess)
|
||||
}
|
||||
}
|
||||
|
||||
fn enable_privilege(token: HANDLE, name: &str) {
|
||||
unsafe {
|
||||
let wname = wide(name);
|
||||
let mut luid = LUID::default();
|
||||
if LookupPrivilegeValueW(PCWSTR::null(), PCWSTR(wname.as_ptr()), &mut luid).is_err() {
|
||||
return;
|
||||
}
|
||||
let tp = TOKEN_PRIVILEGES {
|
||||
PrivilegeCount: 1,
|
||||
Privileges: [LUID_AND_ATTRIBUTES {
|
||||
Luid: luid,
|
||||
Attributes: SE_PRIVILEGE_ENABLED,
|
||||
}],
|
||||
};
|
||||
let _ = AdjustTokenPrivileges(token, false, Some(&tp), 0, None, None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
//! Host system controls (RustDesk/TeamViewer-grade): process list + kill,
|
||||
//! block-local-input, and wallpaper hide/restore. Native Win32 via the windows
|
||||
//! crate — no shelling out, no Python.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use std::ffi::c_void;
|
||||
use std::os::windows::process::CommandExt;
|
||||
use windows::Win32::Foundation::CloseHandle;
|
||||
use windows::Win32::System::Diagnostics::ToolHelp::{
|
||||
CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS,
|
||||
};
|
||||
use windows::Win32::System::ProcessStatus::{GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS};
|
||||
use windows::Win32::System::Threading::{
|
||||
OpenProcess, TerminateProcess, PROCESS_QUERY_LIMITED_INFORMATION, PROCESS_TERMINATE,
|
||||
};
|
||||
use windows::Win32::UI::Input::KeyboardAndMouse::BlockInput;
|
||||
use windows::Win32::UI::WindowsAndMessaging::{
|
||||
SystemParametersInfoW, SPIF_UPDATEINIFILE, SPI_GETDESKWALLPAPER,
|
||||
SPI_SETDESKWALLPAPER,
|
||||
};
|
||||
|
||||
/// Running processes as `{name, pid, mem}` (mem in KB), heaviest first, capped
|
||||
/// at 300 — the shape the viewer's process panel expects.
|
||||
pub fn process_list() -> Vec<Value> {
|
||||
let mut items = Vec::new();
|
||||
unsafe {
|
||||
let snap = match CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) {
|
||||
Ok(h) => h,
|
||||
Err(_) => return items,
|
||||
};
|
||||
let mut entry = PROCESSENTRY32W {
|
||||
dwSize: std::mem::size_of::<PROCESSENTRY32W>() as u32,
|
||||
..Default::default()
|
||||
};
|
||||
if Process32FirstW(snap, &mut entry).is_ok() {
|
||||
loop {
|
||||
let len = entry
|
||||
.szExeFile
|
||||
.iter()
|
||||
.position(|&c| c == 0)
|
||||
.unwrap_or(entry.szExeFile.len());
|
||||
let name = String::from_utf16_lossy(&entry.szExeFile[..len]);
|
||||
let pid = entry.th32ProcessID;
|
||||
let mem = process_mem_kb(pid);
|
||||
items.push(json!({ "name": name, "pid": pid, "mem": mem }));
|
||||
if Process32NextW(snap, &mut entry).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = CloseHandle(snap);
|
||||
}
|
||||
items.sort_by_key(|i| std::cmp::Reverse(i.get("mem").and_then(|m| m.as_u64()).unwrap_or(0)));
|
||||
items.truncate(300);
|
||||
items
|
||||
}
|
||||
|
||||
unsafe fn process_mem_kb(pid: u32) -> u64 {
|
||||
if pid == 0 {
|
||||
return 0;
|
||||
}
|
||||
let h = match OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid) {
|
||||
Ok(h) => h,
|
||||
Err(_) => return 0,
|
||||
};
|
||||
let mut c = PROCESS_MEMORY_COUNTERS::default();
|
||||
let r = GetProcessMemoryInfo(h, &mut c, std::mem::size_of::<PROCESS_MEMORY_COUNTERS>() as u32);
|
||||
let _ = CloseHandle(h);
|
||||
if r.is_ok() {
|
||||
(c.WorkingSetSize / 1024) as u64
|
||||
} else {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
/// Terminate a process by pid. Returns (ok, message).
|
||||
pub fn process_kill(pid: u32) -> (bool, String) {
|
||||
if pid == 0 {
|
||||
return (false, "invalid pid".into());
|
||||
}
|
||||
unsafe {
|
||||
match OpenProcess(PROCESS_TERMINATE, false, pid) {
|
||||
Ok(h) => {
|
||||
let r = TerminateProcess(h, 1);
|
||||
let _ = CloseHandle(h);
|
||||
match r {
|
||||
Ok(()) => (true, format!("terminated {pid}")),
|
||||
Err(e) => (false, format!("terminate failed: {e}")),
|
||||
}
|
||||
}
|
||||
Err(e) => (false, format!("cannot open process {pid}: {e}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Block or unblock local keyboard/mouse so the person at the host cannot
|
||||
/// interfere during a control session. Returns true on success.
|
||||
pub fn set_block_input(on: bool) -> bool {
|
||||
unsafe { BlockInput(on).is_ok() }
|
||||
}
|
||||
|
||||
/// Current desktop wallpaper path (empty if none / unreadable).
|
||||
pub fn wallpaper_get() -> String {
|
||||
unsafe {
|
||||
let mut buf = [0u16; 520];
|
||||
if SystemParametersInfoW(
|
||||
SPI_GETDESKWALLPAPER,
|
||||
buf.len() as u32,
|
||||
Some(buf.as_mut_ptr() as *mut c_void),
|
||||
Default::default(),
|
||||
)
|
||||
.is_ok()
|
||||
{
|
||||
let len = buf.iter().position(|&c| c == 0).unwrap_or(0);
|
||||
String::from_utf16_lossy(&buf[..len])
|
||||
} else {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the desktop wallpaper (empty string clears it to the background colour).
|
||||
pub fn wallpaper_set(path: &str) -> bool {
|
||||
unsafe {
|
||||
let mut wide: Vec<u16> = path.encode_utf16().chain(std::iter::once(0)).collect();
|
||||
SystemParametersInfoW(
|
||||
SPI_SETDESKWALLPAPER,
|
||||
0,
|
||||
Some(wide.as_mut_ptr() as *mut c_void),
|
||||
SPIF_UPDATEINIFILE,
|
||||
)
|
||||
.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Windows process-creation flag: run without flashing a console window.
|
||||
const CREATE_NO_WINDOW: u32 = 0x0800_0000;
|
||||
|
||||
/// Run a shell command on the host and capture stdout+stderr (size-capped).
|
||||
/// Runs with the host process's privileges — full Administrator/SYSTEM when the
|
||||
/// host is elevated (service or run-as-admin), standard user otherwise. This is
|
||||
/// the "run CMD / PowerShell on the remote machine" feature.
|
||||
pub fn run_command(cmd: &str) -> (bool, String) {
|
||||
let out = std::process::Command::new("cmd")
|
||||
.args(["/c", cmd])
|
||||
.creation_flags(CREATE_NO_WINDOW)
|
||||
.output();
|
||||
match out {
|
||||
Ok(o) => {
|
||||
let mut s = String::from_utf8_lossy(&o.stdout).into_owned();
|
||||
let err = String::from_utf8_lossy(&o.stderr);
|
||||
if !err.trim().is_empty() {
|
||||
if !s.is_empty() {
|
||||
s.push('\n');
|
||||
}
|
||||
s.push_str(&err);
|
||||
}
|
||||
if s.len() > 60_000 {
|
||||
s.truncate(60_000);
|
||||
s.push_str("\n…[output truncated]");
|
||||
}
|
||||
if s.trim().is_empty() {
|
||||
s = format!("(exit {}, no output)", o.status.code().unwrap_or(-1));
|
||||
}
|
||||
(o.status.success(), s)
|
||||
}
|
||||
Err(e) => (false, format!("failed to launch: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_detached(program: &str, args: &[&str]) -> bool {
|
||||
std::process::Command::new(program)
|
||||
.args(args)
|
||||
.creation_flags(CREATE_NO_WINDOW)
|
||||
.spawn()
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
/// Lock the workstation (screen-lock). Works from the interactive user session.
|
||||
pub fn lock_workstation() -> bool {
|
||||
spawn_detached("rundll32.exe", &["user32.dll,LockWorkStation"])
|
||||
}
|
||||
|
||||
/// Log the interactive user off.
|
||||
pub fn logoff() -> bool {
|
||||
spawn_detached("shutdown.exe", &["/l", "/f"])
|
||||
}
|
||||
|
||||
/// Restart the machine now.
|
||||
pub fn reboot() -> bool {
|
||||
spawn_detached("shutdown.exe", &["/r", "/t", "0"])
|
||||
}
|
||||
|
||||
/// Shut the machine down now.
|
||||
pub fn shutdown() -> bool {
|
||||
spawn_detached("shutdown.exe", &["/s", "/t", "0"])
|
||||
}
|
||||
|
||||
/// Is this host process running with an elevated (Administrator) token? Reported
|
||||
/// to the viewer so the UI can show whether `run` commands have admin rights.
|
||||
pub fn is_elevated() -> bool {
|
||||
use windows::Win32::Foundation::HANDLE;
|
||||
use windows::Win32::Security::{
|
||||
GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY,
|
||||
};
|
||||
use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
|
||||
unsafe {
|
||||
let mut token = HANDLE::default();
|
||||
if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token).is_err() {
|
||||
return false;
|
||||
}
|
||||
let mut elev = TOKEN_ELEVATION::default();
|
||||
let mut ret = 0u32;
|
||||
let ok = GetTokenInformation(
|
||||
token,
|
||||
TokenElevation,
|
||||
Some(&mut elev as *mut _ as *mut c_void),
|
||||
std::mem::size_of::<TOKEN_ELEVATION>() as u32,
|
||||
&mut ret,
|
||||
)
|
||||
.is_ok();
|
||||
let _ = CloseHandle(token);
|
||||
ok && elev.TokenIsElevated != 0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod admin_tests {
|
||||
use super::*;
|
||||
|
||||
// Safe: runs `whoami` only — never fires lock/logoff/reboot/shutdown.
|
||||
#[test]
|
||||
fn run_command_captures_output() {
|
||||
let (ok, out) = run_command("whoami");
|
||||
assert!(ok, "whoami should succeed: {out}");
|
||||
assert!(!out.trim().is_empty(), "expected username output");
|
||||
let _ = is_elevated(); // must not panic
|
||||
eprintln!("whoami -> {} | elevated={}", out.trim(), is_elevated());
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a Secure Attention Sequence (Ctrl+Alt+Del) to the active session. The
|
||||
/// host must be elevated and the SoftwareSASGeneration policy must allow
|
||||
/// software SAS (set by `--install-admin-task`); fails gracefully otherwise.
|
||||
pub fn send_cad() -> bool {
|
||||
use windows::Win32::Security::Authentication::Identity::SendSAS;
|
||||
// AsUser = TRUE: generated as the interactive (elevated) user.
|
||||
unsafe { SendSAS(windows::Win32::Foundation::TRUE) };
|
||||
true
|
||||
}
|
||||
|
||||
/// Make UAC consent prompts appear on the NORMAL desktop (so a remote viewer can
|
||||
/// see and click them — "enter admin pass" over a session) instead of the
|
||||
/// isolated secure desktop. `visible=false` restores the Windows default.
|
||||
/// Writes HKLM, so it needs an elevated host; fails gracefully otherwise.
|
||||
pub fn set_uac_visible(visible: bool) -> bool {
|
||||
let val = if visible { "0" } else { "1" };
|
||||
std::process::Command::new("reg")
|
||||
.args([
|
||||
"add",
|
||||
r"HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System",
|
||||
"/v", "PromptOnSecureDesktop", "/t", "REG_DWORD", "/d", val, "/f",
|
||||
])
|
||||
.creation_flags(CREATE_NO_WINDOW)
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Approve-before-connect prompt: a topmost Yes/No box on the interactive
|
||||
/// desktop ("<peer> (<user>) wants to control this PC"). Auto-dismisses after
|
||||
/// 20s (returns false), matching the engine's own approve timeout, via
|
||||
/// user32's MessageBoxTimeoutW (exported by name; stable since XP). Returns
|
||||
/// true only on an explicit Yes.
|
||||
pub fn approve_prompt(peer: &str, user: &str) -> bool {
|
||||
#[link(name = "user32")]
|
||||
extern "system" {
|
||||
fn MessageBoxTimeoutW(
|
||||
hwnd: isize,
|
||||
lptext: *const u16,
|
||||
lpcaption: *const u16,
|
||||
utype: u32,
|
||||
wlangid: u16,
|
||||
dwmilliseconds: u32,
|
||||
) -> i32;
|
||||
}
|
||||
let from = match (peer.is_empty(), user.is_empty()) {
|
||||
(true, true) => "an unknown viewer".to_string(),
|
||||
(false, true) => peer.to_string(),
|
||||
(true, false) => user.to_string(),
|
||||
(false, false) => format!("{peer} ({user})"),
|
||||
};
|
||||
let text = format!(
|
||||
"Incoming NetDesk connection\n\nFrom: {from}\n\nThe viewer entered the correct password.\nAccept remote control of this PC?"
|
||||
);
|
||||
let wide = |s: &str| s.encode_utf16().chain(std::iter::once(0)).collect::<Vec<u16>>();
|
||||
let t = wide(&text);
|
||||
let c = wide("NetDesk — accept connection?");
|
||||
const MB_YESNO: u32 = 0x0000_0004;
|
||||
const MB_ICONQUESTION: u32 = 0x0000_0020;
|
||||
const MB_SYSTEMMODAL: u32 = 0x0000_1000; // topmost, above fullscreen apps
|
||||
const MB_SETFOREGROUND: u32 = 0x0001_0000;
|
||||
const MB_DEFBUTTON2: u32 = 0x0000_0100; // default = No (Enter won't accept)
|
||||
const IDYES: i32 = 6;
|
||||
let r = unsafe {
|
||||
MessageBoxTimeoutW(
|
||||
0,
|
||||
t.as_ptr(),
|
||||
c.as_ptr(),
|
||||
MB_YESNO | MB_ICONQUESTION | MB_SYSTEMMODAL | MB_SETFOREGROUND | MB_DEFBUTTON2,
|
||||
0,
|
||||
20_000,
|
||||
)
|
||||
};
|
||||
r == IDYES
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
//! Lightweight host metrics (CPU %, RAM, disk, battery, uptime) for the portal's
|
||||
//! live monitoring. Windows-only; sampled every few seconds by the engine and
|
||||
//! pushed to the server as a `host.telemetry` control message.
|
||||
#![cfg(windows)]
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// Keeps the previous CPU time sample so each `sample()` yields a real % over
|
||||
/// the interval since the last call.
|
||||
pub struct Telemetry {
|
||||
prev_idle: u64,
|
||||
prev_total: u64,
|
||||
}
|
||||
|
||||
impl Telemetry {
|
||||
pub fn new() -> Self {
|
||||
Self { prev_idle: 0, prev_total: 0 }
|
||||
}
|
||||
|
||||
pub fn sample(&mut self) -> Value {
|
||||
let cpu = self.cpu_percent();
|
||||
let (ram_used, ram_total) = mem();
|
||||
let (disk_used, disk_total) = disk();
|
||||
let (battery, charging) = power();
|
||||
json!({
|
||||
"t": "host.telemetry",
|
||||
"cpu": cpu,
|
||||
"ram_used": ram_used, "ram_total": ram_total,
|
||||
"disk_used": disk_used, "disk_total": disk_total,
|
||||
"battery": battery, "charging": charging,
|
||||
"uptime": uptime_secs(),
|
||||
})
|
||||
}
|
||||
|
||||
fn cpu_percent(&mut self) -> u32 {
|
||||
use windows::Win32::Foundation::FILETIME;
|
||||
use windows::Win32::System::Threading::GetSystemTimes;
|
||||
unsafe {
|
||||
let mut idle = FILETIME::default();
|
||||
let mut kern = FILETIME::default();
|
||||
let mut user = FILETIME::default();
|
||||
if GetSystemTimes(Some(&mut idle), Some(&mut kern), Some(&mut user)).is_err() {
|
||||
return 0;
|
||||
}
|
||||
let ft = |f: FILETIME| ((f.dwHighDateTime as u64) << 32) | f.dwLowDateTime as u64;
|
||||
let idle_t = ft(idle);
|
||||
// Kernel time INCLUDES idle, so kernel+user is the full busy+idle total.
|
||||
let total_t = ft(kern) + ft(user);
|
||||
let d_total = total_t.saturating_sub(self.prev_total);
|
||||
let d_idle = idle_t.saturating_sub(self.prev_idle);
|
||||
self.prev_total = total_t;
|
||||
self.prev_idle = idle_t;
|
||||
if d_total == 0 {
|
||||
return 0;
|
||||
}
|
||||
let busy = d_total.saturating_sub(d_idle);
|
||||
((busy as f64 / d_total as f64) * 100.0).round().clamp(0.0, 100.0) as u32
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn mem() -> (u64, u64) {
|
||||
use windows::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
|
||||
unsafe {
|
||||
let mut m = MEMORYSTATUSEX {
|
||||
dwLength: std::mem::size_of::<MEMORYSTATUSEX>() as u32,
|
||||
..Default::default()
|
||||
};
|
||||
if GlobalMemoryStatusEx(&mut m).is_ok() {
|
||||
(m.ullTotalPhys.saturating_sub(m.ullAvailPhys), m.ullTotalPhys)
|
||||
} else {
|
||||
(0, 0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn disk() -> (u64, u64) {
|
||||
use windows::core::PCWSTR;
|
||||
use windows::Win32::Storage::FileSystem::GetDiskFreeSpaceExW;
|
||||
unsafe {
|
||||
let path: Vec<u16> = "C:\\".encode_utf16().chain(std::iter::once(0)).collect();
|
||||
let mut free_avail = 0u64;
|
||||
let mut total = 0u64;
|
||||
let mut total_free = 0u64;
|
||||
if GetDiskFreeSpaceExW(
|
||||
PCWSTR(path.as_ptr()),
|
||||
Some(&mut free_avail),
|
||||
Some(&mut total),
|
||||
Some(&mut total_free),
|
||||
)
|
||||
.is_ok()
|
||||
{
|
||||
(total.saturating_sub(total_free), total)
|
||||
} else {
|
||||
(0, 0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn power() -> (i32, bool) {
|
||||
use windows::Win32::System::Power::GetSystemPowerStatus;
|
||||
use windows::Win32::System::Power::SYSTEM_POWER_STATUS;
|
||||
unsafe {
|
||||
let mut s = SYSTEM_POWER_STATUS::default();
|
||||
if GetSystemPowerStatus(&mut s).is_ok() {
|
||||
let pct = if s.BatteryLifePercent == 255 { -1 } else { s.BatteryLifePercent as i32 };
|
||||
let charging = s.ACLineStatus == 1; // 1 = plugged in
|
||||
(pct, charging)
|
||||
} else {
|
||||
(-1, true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn uptime_secs() -> u64 {
|
||||
use windows::Win32::System::SystemInformation::GetTickCount64;
|
||||
unsafe { GetTickCount64() / 1000 }
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
//! Interactive terminal via Windows ConPTY (pseudo-console) — a real shell the
|
||||
//! viewer drives, not the one-shot `run` command. Mirrors RustDesk 1.5's
|
||||
//! terminal: `term.open` spawns cmd.exe attached to a pseudo-console, its output
|
||||
//! streams back as `term.data`, the viewer's keystrokes are written to its
|
||||
//! input, `term.resize` resizes it, and `term.close` / shell exit tears it down.
|
||||
|
||||
#![cfg(windows)]
|
||||
|
||||
use std::ffi::c_void;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use windows::core::PWSTR;
|
||||
use windows::Win32::Foundation::{CloseHandle, HANDLE};
|
||||
use windows::Win32::Storage::FileSystem::{ReadFile, WriteFile};
|
||||
use windows::Win32::System::Console::{
|
||||
ClosePseudoConsole, CreatePseudoConsole, ResizePseudoConsole, COORD, HPCON,
|
||||
};
|
||||
use windows::Win32::System::Pipes::CreatePipe;
|
||||
use windows::Win32::System::Threading::{
|
||||
CreateProcessW, DeleteProcThreadAttributeList, InitializeProcThreadAttributeList,
|
||||
TerminateProcess, UpdateProcThreadAttribute, EXTENDED_STARTUPINFO_PRESENT, PROCESS_INFORMATION,
|
||||
STARTUPINFOEXW, STARTUPINFOW,
|
||||
};
|
||||
|
||||
// PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE = 0x00020016 (not a named const in the
|
||||
// windows crate's Threading module); value from the Windows SDK.
|
||||
const PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE: usize = 0x0002_0016;
|
||||
|
||||
/// A running ConPTY terminal: write to its input, resize, or kill it. The
|
||||
/// output is pumped by a reader thread (see [`Terminal::spawn`]).
|
||||
pub struct Terminal {
|
||||
hpc: HPCON,
|
||||
input_write: HANDLE,
|
||||
proc: HANDLE,
|
||||
thread: HANDLE,
|
||||
stop: Arc<AtomicBool>,
|
||||
attr_list: Vec<u8>, // keep the proc-thread attribute buffer alive til drop
|
||||
}
|
||||
|
||||
// The handles are owned by this struct + used from the engine thread; the reader
|
||||
// thread only touches a cloned output handle. Mark Send so run.rs can hold it in
|
||||
// its HashMap across .await points (the engine loop is single-threaded per host).
|
||||
unsafe impl Send for Terminal {}
|
||||
|
||||
impl Terminal {
|
||||
/// Spawn `cmd.exe` on a fresh pseudo-console sized `cols`x`rows`. `on_output`
|
||||
/// is called (on a reader thread) with each chunk of raw terminal bytes
|
||||
/// (VT-encoded); `on_exit` is called once when the reader sees EOF. Returns
|
||||
/// the live terminal handle.
|
||||
pub fn spawn(
|
||||
cols: u32,
|
||||
rows: u32,
|
||||
on_output: impl Fn(Vec<u8>) + Send + 'static,
|
||||
on_exit: impl FnOnce() + Send + 'static,
|
||||
) -> std::io::Result<Terminal> {
|
||||
unsafe {
|
||||
// Two pipes: viewer->pty input, pty->viewer output.
|
||||
let mut in_r = HANDLE::default();
|
||||
let mut in_w = HANDLE::default();
|
||||
let mut out_r = HANDLE::default();
|
||||
let mut out_w = HANDLE::default();
|
||||
CreatePipe(&mut in_r, &mut in_w, None, 0)
|
||||
.map_err(|e| std::io::Error::other(format!("CreatePipe(in): {e}")))?;
|
||||
CreatePipe(&mut out_r, &mut out_w, None, 0)
|
||||
.map_err(|e| std::io::Error::other(format!("CreatePipe(out): {e}")))?;
|
||||
|
||||
let size = COORD {
|
||||
X: cols.clamp(1, 1000) as i16,
|
||||
Y: rows.clamp(1, 1000) as i16,
|
||||
};
|
||||
let hpc: HPCON = CreatePseudoConsole(size, in_r, out_w, 0)
|
||||
.map_err(|e| std::io::Error::other(format!("CreatePseudoConsole: {e}")))?;
|
||||
// The pseudo-console now owns these ends; we keep in_w + out_r.
|
||||
let _ = CloseHandle(in_r);
|
||||
let _ = CloseHandle(out_w);
|
||||
|
||||
// Build STARTUPINFOEX with the PSEUDOCONSOLE attribute.
|
||||
let mut attr_size: usize = 0;
|
||||
let _ = InitializeProcThreadAttributeList(
|
||||
windows::Win32::System::Threading::LPPROC_THREAD_ATTRIBUTE_LIST(std::ptr::null_mut()),
|
||||
1,
|
||||
0,
|
||||
&mut attr_size,
|
||||
);
|
||||
let mut attr_list = vec![0u8; attr_size];
|
||||
let attr =
|
||||
windows::Win32::System::Threading::LPPROC_THREAD_ATTRIBUTE_LIST(attr_list.as_mut_ptr() as *mut c_void);
|
||||
InitializeProcThreadAttributeList(attr, 1, 0, &mut attr_size)
|
||||
.map_err(|e| std::io::Error::other(format!("InitAttrList: {e}")))?;
|
||||
UpdateProcThreadAttribute(
|
||||
attr,
|
||||
0,
|
||||
PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE,
|
||||
// Per the Win32 ConPTY sample (EchoCon), lpValue is the HPCON
|
||||
// handle value passed as the pointer, cbValue = sizeof(HPCON).
|
||||
Some(hpc.0 as *const c_void),
|
||||
std::mem::size_of::<HPCON>(),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.map_err(|e| std::io::Error::other(format!("UpdateAttr: {e}")))?;
|
||||
|
||||
let mut si = STARTUPINFOEXW {
|
||||
StartupInfo: STARTUPINFOW {
|
||||
cb: std::mem::size_of::<STARTUPINFOEXW>() as u32,
|
||||
..Default::default()
|
||||
},
|
||||
lpAttributeList: attr,
|
||||
};
|
||||
let mut pi = PROCESS_INFORMATION::default();
|
||||
// Mutable, NUL-terminated command line for CreateProcessW.
|
||||
let mut cmd: Vec<u16> = "cmd.exe".encode_utf16().chain(std::iter::once(0)).collect();
|
||||
let created = CreateProcessW(
|
||||
None,
|
||||
PWSTR(cmd.as_mut_ptr()),
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
EXTENDED_STARTUPINFO_PRESENT,
|
||||
None,
|
||||
None,
|
||||
&si.StartupInfo,
|
||||
&mut pi,
|
||||
);
|
||||
// DeleteProcThreadAttributeList uses lpAttributeList; keep attr_list alive.
|
||||
let _ = &mut si;
|
||||
if let Err(e) = created {
|
||||
DeleteProcThreadAttributeList(attr);
|
||||
let _ = CloseHandle(in_w);
|
||||
let _ = CloseHandle(out_r);
|
||||
ClosePseudoConsole(hpc);
|
||||
return Err(std::io::Error::other(format!("CreateProcessW(cmd): {e}")));
|
||||
}
|
||||
|
||||
// Reader thread: pump ConPTY output -> on_output, EOF -> on_exit.
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let r_stop = stop.clone();
|
||||
let out_r_raw = out_r.0 as isize;
|
||||
std::thread::Builder::new()
|
||||
.name("netdesk-term-read".into())
|
||||
.spawn(move || {
|
||||
let out_r = HANDLE(out_r_raw as *mut c_void);
|
||||
let mut buf = [0u8; 8192];
|
||||
loop {
|
||||
if r_stop.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
let mut read = 0u32;
|
||||
let ok = unsafe { ReadFile(out_r, Some(&mut buf), Some(&mut read), None) };
|
||||
if ok.is_err() || read == 0 {
|
||||
break; // pipe closed / shell exited
|
||||
}
|
||||
on_output(buf[..read as usize].to_vec());
|
||||
}
|
||||
unsafe {
|
||||
let _ = CloseHandle(out_r);
|
||||
}
|
||||
on_exit();
|
||||
})
|
||||
.ok();
|
||||
|
||||
Ok(Terminal {
|
||||
hpc,
|
||||
input_write: in_w,
|
||||
proc: pi.hProcess,
|
||||
thread: pi.hThread,
|
||||
stop,
|
||||
attr_list,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Write viewer keystrokes / command bytes to the shell.
|
||||
pub fn write_input(&self, data: &[u8]) {
|
||||
unsafe {
|
||||
let mut written = 0u32;
|
||||
let _ = WriteFile(self.input_write, Some(data), Some(&mut written), None);
|
||||
}
|
||||
}
|
||||
|
||||
/// Resize the pseudo-console.
|
||||
pub fn resize(&self, cols: u32, rows: u32) {
|
||||
unsafe {
|
||||
let _ = ResizePseudoConsole(
|
||||
self.hpc,
|
||||
COORD {
|
||||
X: cols.clamp(1, 1000) as i16,
|
||||
Y: rows.clamp(1, 1000) as i16,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Terminal {
|
||||
fn drop(&mut self) {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
unsafe {
|
||||
// Closing the PTY + input signals the shell to exit; then reap it.
|
||||
ClosePseudoConsole(self.hpc);
|
||||
let _ = CloseHandle(self.input_write);
|
||||
let _ = TerminateProcess(self.proc, 0);
|
||||
let _ = CloseHandle(self.proc);
|
||||
let _ = CloseHandle(self.thread);
|
||||
}
|
||||
let _ = &self.attr_list; // kept alive until here
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
//! Wake-on-LAN magic-packet sender. An ONLINE NetDesk host can wake another PC
|
||||
//! on its LAN on command (sys action `wol`), so a remote user whose own PC is
|
||||
//! asleep can wake it through any online peer on the same network — the server
|
||||
//! routes the `wol` command to a same-LAN peer, which calls [`send_wol`].
|
||||
//!
|
||||
//! Driverless + dependency-free: a standard WoL magic packet is 6×0xFF followed
|
||||
//! by the target MAC repeated 16 times (102 bytes), UDP-broadcast to the LAN.
|
||||
//! Developed by Akronlogic.
|
||||
|
||||
use std::net::UdpSocket;
|
||||
|
||||
/// Parse "AA:BB:CC:DD:EE:FF" or "AA-BB-..." into 6 bytes.
|
||||
pub fn parse_mac(s: &str) -> Option<[u8; 6]> {
|
||||
let parts: Vec<&str> = s.split(|c| c == ':' || c == '-').collect();
|
||||
if parts.len() != 6 {
|
||||
return None;
|
||||
}
|
||||
let mut mac = [0u8; 6];
|
||||
for (i, p) in parts.iter().enumerate() {
|
||||
mac[i] = u8::from_str_radix(p.trim(), 16).ok()?;
|
||||
}
|
||||
Some(mac)
|
||||
}
|
||||
|
||||
/// Build the 102-byte WoL magic packet: 6×0xFF, then the MAC ×16.
|
||||
pub fn magic_packet(mac: &[u8; 6]) -> [u8; 102] {
|
||||
let mut pkt = [0u8; 102];
|
||||
for b in pkt.iter_mut().take(6) {
|
||||
*b = 0xFF;
|
||||
}
|
||||
for i in 0..16 {
|
||||
pkt[6 + i * 6..6 + i * 6 + 6].copy_from_slice(mac);
|
||||
}
|
||||
pkt
|
||||
}
|
||||
|
||||
/// Broadcast the magic packet for `mac_str` to the LAN. Sends to the limited
|
||||
/// broadcast on the two conventional WoL ports (9 and 7). Returns how many of
|
||||
/// the sends succeeded (0 = total failure).
|
||||
pub fn send_wol(mac_str: &str) -> anyhow::Result<usize> {
|
||||
let mac = parse_mac(mac_str).ok_or_else(|| anyhow::anyhow!("bad MAC: {mac_str}"))?;
|
||||
let pkt = magic_packet(&mac);
|
||||
let sock = UdpSocket::bind("0.0.0.0:0")?;
|
||||
sock.set_broadcast(true)?;
|
||||
let mut sent = 0usize;
|
||||
for addr in ["255.255.255.255:9", "255.255.255.255:7"] {
|
||||
if sock.send_to(&pkt, addr).is_ok() {
|
||||
sent += 1;
|
||||
}
|
||||
}
|
||||
if sent == 0 {
|
||||
anyhow::bail!("no broadcast send succeeded");
|
||||
}
|
||||
Ok(sent)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_and_builds_magic_packet() {
|
||||
assert_eq!(parse_mac("01:02:03:04:05:06").unwrap(), [1, 2, 3, 4, 5, 6]);
|
||||
assert_eq!(parse_mac("AA-BB-CC-DD-EE-FF").unwrap(), [0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]);
|
||||
assert!(parse_mac("nope").is_none());
|
||||
let p = magic_packet(&[1, 2, 3, 4, 5, 6]);
|
||||
assert_eq!(p.len(), 102);
|
||||
assert_eq!(&p[0..6], &[0xFF; 6]);
|
||||
assert_eq!(&p[6..12], &[1, 2, 3, 4, 5, 6]);
|
||||
assert_eq!(&p[96..102], &[1, 2, 3, 4, 5, 6]); // 16th repeat
|
||||
}
|
||||
|
||||
// Prove the packet is actually built + transmitted correctly by sending it
|
||||
// to a bound loopback receiver and reading the 102 bytes back.
|
||||
#[test]
|
||||
fn sends_magic_packet_over_udp_loopback() {
|
||||
let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
|
||||
let port = rx.local_addr().unwrap().port();
|
||||
let mac = parse_mac("aa:bb:cc:dd:ee:ff").unwrap();
|
||||
let pkt = magic_packet(&mac);
|
||||
let tx = UdpSocket::bind("0.0.0.0:0").unwrap();
|
||||
tx.send_to(&pkt, ("127.0.0.1", port)).unwrap();
|
||||
rx.set_read_timeout(Some(std::time::Duration::from_secs(2))).unwrap();
|
||||
let mut buf = [0u8; 256];
|
||||
let (n, _) = rx.recv_from(&mut buf).unwrap();
|
||||
assert_eq!(n, 102);
|
||||
assert_eq!(&buf[0..6], &[0xFF; 6]);
|
||||
assert_eq!(&buf[6..12], &mac);
|
||||
assert_eq!(&buf[96..102], &mac);
|
||||
println!("WOL_SEND_OK 102-byte magic packet transmitted + verified over UDP");
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user