use super::{CaptureRect, ScreenCapture};
use crate::error::{Result, VncError};
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc, Condvar, Mutex,
};
use std::time::Duration;
use windows_capture::{
capture::{Context, GraphicsCaptureApiHandler},
frame::Frame,
graphics_capture_api::InternalCaptureControl,
monitor::Monitor,
settings::{
ColorFormat, CursorCaptureSettings, DirtyRegionSettings, DrawBorderSettings,
MinimumUpdateIntervalSettings, SecondaryWindowSettings, Settings,
},
};
#[derive(Default)]
struct FramePayload {
data: Vec<u8>,
dirty: Vec<CaptureRect>,
had_dirty_api: bool,
generation: u64,
}
struct CaptureState {
payload: Mutex<Option<FramePayload>>,
active: AtomicBool,
}
struct Handler {
state: Arc<CaptureState>,
staging: Vec<u8>,
dirty_staging: Vec<CaptureRect>,
width: u16,
height: u16,
generation: u64,
}
impl GraphicsCaptureApiHandler for Handler {
type Flags = HandlerFlags;
type Error = Box<dyn std::error::Error + Send + Sync>;
fn new(ctx: Context<Self::Flags>) -> std::result::Result<Self, Self::Error> {
let (lock, cvar) = &*ctx.flags.init_signal;
let mut started = lock.lock().map_err(|e| e.to_string())?;
*started = true;
cvar.notify_all();
Ok(Self {
state: ctx.flags.state,
staging: Vec::new(),
dirty_staging: Vec::with_capacity(64),
width: ctx.flags.width,
height: ctx.flags.height,
generation: 0,
})
}
fn on_frame_arrived(
&mut self,
frame: &mut Frame,
capture_control: InternalCaptureControl,
) -> std::result::Result<(), Self::Error> {
if !self.state.active.load(Ordering::Relaxed) {
capture_control.stop();
return Ok(());
}
self.dirty_staging.clear();
let had_dirty_api = match frame.dirty_regions() {
Ok(regions) => {
for r in regions {
let x = r.x.max(0) as u32;
let y = r.y.max(0) as u32;
let w = r.width.max(0) as u32;
let h = r.height.max(0) as u32;
if w == 0 || h == 0 {
continue;
}
let cx = x.min(self.width as u32) as u16;
let cy = y.min(self.height as u32) as u16;
let cw = w.min(self.width as u32 - cx as u32) as u16;
let ch = h.min(self.height as u32 - cy as u32) as u16;
if cw > 0 && ch > 0 {
self.dirty_staging.push(CaptureRect { x: cx, y: cy, w: cw, h: ch });
}
}
true
}
Err(_) => false,
};
let mut buffer = frame.buffer()?;
let raw = buffer.as_raw_buffer();
self.staging.clear();
self.staging.extend_from_slice(raw);
self.generation = self.generation.wrapping_add(1);
let payload = FramePayload {
data: std::mem::take(&mut self.staging),
dirty: std::mem::take(&mut self.dirty_staging),
had_dirty_api,
generation: self.generation,
};
let mut guard = self.state.payload.lock().map_err(|e| e.to_string())?;
let old = guard.replace(payload);
drop(guard);
if let Some(mut prev) = old {
prev.data.clear();
self.staging = prev.data;
} else {
self.staging.reserve(raw.len());
}
self.dirty_staging.reserve(64);
Ok(())
}
fn on_closed(&mut self) -> std::result::Result<(), Self::Error> {
log::info!("Capture session closed.");
self.state.active.store(false, Ordering::Relaxed);
Ok(())
}
}
#[derive(Clone)]
struct HandlerFlags {
state: Arc<CaptureState>,
init_signal: Arc<(Mutex<bool>, Condvar)>,
width: u16,
height: u16,
}
pub struct WindowsCapture {
state: Arc<CaptureState>,
width: u16,
height: u16,
stride: usize,
last_generation: u64,
last_dirty: Vec<CaptureRect>,
last_had_dirty_api: bool,
capture_thread: Option<std::thread::JoinHandle<()>>,
}
impl WindowsCapture {
pub fn new() -> Result<Self> {
let monitor = Monitor::primary().map_err(|e| VncError::Capture(e.to_string()))?;
let width = monitor
.width()
.map_err(|e| VncError::Capture(e.to_string()))? as u16;
let height = monitor
.height()
.map_err(|e| VncError::Capture(e.to_string()))? as u16;
let stride = width as usize * 4;
let state = Arc::new(CaptureState {
payload: Mutex::new(None),
active: AtomicBool::new(true),
});
let init_signal = Arc::new((Mutex::new(false), Condvar::new()));
let thread_state = state.clone();
let flags = HandlerFlags {
state: thread_state.clone(),
init_signal: init_signal.clone(),
width,
height,
};
let init_signal_err = init_signal.clone();
let capture_thread = std::thread::spawn(move || {
let settings = Settings::new(
monitor,
CursorCaptureSettings::WithoutCursor,
DrawBorderSettings::WithoutBorder,
SecondaryWindowSettings::Default,
MinimumUpdateIntervalSettings::Default,
DirtyRegionSettings::ReportOnly,
ColorFormat::Bgra8,
flags,
);
if let Err(e) = Handler::start_free_threaded(settings) {
log::error!("Capture error: {}", e);
thread_state.active.store(false, Ordering::Relaxed);
let (lock, cvar) = &*init_signal_err;
if let Ok(mut started) = lock.lock() {
*started = true;
cvar.notify_all();
}
}
});
let (lock, cvar) = &*init_signal;
let mut started = lock
.lock()
.map_err(|_| VncError::Capture("Mutex poisoned during init".into()))?;
while !*started {
let result = cvar
.wait_timeout(started, Duration::from_secs(3))
.map_err(|_| VncError::Capture("Condvar poisoned".into()))?;
if result.1.timed_out() {
return Err(VncError::Capture("Capture initialization timed out".into()));
}
started = result.0;
}
if !state.active.load(Ordering::Relaxed) {
return Err(VncError::Capture("Capture failed to start".into()));
}
Ok(Self {
state,
width,
height,
stride,
last_generation: 0,
last_dirty: Vec::with_capacity(64),
last_had_dirty_api: false,
capture_thread: Some(capture_thread),
})
}
}
impl Drop for WindowsCapture {
fn drop(&mut self) {
self.state.active.store(false, Ordering::Relaxed);
if let Some(handle) = self.capture_thread.take() {
let _ = handle.join();
}
}
}
impl ScreenCapture for WindowsCapture {
fn width(&self) -> u16 {
self.width
}
fn height(&self) -> u16 {
self.height
}
fn stride(&self) -> usize {
self.stride
}
fn swap_frame(&mut self, buf: &mut Vec<u8>) -> Result<bool> {
if !self.state.active.load(Ordering::Relaxed) {
return Err(VncError::Capture("Capture session ended".into()));
}
let mut guard = self
.state
.payload
.lock()
.map_err(|_| VncError::Capture("Capture state mutex poisoned".into()))?;
let payload = match guard.as_mut() {
Some(p) => p,
None => return Ok(false),
};
if payload.generation == self.last_generation {
return Ok(false);
}
self.last_generation = payload.generation;
std::mem::swap(buf, &mut payload.data);
self.last_dirty.clear();
self.last_dirty.extend_from_slice(&payload.dirty);
self.last_had_dirty_api = payload.had_dirty_api;
Ok(true)
}
fn take_dirty_hints(&mut self, out: &mut Vec<CaptureRect>) -> bool {
out.clear();
if !self.last_had_dirty_api {
return false;
}
out.extend_from_slice(&self.last_dirty);
true
}
}