Underlay/src/windows/mod.rs

use windows::Win32::Graphics::Gdi::ClientToScreen;
use windows::Win32::System::Com::VT_I4;
use windows::Win32::System::Ole::{VariantClear, VariantInit};
use windows::Win32::UI::Accessibility::{
    AccessibleObjectFromEvent, AccessibleObjectFromWindow, IAccessible,
};

use std::sync::atomic::AtomicPtr;
use std::sync::{Arc, Mutex};
use windows::core::{Error, PCWSTR};
use windows::Win32::Foundation::{
    GetLastError, SetLastError, ERROR_ACCESS_DENIED, LPARAM, POINT, RECT, WPARAM,
};
use windows::Win32::UI::WindowsAndMessaging::{
    DispatchMessageW, GetClientRect, GetForegroundWindow, GetMessageW, GetWindowTextW,
    GetWindowThreadProcessId, IsHungAppWindow, PostMessageW, RegisterWindowMessageW,
    TranslateMessage, CHILDID_SELF, EVENT_OBJECT_DESTROY, EVENT_OBJECT_LOCATIONCHANGE,
    EVENT_OBJECT_NAMECHANGE, EVENT_SYSTEM_FOREGROUND, EVENT_SYSTEM_MINIMIZEEND, MSG,
    OBJECT_IDENTIFIER, OBJID_WINDOW, STATE_SYSTEM_FOCUSED, WINEVENT_OUTOFCONTEXT,
};
use windows::Win32::{
    Foundation::HWND,
    UI::Accessibility::{SetWinEventHook, UnhookWinEvent, HWINEVENTHOOK},
};

use crate::{Bounds, UnderlayEvent};

static GLOBAL_CALLBACK_SENDER: AtomicPtr<Arc<Mutex<Runtime>>> =
    AtomicPtr::new(std::ptr::null_mut());

fn constant_name(value: u32) -> String {
    match value {
        EVENT_OBJECT_DESTROY => "EVENT_OBJECT_DESTROY".into(),
        EVENT_OBJECT_LOCATIONCHANGE => "EVENT_OBJECT_LOCATIONCHANGE".into(),
        EVENT_OBJECT_NAMECHANGE => "EVENT_OBJECT_NAMECHANGE".into(),
        EVENT_SYSTEM_FOREGROUND => "EVENT_SYSTEM_FOREGROUND".into(),
        EVENT_SYSTEM_MINIMIZEEND => "EVENT_SYSTEM_MINIMIZEEND".into(),
        WINEVENT_OUTOFCONTEXT => "WINEVENT_OUTOFCONTEXT".into(),
        _ => format!("Unknown constant ({value})"),
    }
}

#[derive(Clone)]
struct HookInfo {
    pub id: HWINEVENTHOOK,
    pub min: u32,
    pub max: u32,
    pub idprocess: u32,
    pub idthread: u32,
}

impl std::fmt::Debug for HookInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("HookInfo")
            .field("id", &self.id)
            .field("min", &constant_name(self.min))
            .field("max", &constant_name(self.max))
            .field("idprocess", &self.idprocess)
            .field("idthread", &self.idthread)
            .finish()
    }
}

impl HookInfo {
    pub fn new(min: u32, max: u32) -> Self {
        Self {
            id: Default::default(),
            min,
            max,
            idprocess: Default::default(),
            idthread: Default::default(),
        }
    }

    pub fn bind(&mut self) {
        log::trace!("Binding hook: {self:?}");

        if self.id != Default::default() {
            log::warn!(
                "Binding {} to {} without unbinding the previous one",
                constant_name(self.min),
                constant_name(self.max)
            )
        }

        unsafe {
            self.id = SetWinEventHook(
                self.min,
                self.max,
                None,
                Some(global_hook),
                self.idprocess,
                self.idthread,
                WINEVENT_OUTOFCONTEXT,
            );
        }
    }

    pub fn unbind(&mut self) {
        if self.id == Default::default() {
            log::warn!("Unbinding unbind hook...");
            return;
        };

        unsafe { UnhookWinEvent(self.id) };
        self.id = HWINEVENTHOOK::default();
    }
}

impl Drop for HookInfo {
    fn drop(&mut self) {
        log::trace!("Dropping hook: {self:?}");
        self.unbind();
    }
}

impl Event {
    fn from_data(event: u32, hwnd: HWND, idobject: OBJECT_IDENTIFIER, idchild: u32) -> Self {
        match event {
            EVENT_OBJECT_DESTROY => Self::Destroyed {
                hwnd,
                idobject,
                idchild,
            },
            EVENT_OBJECT_LOCATIONCHANGE => Self::LocationChange {
                hwnd,
                idobject,
                idchild,
            },
            EVENT_OBJECT_NAMECHANGE => Self::NameChange {
                hwnd,
                idobject,
                idchild,
            },
            EVENT_SYSTEM_FOREGROUND | EVENT_SYSTEM_MINIMIZEEND => Event::ForegroundChange { hwnd },
            _ => Event::Unrepresented(event),
        }
    }
}

#[derive(Debug, Clone)]
enum Event {
    ForegroundChange {
        hwnd: HWND,
    },
    NameChange {
        hwnd: HWND,
        idobject: OBJECT_IDENTIFIER,
        idchild: u32,
    },
    LocationChange {
        hwnd: HWND,
        idobject: OBJECT_IDENTIFIER,
        idchild: u32,
    },
    Destroyed {
        hwnd: HWND,
        idobject: OBJECT_IDENTIFIER,
        idchild: u32,
    },
    Unrepresented(u32),
}

unsafe extern "system" fn global_hook(
    _hwineventhook: HWINEVENTHOOK,
    event: u32,
    hwnd: HWND,
    idobject: i32,
    idchild: i32,
    _ideventthread: u32,
    _dwmseventtime: u32,
) {
    if let Ok(mut runtime) =
        (*GLOBAL_CALLBACK_SENDER.load(std::sync::atomic::Ordering::Relaxed)).lock()
    {
        runtime.on_event(Event::from_data(
            event,
            hwnd,
            OBJECT_IDENTIFIER(idobject),
            idchild as u32,
        ));
    }
}

unsafe extern "system" fn global_timer_hook(
    _hwnd: HWND,
    _msg: u32,
    _timer_id: usize,
    _dwms_event_time: u32,
) {
    if let Ok(mut runtime) =
        (*GLOBAL_CALLBACK_SENDER.load(std::sync::atomic::Ordering::Relaxed)).lock()
    {
        let system_foreground_window = GetForegroundWindow();
        if runtime.foreground_window != system_foreground_window
            && runtime.msaa_check_window_focused_state(system_foreground_window)
        {
            runtime.on_new_foreground(system_foreground_window);
        }
    }
}

struct Runtime {
    foreground_change: HookInfo,
    unminimize_change: HookInfo,

    foreground_name_change: HookInfo,
    foreground_window: HWND,

    target: TargetWindow,

    underlay_tx: std::sync::mpsc::Sender<UnderlayEvent>,
    uipi_msg_id: u32,
}

struct TargetWindow {
    hwnd: HWND,
    title: String,
    location_change: HookInfo,
    destroyed_change: HookInfo,
    is_destroyed: bool,
    is_focused: bool,
}

trait Valid {
    fn is_valid(&self) -> bool;
}

impl Valid for HWND {
    fn is_valid(&self) -> bool {
        *self != Default::default()
    }
}

impl Runtime {
    fn on_event(&mut self, event: Event) {
        log::trace!("Got event: {event:?}");
        match event {
            Event::ForegroundChange { hwnd } => {
                log::trace!("ForegroundChange: {hwnd:?}");

                //Sometimes focus changes doesn't keep up, check manually.
                if unsafe { GetForegroundWindow() } != hwnd {
                    if !self.msaa_check_window_focused_state(hwnd) {
                        return;
                    }
                }

                self.on_new_foreground(hwnd)
            }
            Event::NameChange {
                hwnd,
                idobject,
                idchild,
            } => {
                log::trace!("NameChange: {hwnd:?}");
                if self.target.hwnd == hwnd && idobject == OBJID_WINDOW && idchild == CHILDID_SELF {
                    self.check_and_handle(self.foreground_window)
                }
            }
            Event::LocationChange {
                hwnd,
                idobject,
                idchild,
            } => {
                log::trace!("LocationChange: {hwnd:?}");
                if self.target.hwnd == hwnd && idobject == OBJID_WINDOW && idchild == CHILDID_SELF {
                    self.handle_movesize();
                }
            }
            Event::Destroyed {
                hwnd,
                idobject,
                idchild,
            } => {
                log::trace!("Destroy: {hwnd:?}");
                if self.target.hwnd == hwnd && idobject == OBJID_WINDOW && idchild == CHILDID_SELF {
                    self.target.is_destroyed = true;
                    log::trace!("target.hwnd == destroyed hwnd");
                    self.check_and_handle(Default::default());
                }
            }
            Event::Unrepresented(event_code) => {
                log::trace!("Unrepresented event: {event_code}");
            }
        };
    }

    fn on_new_foreground(&mut self, hwnd: HWND) {
        log::trace!(
            "Got new foreground: {hwnd:?} - {}",
            self.get_title(hwnd).unwrap_or("".into())
        );
        self.foreground_window = hwnd;

        self.foreground_name_change.unbind();

        if self.foreground_window.is_valid() && self.foreground_window == self.target.hwnd {
            self.foreground_name_change.idthread =
                unsafe { GetWindowThreadProcessId(self.foreground_window, None) };
            self.foreground_name_change.bind();
        }

        self.check_and_handle(self.foreground_window);
    }

    fn get_title(&mut self, window: HWND) -> Result<String, Error> {
        //Allow some extra space in case of undocumented higher limit / little bit of safety for newer versions
        const BUFSIZE: usize = 512;
        let mut buffer: [u16; 512] = [0_u16; BUFSIZE];

        unsafe { SetLastError(Default::default()) };
        let num_stored = unsafe { GetWindowTextW(window, &mut buffer) };
        if num_stored == 0 {
            unsafe { GetLastError() }.ok()?;
            return Ok("".into());
        }

        Ok(String::from_utf16_lossy(&buffer)
            .trim_end_matches(char::from(0))
            .to_string())
    }

    fn check_and_handle(&mut self, hwnd: HWND) {
        if unsafe { IsHungAppWindow(hwnd).into() } {
            return;
        }

        if self.target.hwnd.is_valid() {
            if self.target.hwnd == hwnd {
                if !self.target.is_focused {
                    self.target.is_focused = true;
                    self.underlay_tx.send(UnderlayEvent::Focus).ok();
                }
                return;
            }

            if self.target.is_focused {
                self.target.is_focused = false;
                self.underlay_tx.send(UnderlayEvent::Blur).ok();
            }

            if self.target.is_destroyed {
                self.target.hwnd = Default::default();
                self.target.is_destroyed = false;
                self.underlay_tx.send(UnderlayEvent::Detach).ok();
            }
        }

        log::trace!("Checking and handling pre get_title");

        match self.get_title(hwnd) {
            Ok(title) if title == self.target.title => {}
            _ => return,
        }

        log::info!("Got target window!");

        if self.target.hwnd.is_valid() {
            self.target.location_change.unbind();
            self.target.destroyed_change.unbind();
        }

        self.target.hwnd = hwnd;
        let thread_id = unsafe { GetWindowThreadProcessId(hwnd, None) };

        self.target.location_change.idthread = thread_id;
        self.target.location_change.bind();
        self.target.destroyed_change.idthread = thread_id;
        self.target.destroyed_change.bind();

        let has_access = self.has_uipi_access();
        if let Some(bounds) = self.get_content_bounds() {
            self.underlay_tx
                .send(UnderlayEvent::Attach { bounds, has_access })
                .ok();

            self.target.is_focused = true;
            self.underlay_tx.send(UnderlayEvent::Focus).ok();
        } else {
            self.target.hwnd = Default::default();
        }
    }

    fn msaa_check_window_focused_state(&mut self, hwnd: HWND) -> bool {
        log::trace!("msaa_check_window_focused_state");
        let mut var_child_self = unsafe { VariantInit() };
        let mut p_acc: Option<IAccessible> = None;
        let hr = unsafe {
            AccessibleObjectFromEvent(
                hwnd,
                OBJID_WINDOW.0 as u32,
                CHILDID_SELF,
                &mut p_acc,
                &mut var_child_self,
            )
        };
        if hr.is_err() || p_acc.is_none() {
            unsafe { VariantClear(&mut var_child_self) }.ok();
            return false;
        }
        let p_acc = p_acc.unwrap();
        let var_state = unsafe { p_acc.get_accState(var_child_self) };

        let mut is_focused = false;

        if let Ok(mut var_state) = var_state {
            if unsafe { var_state.Anonymous.Anonymous.vt } == VT_I4 {
                let value = unsafe { var_state.Anonymous.Anonymous.Anonymous.lVal };
                is_focused = (value & STATE_SYSTEM_FOCUSED as i32) != 0;
            }
            unsafe {
                VariantClear(&mut var_state).ok();
            };
        }

        return is_focused;
    }

    fn handle_movesize(&self) {
        if let Some(bounds) = self.get_content_bounds() {
            self.underlay_tx
                .send(UnderlayEvent::MoveResize(bounds))
                .ok();
        }
    }

    fn has_uipi_access(&self) -> bool {
        unsafe {
            SetLastError(Default::default());
            PostMessageW(
                self.target.hwnd,
                self.uipi_msg_id,
                WPARAM::default(),
                LPARAM::default(),
            );
            GetLastError() != ERROR_ACCESS_DENIED
        }
    }

    fn get_content_bounds(&self) -> Option<Bounds> {
        let mut rect: RECT = unsafe { std::mem::zeroed() };
        if !unsafe { GetClientRect(self.target.hwnd, &mut rect).into() } {
            return None;
        }

        let mut pt_client_ul: POINT = POINT {
            x: rect.left,
            y: rect.top,
        };
        if !unsafe { ClientToScreen(self.target.hwnd, &mut pt_client_ul).into() } {
            return None;
        }

        Some(Bounds {
            x: pt_client_ul.x,
            y: pt_client_ul.y,
            width: rect.right,
            height: rect.bottom,
        })
    }
}

pub fn init(
    window_title: String,
    tx: std::sync::mpsc::Sender<crate::UnderlayEvent>,
) -> Result<(), String> {
    std::thread::spawn(move || {
        let uipi_msg_id = register_message("UNDERLAY_UIPI_TEST");

        let mut runtime = Runtime {
            foreground_change: HookInfo::new(EVENT_SYSTEM_FOREGROUND, EVENT_SYSTEM_FOREGROUND),
            unminimize_change: HookInfo::new(EVENT_SYSTEM_MINIMIZEEND, EVENT_SYSTEM_MINIMIZEEND),

            foreground_name_change: HookInfo::new(EVENT_OBJECT_NAMECHANGE, EVENT_OBJECT_NAMECHANGE),
            foreground_window: Default::default(),

            target: TargetWindow {
                hwnd: Default::default(),
                title: window_title,
                location_change: HookInfo::new(
                    EVENT_OBJECT_LOCATIONCHANGE,
                    EVENT_OBJECT_LOCATIONCHANGE,
                ),
                destroyed_change: HookInfo::new(EVENT_OBJECT_DESTROY, EVENT_OBJECT_DESTROY),
                is_destroyed: false,
                is_focused: false,
            },

            uipi_msg_id,

            underlay_tx: tx,
        };
        runtime.foreground_change.bind();
        runtime.unminimize_change.bind();

        GLOBAL_CALLBACK_SENDER.store(
            Box::into_raw(Box::new(Arc::new(Mutex::new(runtime)))),
            std::sync::atomic::Ordering::Relaxed,
        );

        let mut msg = MSG::default();
        while unsafe { GetMessageW(&mut msg, HWND(0), 0, 0) } != windows::Win32::Foundation::FALSE {
            log::trace!("Got message");
            unsafe { TranslateMessage(&msg) };
            log::trace!("Translated message");
            unsafe { DispatchMessageW(&msg) };
            log::trace!("Dispatched message");
        }
        log::trace!("GetMessageW caused loop to exit");
    });

    Ok(())
}

fn register_message<T: Into<String>>(message: T) -> u32 {
    let message_wide = windows::core::HSTRING::from(message.into());
    let result = unsafe { RegisterWindowMessageW(PCWSTR(message_wide.as_wide().as_ptr())) };
    result
}