AutoMuteTool/src/sm_session_notifier.rs

use std::{
    error::Error,
    sync::mpsc::{self, Receiver, Sender},
    thread::{self, JoinHandle},
};

use crate::{
    device_notification_client::{DeviceNotificationClient, DeviceNotificationObserver},
    session_notification::{SessionNotification, SessionObserver},
};
use windows::{
    core::{Interface, PCWSTR},
    Win32::{
        Devices::FunctionDiscovery::PKEY_Device_FriendlyName,
        Media::Audio::{
            eRender, IAudioSessionControl2, IAudioSessionManager2, IAudioSessionNotification,
            IMMDeviceEnumerator, IMMNotificationClient, MMDeviceEnumerator, DEVICE_STATE_ACTIVE,
        },
        System::Com::{
            CoCreateInstance, StructuredStorage::PropVariantClear, CLSCTX_ALL, STGM_READ,
        },
    },
};

pub enum SMMessage {
    Session(IAudioSessionControl2),
    Device(PCWSTR),
    Exit(),
}

unsafe impl Send for SMMessage {}

struct SessionToMessage {
    sender: Sender<SMMessage>,
}

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<SMMessage>,
}

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<IAudioSessionManager2>,
    notification_function: Box<dyn Fn(IAudioSessionControl2)>,
    receiver: Receiver<SMMessage>,
}

impl SMSessionNotifier {
    pub(crate) fn new(
        callback: Box<dyn Fn(IAudioSessionControl2)>,
        sender: mpsc::Sender<SMMessage>,
        receiver: mpsc::Receiver<SMMessage>,
    ) -> SMSessionNotifier {
        SMSessionNotifier {
            session_managers: Vec::new(),
            device_enumerator: unsafe {
                CoCreateInstance(&MMDeviceEnumerator, None, CLSCTX_ALL).unwrap()
            },
            device_notification_client: IMMNotificationClient::from(DeviceNotificationClient {
                observer: Box::new(DeviceToMessage {
                    sender: sender.clone(),
                }),
            }),
            session_notification: IAudioSessionNotification::from(SessionNotification {
                observer: Box::new(SessionToMessage { sender }),
            }),
            notification_function: callback,
            receiver,
        }
    }

    pub(crate) fn boot_devices(self: &mut SMSessionNotifier) -> Result<(), Box<dyn Error>> {
        unsafe {
            self.device_enumerator
                .RegisterEndpointNotificationCallback(&self.device_notification_client)?;
            let device_collection = self
                .device_enumerator
                .EnumAudioEndpoints(eRender, DEVICE_STATE_ACTIVE)?;
            for device in (0..device_collection.GetCount()?).map(|x| device_collection.Item(x)) {
                let mmdevice = device?;
                self.add_device(mmdevice)?;
            }
            Ok(())
        }
    }

    fn add_device_by_id(self: &mut SMSessionNotifier, id: &PCWSTR) -> Result<(), Box<dyn Error>> {
        unsafe {
            let device_enumerator: IMMDeviceEnumerator =
                CoCreateInstance(&MMDeviceEnumerator, None, CLSCTX_ALL)?;
            let device = device_enumerator.GetDevice(id)?;
            self.add_device(device)?;
            Ok(())
        }
    }

    unsafe fn add_device(
        self: &mut SMSessionNotifier,
        device: windows::Win32::Media::Audio::IMMDevice,
    ) -> Result<(), Box<dyn Error>> {
        let session_manager: IAudioSessionManager2 = device.Activate(CLSCTX_ALL, None)?;
        session_manager.RegisterSessionNotification(&self.session_notification)?;
        let session_enumerator = session_manager.GetSessionEnumerator()?;
        let session_count = session_enumerator.GetCount()?;
        let device_sessions = (0..session_count)
            .map(|idx| {
                session_enumerator
                    .GetSession(idx)
                    .and_then(|session| session.cast())
            })
            .collect::<Result<Vec<IAudioSessionControl2>, _>>()?;
        for session in device_sessions {
            (self.notification_function)(session);
        }
        let prop_store = device.OpenPropertyStore(STGM_READ)?;
        let mut prop_var = prop_store.GetValue(&PKEY_Device_FriendlyName)?;
        println!(
            "Device Added: {} Existing Sessions: {}",
            prop_var.Anonymous.Anonymous.Anonymous.pwszVal.to_string()?,
            session_count
        );
        PropVariantClear(&mut prop_var)?;

        self.session_managers.push(session_manager);
        Ok(())
    }

    pub fn run(&mut self) -> Result<(), Box<dyn Error>> {
        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(())
    }
}

impl Drop for SMSessionNotifier {
    fn drop(&mut self) {
        unsafe {
            self.session_managers
                .drain(..)
                .map(|x| x.UnregisterSessionNotification(&self.session_notification))
                .collect::<Result<(), _>>()
                .unwrap();
            self.device_enumerator
                .UnregisterEndpointNotificationCallback(&self.device_notification_client)
                .unwrap();
        }
    }
}

pub struct SMSessionNotifierThread {
    handle: Option<JoinHandle<()>>,
    sender: Sender<SMMessage>,
}

impl SMSessionNotifierThread {
    pub fn new(s: Box<dyn Fn(IAudioSessionControl2) + Send>) -> SMSessionNotifierThread {
        let (sender, receiver) = mpsc::channel::<SMMessage>();
        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()
        }
    }
}