fufesou 2aaca0c54a Remove get_error() in windows.rs
Signed-off-by: fufesou <shuanglongchen@yeah.net>
2023-11-20 17:38:53 +08:00

650 lines
21 KiB
Rust

// logic from webrtc -- https://github.com/shiguredo/libwebrtc/blob/main/modules/desktop_capture/win/screen_capturer_win_magnifier.cc
use lazy_static;
use std::{
ffi::CString,
io::{Error, ErrorKind, Result},
mem::size_of,
sync::Mutex,
};
use winapi::{
shared::{
basetsd::SIZE_T,
guiddef::{IsEqualGUID, GUID},
minwindef::{BOOL, DWORD, FALSE, FARPROC, HINSTANCE, HMODULE, HRGN, TRUE, UINT},
ntdef::{LONG, NULL},
windef::{HWND, RECT},
winerror::ERROR_CLASS_ALREADY_EXISTS,
},
um::{
errhandlingapi::GetLastError,
libloaderapi::{FreeLibrary, GetModuleHandleExA, GetProcAddress, LoadLibraryExA},
winuser::*,
},
};
pub const MW_FILTERMODE_EXCLUDE: u32 = 0;
pub const MW_FILTERMODE_INCLUDE: u32 = 1;
pub const GET_MODULE_HANDLE_EX_FLAG_PIN: u32 = 1;
pub const GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT: u32 = 2;
pub const GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS: u32 = 4;
pub const LOAD_LIBRARY_AS_DATAFILE: u32 = 2;
pub const LOAD_WITH_ALTERED_SEARCH_PATH: u32 = 8;
pub const LOAD_IGNORE_CODE_AUTHZ_LEVEL: u32 = 16;
pub const LOAD_LIBRARY_AS_IMAGE_RESOURCE: u32 = 32;
pub const LOAD_LIBRARY_AS_DATAFILE_EXCLUSIVE: u32 = 64;
pub const LOAD_LIBRARY_REQUIRE_SIGNED_TARGET: u32 = 128;
pub const LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR: u32 = 256;
pub const LOAD_LIBRARY_SEARCH_APPLICATION_DIR: u32 = 512;
pub const LOAD_LIBRARY_SEARCH_USER_DIRS: u32 = 1024;
pub const LOAD_LIBRARY_SEARCH_SYSTEM32: u32 = 2048;
pub const LOAD_LIBRARY_SEARCH_DEFAULT_DIRS: u32 = 4096;
pub const LOAD_LIBRARY_SAFE_CURRENT_DIRS: u32 = 8192;
pub const LOAD_LIBRARY_SEARCH_SYSTEM32_NO_FORWARDER: u32 = 16384;
pub const LOAD_LIBRARY_OS_INTEGRITY_CONTINUITY: u32 = 32768;
extern "C" {
pub static GUID_WICPixelFormat32bppRGBA: GUID;
}
lazy_static::lazy_static! {
static ref MAG_BUFFER: Mutex<(bool, Vec<u8>)> = Default::default();
}
pub type REFWICPixelFormatGUID = *const GUID;
pub type WICPixelFormatGUID = GUID;
#[allow(non_snake_case)]
#[repr(C)]
#[derive(Copy, Clone)]
pub struct tagMAGIMAGEHEADER {
pub width: UINT,
pub height: UINT,
pub format: WICPixelFormatGUID,
pub stride: UINT,
pub offset: UINT,
pub cbSize: SIZE_T,
}
pub type MAGIMAGEHEADER = tagMAGIMAGEHEADER;
pub type PMAGIMAGEHEADER = *mut tagMAGIMAGEHEADER;
// Function types
pub type MagImageScalingCallback = ::std::option::Option<
unsafe extern "C" fn(
hwnd: HWND,
srcdata: *mut ::std::os::raw::c_void,
srcheader: MAGIMAGEHEADER,
destdata: *mut ::std::os::raw::c_void,
destheader: MAGIMAGEHEADER,
unclipped: RECT,
clipped: RECT,
dirty: HRGN,
) -> BOOL,
>;
extern "C" {
pub fn MagShowSystemCursor(fShowCursor: BOOL) -> BOOL;
}
pub type MagInitializeFunc = ::std::option::Option<unsafe extern "C" fn() -> BOOL>;
pub type MagUninitializeFunc = ::std::option::Option<unsafe extern "C" fn() -> BOOL>;
pub type MagSetWindowSourceFunc =
::std::option::Option<unsafe extern "C" fn(hwnd: HWND, rect: RECT) -> BOOL>;
pub type MagSetWindowFilterListFunc = ::std::option::Option<
unsafe extern "C" fn(
hwnd: HWND,
dwFilterMode: DWORD,
count: ::std::os::raw::c_int,
pHWND: *mut HWND,
) -> BOOL,
>;
pub type MagSetImageScalingCallbackFunc = ::std::option::Option<
unsafe extern "C" fn(hwnd: HWND, callback: MagImageScalingCallback) -> BOOL,
>;
#[repr(C)]
#[derive(Debug, Clone)]
struct MagInterface {
init_succeeded: bool,
lib_handle: HINSTANCE,
pub mag_initialize_func: MagInitializeFunc,
pub mag_uninitialize_func: MagUninitializeFunc,
pub set_window_source_func: MagSetWindowSourceFunc,
pub set_window_filter_list_func: MagSetWindowFilterListFunc,
pub set_image_scaling_callback_func: MagSetImageScalingCallbackFunc,
}
// NOTE: MagInitialize and MagUninitialize should not be called in global init and uninit.
// If so, strange errors occur.
impl MagInterface {
fn new() -> Result<Self> {
let mut s = MagInterface {
init_succeeded: false,
lib_handle: NULL as _,
mag_initialize_func: None,
mag_uninitialize_func: None,
set_window_source_func: None,
set_window_filter_list_func: None,
set_image_scaling_callback_func: None,
};
s.init_succeeded = false;
unsafe {
// load lib
let lib_file_name = "Magnification.dll";
let lib_file_name_c = CString::new(lib_file_name)?;
s.lib_handle = LoadLibraryExA(
lib_file_name_c.as_ptr() as _,
NULL,
LOAD_WITH_ALTERED_SEARCH_PATH,
);
if s.lib_handle.is_null() {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to LoadLibraryExA {}, error: {}",
lib_file_name,
Error::last_os_error()
),
));
};
// load functions
s.mag_initialize_func = Some(std::mem::transmute(Self::load_func(
s.lib_handle,
"MagInitialize",
)?));
s.mag_uninitialize_func = Some(std::mem::transmute(Self::load_func(
s.lib_handle,
"MagUninitialize",
)?));
s.set_window_source_func = Some(std::mem::transmute(Self::load_func(
s.lib_handle,
"MagSetWindowSource",
)?));
s.set_window_filter_list_func = Some(std::mem::transmute(Self::load_func(
s.lib_handle,
"MagSetWindowFilterList",
)?));
s.set_image_scaling_callback_func = Some(std::mem::transmute(Self::load_func(
s.lib_handle,
"MagSetImageScalingCallback",
)?));
// MagInitialize
if let Some(init_func) = s.mag_initialize_func {
if FALSE == init_func() {
return Err(Error::new(
ErrorKind::Other,
format!("Failed to MagInitialize, error: {}", Error::last_os_error()),
));
} else {
s.init_succeeded = true;
}
} else {
return Err(Error::new(
ErrorKind::Other,
"Unreachable, mag_initialize_func should not be none",
));
}
}
Ok(s)
}
unsafe fn load_func(lib_module: HMODULE, func_name: &str) -> Result<FARPROC> {
let func_name_c = CString::new(func_name)?;
let func = GetProcAddress(lib_module, func_name_c.as_ptr() as _);
if func.is_null() {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to GetProcAddress {}, error {}",
func_name,
Error::last_os_error()
),
));
}
Ok(func)
}
pub(super) fn uninit(&mut self) {
if self.init_succeeded {
if let Some(uninit_func) = self.mag_uninitialize_func {
unsafe {
if FALSE == uninit_func() {
println!("Failed MagUninitialize, error {}", Error::last_os_error())
}
}
}
if !self.lib_handle.is_null() {
unsafe {
if FALSE == FreeLibrary(self.lib_handle) {
println!("Failed FreeLibrary, error {}", Error::last_os_error())
}
}
self.lib_handle = NULL as _;
}
}
self.init_succeeded = false;
}
}
impl Drop for MagInterface {
fn drop(&mut self) {
self.uninit();
}
}
pub struct CapturerMag {
mag_interface: MagInterface,
host_window: HWND,
magnifier_window: HWND,
magnifier_host_class: CString,
host_window_name: CString,
magnifier_window_class: CString,
magnifier_window_name: CString,
rect: RECT,
width: usize,
height: usize,
}
impl Drop for CapturerMag {
fn drop(&mut self) {
self.destroy_windows();
self.mag_interface.uninit();
}
}
impl CapturerMag {
pub(crate) fn is_supported() -> bool {
MagInterface::new().is_ok()
}
pub(crate) fn new(origin: (i32, i32), width: usize, height: usize) -> Result<Self> {
unsafe {
let x = GetSystemMetrics(SM_XVIRTUALSCREEN);
let y = GetSystemMetrics(SM_YVIRTUALSCREEN);
let w = GetSystemMetrics(SM_CXVIRTUALSCREEN);
let h = GetSystemMetrics(SM_CYVIRTUALSCREEN);
if !(origin.0 >= x as i32
&& origin.1 >= y as i32
&& width <= w as usize
&& height <= h as usize)
{
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed Check screen rect ({}, {}, {} , {}) to ({}, {}, {}, {})",
origin.0,
origin.1,
origin.0 + width as i32,
origin.1 + height as i32,
x,
y,
x + w,
y + h
),
));
}
}
let mut s = Self {
mag_interface: MagInterface::new()?,
host_window: 0 as _,
magnifier_window: 0 as _,
magnifier_host_class: CString::new("ScreenCapturerWinMagnifierHost")?,
host_window_name: CString::new("MagnifierHost")?,
magnifier_window_class: CString::new("Magnifier")?,
magnifier_window_name: CString::new("MagnifierWindow")?,
rect: RECT {
left: origin.0 as _,
top: origin.1 as _,
right: origin.0 + width as LONG,
bottom: origin.1 + height as LONG,
},
width,
height,
};
unsafe {
let mut instance = 0 as HMODULE;
if 0 == GetModuleHandleExA(
GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS
| GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
DefWindowProcA as _,
&mut instance as _,
) {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to GetModuleHandleExA, error {}",
Error::last_os_error()
),
));
}
// Register the host window class. See the MSDN documentation of the
// Magnification API for more information.
let wcex = WNDCLASSEXA {
cbSize: size_of::<WNDCLASSEXA>() as _,
style: 0,
lpfnWndProc: Some(DefWindowProcA),
cbClsExtra: 0,
cbWndExtra: 0,
hInstance: instance,
hIcon: 0 as _,
hCursor: LoadCursorA(NULL as _, IDC_ARROW as _),
hbrBackground: 0 as _,
lpszClassName: s.magnifier_host_class.as_ptr() as _,
lpszMenuName: 0 as _,
hIconSm: 0 as _,
};
// Ignore the error which may happen when the class is already registered.
if 0 == RegisterClassExA(&wcex) {
let code = GetLastError();
if code != ERROR_CLASS_ALREADY_EXISTS {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to RegisterClassExA, error {:?}",
Error::from_raw_os_error(code as _)
),
));
}
}
// Create the host window.
s.host_window = CreateWindowExA(
WS_EX_LAYERED,
s.magnifier_host_class.as_ptr(),
s.host_window_name.as_ptr(),
WS_POPUP,
0,
0,
0,
0,
NULL as _,
NULL as _,
instance,
NULL,
);
if s.host_window.is_null() {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to CreateWindowExA host_window, error {}",
Error::last_os_error()
),
));
}
// Create the magnifier control.
s.magnifier_window = CreateWindowExA(
0,
s.magnifier_window_class.as_ptr(),
s.magnifier_window_name.as_ptr(),
WS_CHILD | WS_VISIBLE,
0,
0,
0,
0,
s.host_window,
NULL as _,
instance,
NULL,
);
if s.magnifier_window.is_null() {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed CreateWindowA magnifier_window, error {}",
Error::last_os_error()
),
));
}
// Hide the host window.
let _ = ShowWindow(s.host_window, SW_HIDE);
// Set the scaling callback to receive captured image.
if let Some(set_callback_func) = s.mag_interface.set_image_scaling_callback_func {
if FALSE
== set_callback_func(
s.magnifier_window,
Some(Self::on_gag_image_scaling_callback),
)
{
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to MagSetImageScalingCallback, error {}",
Error::last_os_error()
),
));
}
} else {
return Err(Error::new(
ErrorKind::Other,
"Unreachable, set_image_scaling_callback_func should not be none",
));
}
}
Ok(s)
}
pub(crate) fn exclude(&mut self, cls: &str, name: &str) -> Result<bool> {
let name_c = CString::new(name)?;
unsafe {
let mut hwnd = if cls.len() == 0 {
FindWindowExA(NULL as _, NULL as _, NULL as _, name_c.as_ptr())
} else {
let cls_c = CString::new(cls).unwrap();
FindWindowExA(NULL as _, NULL as _, cls_c.as_ptr(), name_c.as_ptr())
};
if hwnd.is_null() {
return Ok(false);
}
if let Some(set_window_filter_list_func) =
self.mag_interface.set_window_filter_list_func
{
if FALSE
== set_window_filter_list_func(
self.magnifier_window,
MW_FILTERMODE_EXCLUDE,
1,
&mut hwnd,
)
{
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed MagSetWindowFilterList for cls {} name {}, error {}",
cls,
name,
Error::last_os_error()
),
));
}
} else {
return Err(Error::new(
ErrorKind::Other,
"Unreachable, MagSetWindowFilterList should not be none",
));
}
}
Ok(true)
}
pub(crate) fn get_rect(&self) -> ((i32, i32), usize, usize) {
(
(self.rect.left as _, self.rect.top as _),
self.width as _,
self.height as _,
)
}
fn clear_data() {
let mut lock = MAG_BUFFER.lock().unwrap();
lock.0 = false;
lock.1.clear();
}
pub(crate) fn frame(&mut self, data: &mut Vec<u8>) -> Result<()> {
Self::clear_data();
unsafe {
let x = GetSystemMetrics(SM_XVIRTUALSCREEN);
let y = GetSystemMetrics(SM_YVIRTUALSCREEN);
let w = GetSystemMetrics(SM_CXVIRTUALSCREEN);
let h = GetSystemMetrics(SM_CYVIRTUALSCREEN);
if !(self.rect.left >= x as i32
&& self.rect.top >= y as i32
&& self.rect.right <= (x + w) as i32
&& self.rect.bottom <= (y + h) as i32)
{
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed Check screen rect ({}, {}, {} , {}) to ({}, {}, {}, {})",
self.rect.left,
self.rect.top,
self.rect.right,
self.rect.bottom,
x,
y,
x + w,
y + h
),
));
}
if FALSE
== SetWindowPos(
self.magnifier_window,
HWND_TOP,
self.rect.left,
self.rect.top,
self.rect.right - self.rect.left,
self.rect.bottom - self.rect.top,
0,
)
{
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed SetWindowPos (x, y, w , h) - ({}, {}, {}, {}), error {}",
self.rect.left,
self.rect.top,
self.rect.right - self.rect.left,
self.rect.bottom - self.rect.top,
Error::last_os_error()
),
));
}
// on_gag_image_scaling_callback will be called and fill in the
// frame before set_window_source_func_ returns.
if let Some(set_window_source_func) = self.mag_interface.set_window_source_func {
if FALSE == set_window_source_func(self.magnifier_window, self.rect) {
return Err(Error::new(
ErrorKind::Other,
format!(
"Failed to MagSetWindowSource, error {}",
Error::last_os_error()
),
));
}
} else {
return Err(Error::new(
ErrorKind::Other,
"Unreachable, set_window_source_func should not be none",
));
}
}
let mut lock = MAG_BUFFER.lock().unwrap();
if !lock.0 {
return Err(Error::new(
ErrorKind::Other,
"No data captured by magnifier",
));
}
data.resize(lock.1.len(), 0);
unsafe {
std::ptr::copy_nonoverlapping(&mut lock.1[0], &mut data[0], data.len());
}
Ok(())
}
fn destroy_windows(&mut self) {
if !self.magnifier_window.is_null() {
unsafe {
if FALSE == DestroyWindow(self.magnifier_window) {
//
println!(
"Failed DestroyWindow magnifier window, error {}",
Error::last_os_error()
)
}
}
}
self.magnifier_window = NULL as _;
if !self.host_window.is_null() {
unsafe {
if FALSE == DestroyWindow(self.host_window) {
//
println!(
"Failed DestroyWindow host window, error {}",
Error::last_os_error()
)
}
}
}
self.host_window = NULL as _;
}
unsafe extern "C" fn on_gag_image_scaling_callback(
_hwnd: HWND,
srcdata: *mut ::std::os::raw::c_void,
srcheader: MAGIMAGEHEADER,
_destdata: *mut ::std::os::raw::c_void,
_destheader: MAGIMAGEHEADER,
_unclipped: RECT,
_clipped: RECT,
_dirty: HRGN,
) -> BOOL {
Self::clear_data();
if !IsEqualGUID(&srcheader.format, &GUID_WICPixelFormat32bppRGBA) {
// log warning?
return FALSE;
}
let mut lock = MAG_BUFFER.lock().unwrap();
lock.1.resize(srcheader.cbSize, 0);
std::ptr::copy_nonoverlapping(srcdata as _, &mut lock.1[0], srcheader.cbSize);
lock.0 = true;
TRUE
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test() {
let mut capture_mag = CapturerMag::new((0, 0), 1920, 1080).unwrap();
capture_mag.exclude("", "RustDeskPrivacyWindow").unwrap();
std::thread::sleep(std::time::Duration::from_millis(1000 * 10));
let mut data = Vec::new();
capture_mag.frame(&mut data).unwrap();
println!("capture data len: {}", data.len());
}
}