use hbb_common::{log, thiserror}; use std::{ ffi::OsStr, io, ops::{Deref, DerefMut}, os::windows::ffi::OsStrExt, ptr::null_mut, result::Result, }; use winapi::{ shared::{ guiddef::GUID, minwindef::{BOOL, DWORD, FALSE, MAX_PATH, PBOOL, TRUE}, ntdef::{HANDLE, LPCWSTR, NULL}, windef::HWND, winerror::{ERROR_INSUFFICIENT_BUFFER, ERROR_NO_MORE_ITEMS}, }, um::{ cfgmgr32::MAX_DEVICE_ID_LEN, fileapi::{CreateFileW, OPEN_EXISTING}, handleapi::{CloseHandle, INVALID_HANDLE_VALUE}, ioapiset::DeviceIoControl, setupapi::*, winnt::{GENERIC_READ, GENERIC_WRITE}, }, }; #[link(name = "Newdev")] extern "system" { fn UpdateDriverForPlugAndPlayDevicesW( hwnd_parent: HWND, hardware_id: LPCWSTR, full_inf_path: LPCWSTR, install_flags: DWORD, b_reboot_required: PBOOL, ) -> BOOL; } #[derive(thiserror::Error, Debug)] pub enum DeviceError { #[error("Failed to call {0}, {1:?}")] WinApiLastErr(String, io::Error), #[error("Failed to call {0}, returns {1}")] WinApiErrCode(String, DWORD), #[error("{0}")] Raw(String), } struct DeviceInfo(HDEVINFO); impl DeviceInfo { fn setup_di_create_device_info_list(class_guid: &mut GUID) -> Result { let dev_info = unsafe { SetupDiCreateDeviceInfoList(class_guid, null_mut()) }; if dev_info == null_mut() { return Err(DeviceError::WinApiLastErr( "SetupDiCreateDeviceInfoList".to_string(), io::Error::last_os_error(), )); } Ok(Self(dev_info)) } fn setup_di_get_class_devs_ex_w( class_guid: *const GUID, flags: DWORD, ) -> Result { let dev_info = unsafe { SetupDiGetClassDevsExW( class_guid, null_mut(), null_mut(), flags, null_mut(), null_mut(), null_mut(), ) }; if dev_info == null_mut() { return Err(DeviceError::WinApiLastErr( "SetupDiGetClassDevsExW".to_string(), io::Error::last_os_error(), )); } Ok(Self(dev_info)) } } impl Drop for DeviceInfo { fn drop(&mut self) { unsafe { SetupDiDestroyDeviceInfoList(self.0); } } } impl Deref for DeviceInfo { type Target = HDEVINFO; fn deref(&self) -> &Self::Target { &self.0 } } impl DerefMut for DeviceInfo { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } pub unsafe fn install_driver( inf_path: &str, hardware_id: &str, reboot_required: &mut bool, ) -> Result<(), DeviceError> { let driver_inf_path = OsStr::new(inf_path) .encode_wide() .chain(Some(0).into_iter()) .collect::>(); let hardware_id = OsStr::new(hardware_id) .encode_wide() .chain(Some(0).into_iter()) .collect::>(); let mut class_guid: GUID = std::mem::zeroed(); let mut class_name: [u16; 32] = [0; 32]; if SetupDiGetINFClassW( driver_inf_path.as_ptr(), &mut class_guid, class_name.as_mut_ptr(), class_name.len() as _, null_mut(), ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiGetINFClassW".to_string(), io::Error::last_os_error(), )); } let dev_info = DeviceInfo::setup_di_create_device_info_list(&mut class_guid)?; let mut dev_info_data = SP_DEVINFO_DATA { cbSize: std::mem::size_of::() as _, ClassGuid: class_guid, DevInst: 0, Reserved: 0, }; if SetupDiCreateDeviceInfoW( *dev_info, class_name.as_ptr(), &class_guid, null_mut(), null_mut(), DICD_GENERATE_ID, &mut dev_info_data, ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiCreateDeviceInfoW".to_string(), io::Error::last_os_error(), )); } if SetupDiSetDeviceRegistryPropertyW( *dev_info, &mut dev_info_data, SPDRP_HARDWAREID, hardware_id.as_ptr() as _, (hardware_id.len() * 2) as _, ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiSetDeviceRegistryPropertyW".to_string(), io::Error::last_os_error(), )); } if SetupDiCallClassInstaller(DIF_REGISTERDEVICE, *dev_info, &mut dev_info_data) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiCallClassInstaller".to_string(), io::Error::last_os_error(), )); } let mut reboot_required_ = FALSE; if UpdateDriverForPlugAndPlayDevicesW( null_mut(), hardware_id.as_ptr(), driver_inf_path.as_ptr(), 1, &mut reboot_required_, ) == FALSE { return Err(DeviceError::WinApiLastErr( "UpdateDriverForPlugAndPlayDevicesW".to_string(), io::Error::last_os_error(), )); } *reboot_required = reboot_required_ == TRUE; Ok(()) } unsafe fn is_same_hardware_id( dev_info: &DeviceInfo, devinfo_data: &mut SP_DEVINFO_DATA, hardware_id: &str, ) -> Result { let mut cur_hardware_id = [0u16; MAX_DEVICE_ID_LEN]; if SetupDiGetDeviceRegistryPropertyW( **dev_info, devinfo_data, SPDRP_HARDWAREID, null_mut(), cur_hardware_id.as_mut_ptr() as _, cur_hardware_id.len() as _, null_mut(), ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiGetDeviceRegistryPropertyW".to_string(), io::Error::last_os_error(), )); } let cur_hardware_id = String::from_utf16_lossy(&cur_hardware_id) .trim_end_matches(char::from(0)) .to_string(); Ok(cur_hardware_id == hardware_id) } pub unsafe fn uninstall_driver( hardware_id: &str, reboot_required: &mut bool, ) -> Result<(), DeviceError> { let dev_info = DeviceInfo::setup_di_get_class_devs_ex_w(null_mut(), DIGCF_ALLCLASSES | DIGCF_PRESENT)?; let mut device_info_list_detail = SP_DEVINFO_LIST_DETAIL_DATA_W { cbSize: std::mem::size_of::() as _, ClassGuid: std::mem::zeroed(), RemoteMachineHandle: null_mut(), RemoteMachineName: [0; SP_MAX_MACHINENAME_LENGTH], }; if SetupDiGetDeviceInfoListDetailW(*dev_info, &mut device_info_list_detail) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiGetDeviceInfoListDetailW".to_string(), io::Error::last_os_error(), )); } let mut devinfo_data = SP_DEVINFO_DATA { cbSize: std::mem::size_of::() as _, ClassGuid: std::mem::zeroed(), DevInst: 0, Reserved: 0, }; let mut device_index = 0; loop { if SetupDiEnumDeviceInfo(*dev_info, device_index, &mut devinfo_data) == FALSE { let err = io::Error::last_os_error(); if err.raw_os_error() == Some(ERROR_NO_MORE_ITEMS as _) { break; } return Err(DeviceError::WinApiLastErr( "SetupDiEnumDeviceInfo".to_string(), err, )); } match is_same_hardware_id(&dev_info, &mut devinfo_data, hardware_id) { Ok(false) => { device_index += 1; continue; } Err(e) => { log::error!("Failed to call is_same_hardware_id, {:?}", e); device_index += 1; continue; } _ => {} } let mut remove_device_params = SP_REMOVEDEVICE_PARAMS { ClassInstallHeader: SP_CLASSINSTALL_HEADER { cbSize: std::mem::size_of::() as _, InstallFunction: DIF_REMOVE, }, Scope: DI_REMOVEDEVICE_GLOBAL, HwProfile: 0, }; if SetupDiSetClassInstallParamsW( *dev_info, &mut devinfo_data, &mut remove_device_params.ClassInstallHeader, std::mem::size_of::() as _, ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiSetClassInstallParams".to_string(), io::Error::last_os_error(), )); } if SetupDiCallClassInstaller(DIF_REMOVE, *dev_info, &mut devinfo_data) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiCallClassInstaller".to_string(), io::Error::last_os_error(), )); } let mut device_params = SP_DEVINSTALL_PARAMS_W { cbSize: std::mem::size_of::() as _, Flags: 0, FlagsEx: 0, hwndParent: null_mut(), InstallMsgHandler: None, InstallMsgHandlerContext: null_mut(), FileQueue: null_mut(), ClassInstallReserved: 0, Reserved: 0, DriverPath: [0; MAX_PATH], }; if SetupDiGetDeviceInstallParamsW(*dev_info, &mut devinfo_data, &mut device_params) == FALSE { log::error!( "Failed to call SetupDiGetDeviceInstallParamsW, {:?}", io::Error::last_os_error() ); } else { if device_params.Flags & (DI_NEEDRESTART | DI_NEEDREBOOT) != 0 { *reboot_required = true; } } device_index += 1; } Ok(()) } pub unsafe fn device_io_control( interface_guid: &GUID, control_code: u32, inbuf: &[u8], outbuf_max_len: usize, ) -> Result, DeviceError> { let h_device = open_device_handle(interface_guid)?; let mut bytes_returned = 0; let mut outbuf: Vec = vec![]; let outbuf_ptr = if outbuf_max_len > 0 { outbuf.reserve(outbuf_max_len); outbuf.as_mut_ptr() } else { null_mut() }; let result = DeviceIoControl( h_device, control_code, inbuf.as_ptr() as _, inbuf.len() as _, outbuf_ptr as _, outbuf_max_len as _, &mut bytes_returned, null_mut(), ); CloseHandle(h_device); if result == FALSE { return Err(DeviceError::WinApiLastErr( "DeviceIoControl".to_string(), io::Error::last_os_error(), )); } if outbuf_max_len > 0 { outbuf.set_len(bytes_returned as _); Ok(outbuf) } else { Ok(Vec::new()) } } unsafe fn get_device_path(interface_guid: &GUID) -> Result, DeviceError> { let dev_info = DeviceInfo::setup_di_get_class_devs_ex_w( interface_guid, DIGCF_PRESENT | DIGCF_DEVICEINTERFACE, )?; let mut device_interface_data = SP_DEVICE_INTERFACE_DATA { cbSize: std::mem::size_of::() as _, InterfaceClassGuid: *interface_guid, Flags: 0, Reserved: 0, }; if SetupDiEnumDeviceInterfaces( *dev_info, null_mut(), interface_guid, 0, &mut device_interface_data, ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiEnumDeviceInterfaces".to_string(), io::Error::last_os_error(), )); } let mut required_length = 0; if SetupDiGetDeviceInterfaceDetailW( *dev_info, &mut device_interface_data, null_mut(), 0, &mut required_length, null_mut(), ) == FALSE { let err = io::Error::last_os_error(); if err.raw_os_error() != Some(ERROR_INSUFFICIENT_BUFFER as _) { return Err(DeviceError::WinApiLastErr( "SetupDiGetDeviceInterfaceDetailW".to_string(), err, )); } } let predicted_length = required_length; let mut vec_data: Vec = Vec::with_capacity(required_length as _); let device_interface_detail_data = vec_data.as_mut_ptr(); let device_interface_detail_data = device_interface_detail_data as *mut SP_DEVICE_INTERFACE_DETAIL_DATA_W; (*device_interface_detail_data).cbSize = std::mem::size_of::() as _; if SetupDiGetDeviceInterfaceDetailW( *dev_info, &mut device_interface_data, device_interface_detail_data, predicted_length, &mut required_length, null_mut(), ) == FALSE { return Err(DeviceError::WinApiLastErr( "SetupDiGetDeviceInterfaceDetailW".to_string(), io::Error::last_os_error(), )); } let mut path = Vec::new(); let device_path_ptr = std::ptr::addr_of!((*device_interface_detail_data).DevicePath) as *const u16; let steps = device_path_ptr as usize - vec_data.as_ptr() as usize; for i in 0..(predicted_length - steps as u32) / 2 { if *device_path_ptr.offset(i as _) == 0 { path.push(0); break; } path.push(*device_path_ptr.offset(i as _)); } Ok(path) } unsafe fn open_device_handle(interface_guid: &GUID) -> Result { let device_path = get_device_path(interface_guid)?; println!("device_path: {:?}", String::from_utf16_lossy(&device_path)); let h_device = CreateFileW( device_path.as_ptr(), GENERIC_READ | GENERIC_WRITE, 0, null_mut(), OPEN_EXISTING, 0, null_mut(), ); if h_device == INVALID_HANDLE_VALUE || h_device == NULL { return Err(DeviceError::WinApiLastErr( "CreateFileW".to_string(), io::Error::last_os_error(), )); } Ok(h_device) }