diff --git a/Cargo.toml b/Cargo.toml index 62bda75..39a418a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,5 +20,7 @@ features = [ "Win32_UI_WindowsAndMessaging", "Win32_System_Com", "Win32_System_Com_StructuredStorage", - "Win32_Devices_FunctionDiscovery" + "Win32_Devices_FunctionDiscovery", + "Win32_Storage_FileSystem", + "Win32_System_WindowsProgramming", ] \ No newline at end of file diff --git a/mute.txt b/mute.txt index 2caf36a..45696cc 100644 --- a/mute.txt +++ b/mute.txt @@ -13,5 +13,4 @@ hollow_knight.exe sm64.us.f3dex2e.exe teardown.exe underrail.exe -valheim.exe -Spotify.exe \ No newline at end of file +valheim.exe \ No newline at end of file diff --git a/src/device_notification_client.rs b/src/device_notification_client.rs index 7c30aef..6fe8df6 100644 --- a/src/device_notification_client.rs +++ b/src/device_notification_client.rs @@ -6,9 +6,13 @@ use windows::{ }, }; +pub trait DeviceNotificationObserver { + fn add_device(&self, device_id : &windows::core::PCWSTR); +} + #[windows::core::implement(IMMNotificationClient)] pub(crate) struct DeviceNotificationClient { - pub(crate) callback: Box + pub observer: Box } impl IMMNotificationClient_Impl for DeviceNotificationClient { @@ -21,7 +25,7 @@ impl IMMNotificationClient_Impl for DeviceNotificationClient { } fn OnDeviceAdded(&self, pwstrdeviceid: &windows::core::PCWSTR) -> windows::core::Result<()> { - (self.callback)(pwstrdeviceid); + self.observer.add_device(pwstrdeviceid); Ok(()) } diff --git a/src/main.rs b/src/main.rs index 2af8d48..ffab3c5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,57 +12,74 @@ use std::{ io::{BufRead, BufReader}, path::Path, ptr::null_mut, - sync::{Arc, Mutex}, + sync::mpsc::{self, Receiver, Sender}, + thread::{self, JoinHandle}, }; -use sm_session_notifier::SMSessionNotifier; +use sm_session_notifier::SMSessionNotifierThread; use window_change::WindowChangeMonitor; use windows::{ - core::Interface, + core::{Interface}, Win32::{ Media::Audio::{IAudioSessionControl2, ISimpleAudioVolume}, - System::Com::{CoInitializeEx, COINIT_MULTITHREADED}, - }, + System::{Com::{CoInitializeEx, COINIT_MULTITHREADED}, Threading::WaitForSingleObject, WindowsProgramming::INFINITE}, Storage::FileSystem::{FindFirstChangeNotificationW, FILE_NOTIFY_CHANGE_LAST_WRITE, FindCloseChangeNotification}, Foundation::HANDLE, + }, w, }; use crate::pid_to_exe::pid_to_exe_path; +enum MuterMessage { + WindowChange(String), + AddSession(IAudioSessionControl2), + Exit(), +} + +unsafe impl Send for MuterMessage {} + struct SessionMuter { sessions: Vec, mute_executables: HashSet, mute_flag: bool, - win: Option + _sn: SMSessionNotifierThread, + _wn: WindowChangeMonitor, + rx: Receiver, } impl SessionMuter { - fn start(mute_executables: HashSet) -> Arc> { - let muter = Arc::new(Mutex::new(SessionMuter { + fn new( + mute_executables: HashSet, + rx: Receiver, + tx: Sender, + ) -> SessionMuter { + SessionMuter { sessions: Vec::new(), mute_executables, mute_flag: true, - win: None - })); - { - let muter = Arc::downgrade(&muter); - let muter_internal = muter.clone(); - muter.upgrade().unwrap().lock().as_mut().unwrap().win = Some(WindowChangeMonitor::start(Box::new(move |s| { - muter_internal.upgrade().map(|f| f.lock().unwrap().notify_window_changed(&s)); - }))); + _sn: { + let tx = tx.clone(); + SMSessionNotifierThread::new(Box::new(move |session| { + tx.send(MuterMessage::AddSession(session)).unwrap(); + })) + }, + _wn: { + let tx = tx.clone(); + WindowChangeMonitor::start(Box::new(move |s| { + tx.send(MuterMessage::WindowChange(s.to_owned())).unwrap(); + })) + }, + rx, } + } - { - let muter = Arc::downgrade(&muter); - let sn = SMSessionNotifier::new(Box::new(move |session| { - muter.upgrade().unwrap().lock().unwrap().add_session(session).unwrap() - })); - sn.lock() - .unwrap() - .as_mut() - .unwrap() - .boot_devices() - .expect("failed to get initial audio devices and sessions"); + fn run(&mut self) { + loop { + let msg = self.rx.recv().unwrap(); + match msg { + MuterMessage::WindowChange(win) => self.notify_window_changed(&win), + MuterMessage::AddSession(session) => self.add_session(session).unwrap(), + MuterMessage::Exit() => break, + } } - return muter; } fn add_session( @@ -132,23 +149,55 @@ impl SessionMuter { unsafe impl Send for SessionMuter {} +pub struct MuterThread { + handle: Option>, + sender: Sender, +} + +impl MuterThread { + pub fn new(s: HashSet) -> MuterThread { + let (sender, receiver) = mpsc::channel::(); + MuterThread { + sender: sender.clone(), + handle: Some(thread::spawn(move || { + let mut muter = SessionMuter::new(s, receiver, sender); + muter.run(); + })), + } + } +} + +impl Drop for MuterThread { + fn drop(&mut self) { + if let Some(handle) = self.handle.take() { + self.sender.send(MuterMessage::Exit()).unwrap(); + handle.join().unwrap() + } + } +} fn load_mute_txt() -> HashSet { let file = File::open("mute.txt").unwrap(); return HashSet::from_iter(BufReader::new(file).lines().map(|line| line.unwrap())); } - -fn test_muter() { - SessionMuter::start(load_mute_txt()); +fn await_file_change() -> Result<(), Box> { + unsafe { + let handle = FindFirstChangeNotificationW(w!("."), false, FILE_NOTIFY_CHANGE_LAST_WRITE)?; + WaitForSingleObject(HANDLE{0: handle.0}, INFINITE).ok()?; + FindCloseChangeNotification(handle).ok()?; + Ok(()) + } } fn main() { unsafe { CoInitializeEx(None, COINIT_MULTITHREADED).unwrap(); } - loop { - test_muter(); - println!("Test_muter done"); + loop { + let mut _mt = MuterThread::new(load_mute_txt()); + await_file_change().unwrap(); } + + } diff --git a/src/session_notification.rs b/src/session_notification.rs index 3ad9583..ad4fed6 100644 --- a/src/session_notification.rs +++ b/src/session_notification.rs @@ -10,21 +10,25 @@ use windows::{ #[windows::core::implement(IAudioSessionNotification)] pub(crate) struct SessionNotification { - pub(crate) callback: Box + pub(crate) observer: Box } -impl IAudioSessionNotification_Impl for SessionNotification { +pub trait SessionObserver { + fn add_session(&self, session: IAudioSessionControl2); +} + +impl<'a> IAudioSessionNotification_Impl for SessionNotification { fn OnSessionCreated( self: &SessionNotification, newsession: &core::option::Option, ) -> windows::core::Result<()> { let ses: IAudioSessionControl2 = newsession.as_ref().unwrap().cast().unwrap(); - (self.callback)(ses); + self.observer.add_session(ses); Ok(()) } } -impl Drop for SessionNotification { +impl<'a> Drop for SessionNotification { fn drop(&mut self) { println!("SN drop"); } diff --git a/src/sm_session_notifier.rs b/src/sm_session_notifier.rs index 263bfca..c7e00c1 100644 --- a/src/sm_session_notifier.rs +++ b/src/sm_session_notifier.rs @@ -1,10 +1,12 @@ use std::{ error::Error, - sync::{Arc, Mutex}, + sync::mpsc::{self, Receiver, Sender}, + thread::{self, JoinHandle}, }; use crate::{ - device_notification_client::DeviceNotificationClient, session_notification::SessionNotification, + device_notification_client::{DeviceNotificationClient, DeviceNotificationObserver}, + session_notification::{SessionNotification, SessionObserver}, }; use windows::{ core::{Interface, PCWSTR}, @@ -17,46 +19,72 @@ use windows::{ System::Com::{CoCreateInstance, CLSCTX_ALL, STGM_READ}, }, }; + +pub enum SMMessage { + Session(IAudioSessionControl2), + Device(PCWSTR), + Exit(), +} + +unsafe impl Send for SMMessage {} + +struct SessionToMessage { + sender: Sender, +} + +impl SessionObserver for SessionToMessage { + fn add_session(&self, session: IAudioSessionControl2) { + self.sender + .send(SMMessage::Session(session)) + .unwrap_or_else(|_| println!("Failed to add new session")); + } +} + +struct DeviceToMessage { + sender: Sender, +} + +impl DeviceNotificationObserver for DeviceToMessage { + fn add_device(&self, device_id: &windows::core::PCWSTR) { + self.sender + .send(SMMessage::Device(*device_id)) + .unwrap_or_else(|_| println!("Failed to add new device")); + } +} + pub(crate) struct SMSessionNotifier { device_enumerator: IMMDeviceEnumerator, device_notification_client: IMMNotificationClient, session_notification: IAudioSessionNotification, session_managers: Vec, notification_function: Box, + receiver: Receiver, } impl SMSessionNotifier { pub(crate) fn new( callback: Box, - ) -> Arc>> { - let session_notifier_arc: Arc>> = - Arc::new(Mutex::new(None)); - let sn2 = Arc::downgrade(&session_notifier_arc); - let sn3 = sn2.clone(); - let notifier = SMSessionNotifier { + sender: mpsc::Sender, + receiver: mpsc::Receiver, + ) -> SMSessionNotifier { + SMSessionNotifier { session_managers: Vec::new(), device_enumerator: unsafe { CoCreateInstance(&MMDeviceEnumerator, None, CLSCTX_ALL).unwrap() }, device_notification_client: IMMNotificationClient::from(DeviceNotificationClient { - callback: Box::new(move |id| { - match sn2.upgrade().unwrap().lock().unwrap().as_mut().unwrap().add_device_by_id(id) { - Ok(_) => {}, - Err(e) => { - println!("Failed to add new device: {:?}", e); - }, - } + observer: Box::new(DeviceToMessage { + sender: sender.clone(), }), }), session_notification: IAudioSessionNotification::from(SessionNotification { - callback: Box::new(move |session| { - (sn3.upgrade().unwrap().lock().unwrap().as_mut().unwrap().notification_function)(session) + observer: Box::new(SessionToMessage { + sender: sender.clone(), }), }), notification_function: callback, - }; - *session_notifier_arc.lock().unwrap() = Some(notifier); - return session_notifier_arc; + receiver, + } } pub(crate) fn boot_devices(self: &mut SMSessionNotifier) -> Result<(), Box> { @@ -105,8 +133,57 @@ impl SMSessionNotifier { } let prop_store = device.OpenPropertyStore(STGM_READ)?; let prop_var = prop_store.GetValue(&PKEY_Device_FriendlyName)?; - println!("Device Added: {} Existing Sessions: {}", prop_var.Anonymous.Anonymous.Anonymous.pwszVal.to_string()?, session_count); + println!( + "Device Added: {} Existing Sessions: {}", + prop_var.Anonymous.Anonymous.Anonymous.pwszVal.to_string()?, + session_count + ); self.session_managers.push(sm); Ok(()) } + + pub fn run(&mut self) -> Result<(), Box> { + self.boot_devices()?; + + loop { + let msg = self.receiver.recv()?; + match msg { + SMMessage::Session(session) => { + (self.notification_function)(session); + } + SMMessage::Device(id) => { + self.add_device_by_id(&id)?; + } + SMMessage::Exit() => break, + }; + } + Ok(()) + } +} + +pub struct SMSessionNotifierThread { + handle: Option>, + sender: Sender, +} + +impl SMSessionNotifierThread { + pub fn new(s: Box) -> SMSessionNotifierThread { + let (sender, receiver) = mpsc::channel::(); + SMSessionNotifierThread { + sender: sender.clone(), + handle: Some(thread::spawn(move || { + let mut session_notifier = SMSessionNotifier::new(s, sender, receiver); + session_notifier.run().unwrap(); + })), + } + } +} + +impl Drop for SMSessionNotifierThread { + fn drop(&mut self) { + if let Some(handle) = self.handle.take() { + self.sender.send(SMMessage::Exit()).unwrap(); + handle.join().unwrap() + } + } }