first-draft windows hotplug support using subclass.

windows_hotplug.rs

use bevy::prelude::*;
use bevy::{input::keyboard::KeyboardInput, prelude::*};
use bevy::winit::WinitWindows;
use winit::platform::windows::WindowExtWindows;
use bevy::window::WindowCreated;
use bevy::app::{Events, ManualEventReader};
use winapi::{
    DEFINE_GUID,
    ctypes::c_void,
    shared::{
        basetsd::{UINT_PTR, DWORD_PTR},
        minwindef::{BOOL, DWORD, HIWORD, INT, LOWORD, LPARAM, LRESULT, UINT, WORD, WPARAM},
        usbiodef,
        windef::{HWND, POINT, RECT},
        windowsx, winerror,
    },
    um::{
        commctrl, dbt, libloaderapi, ole2, processthreadsapi, winbase,
        winnt::{HANDLE, LONG, LPCSTR, SHORT},
        winuser,
    },
};
use std::cell::{Cell, RefCell};
use std::mem;
use std::ffi::{CStr, OsString};
use widestring::{U16Str, U16CStr};
use regex::Regex;
use crossbeam_channel::{unbounded, Sender, Receiver, SendError, RecvError, TryRecvError, SendTimeoutError};
use super::GcnHotplugEvent;

const SUBCLASS_ID: UINT_PTR = 0;

struct GcnSubclass {
    // These are needed to track subclass state
    pub subclass_removed: bool,
    pub recurse_depth: u32,
    // Our stuff
    regex: Regex,
    tx: Sender<GcnHotplugEvent>,
}

impl GcnSubclass {
    pub fn new(tx: Sender<GcnHotplugEvent>) -> Self {
        Self {
            subclass_removed: false,
            recurse_depth: 0,
            regex: Regex::new(r"^\\\\\?\\USB#VID_([0-9a-fA-F]+)&PID_([0-9a-fA-F]+)#").unwrap(),
            tx
        }
    }
}

pub struct GcnHotplug {
    hwnd: HWND,
    pub rx: Receiver<GcnHotplugEvent>,
}

/// If a hotplug system exists, will pipe events to bevy
/// Otherwise will attempt to set one up
pub fn hotplug_system(world: &mut World) {
    let mut detached = false;
    if let Some(hotplug) = world.get_non_send_resource::<GcnHotplug>() {
        let mut events = unsafe { world.get_resource_unchecked_mut::<Events<GcnHotplugEvent>>().unwrap() };
        loop {
            match hotplug.rx.try_recv() {
                Ok(event) => {
                    events.send(event);

                    if let GcnHotplugEvent::Attached = event {
                        // attempt to connect to new device
                    }
                },
                Err(TryRecvError::Disconnected) => {
                    debug!("hotplug detached!");
                    detached = true;
                },
                Err(TryRecvError::Empty) => break,
            }
        }
    }
    else {
        let winwin = world.get_resource::<WinitWindows>().unwrap();
        if let Some(window) = winwin.windows.values().nth(0) {
            let hwnd = window.hwnd() as HWND;

            let (tx, rx) = unbounded::<GcnHotplugEvent>();

            // safety: winit gives us a valid hwnd and we're on the main thread
            unsafe {
                setup_subclass(hwnd, tx);
            }

            world.insert_non_send(GcnHotplug { hwnd, rx });
        }
        else {
            debug!("hotplug cannot attach");
        }
    }

    if detached {
        world.remove_non_send::<GcnHotplug>();
    }
}

// SAFETY: you must pass in a valid hwnd and be on the main thread
unsafe fn setup_subclass(hwnd: HWND, tx: Sender<GcnHotplugEvent>) {
    let subclass_result = commctrl::SetWindowSubclass(
        hwnd,
        Some(subclass_callback),
        0,
        Box::into_raw(Box::new(GcnSubclass::new(tx))) as DWORD_PTR,
    );
    assert_eq!(subclass_result, 1);

    let mut filter = dbt::DEV_BROADCAST_DEVICEINTERFACE_A {
        dbcc_size: mem::size_of::<dbt::DEV_BROADCAST_DEVICEINTERFACE_A>() as u32,
        dbcc_devicetype: dbt::DBT_DEVTYP_DEVICEINTERFACE,
        dbcc_reserved: 0,
        dbcc_classguid: usbiodef::GUID_DEVINTERFACE_USB_DEVICE,
        dbcc_name: [0]
    };

    winuser::RegisterDeviceNotificationA(hwnd as *mut c_void, &mut filter as *mut _ as *mut c_void, winuser::DEVICE_NOTIFY_WINDOW_HANDLE);
}

// SAFETY: you are Windows
unsafe extern "system" fn subclass_callback(
    hwnd: HWND,
    msg: UINT,
    wparam: WPARAM,
    lparam: LPARAM,
    _: UINT_PTR,
    subclass_ptr: DWORD_PTR,
) -> LRESULT {
    let subclass_ptr = subclass_ptr as *mut GcnSubclass;

    (*subclass_ptr).recurse_depth += 1;

    let ret = match msg {
        winuser::WM_DEVICECHANGE => {
            match wparam {
                dbt::DBT_DEVICEARRIVAL | dbt::DBT_DEVICEREMOVECOMPLETE => {
                    let data = unsafe { &mut *(lparam as *mut dbt::DEV_BROADCAST_HDR) };
                    match data.dbch_devicetype {
                        dbt::DBT_DEVTYP_DEVICEINTERFACE => {
                            // A lot of bs to extract useful information about the device
                            let data = unsafe { &mut *(lparam as *mut dbt::DEV_BROADCAST_DEVICEINTERFACE_W) };
                            let name = U16CStr::from_ptr_str(&data.dbcc_name as *const u16).to_string_lossy();
                            let captures = (*subclass_ptr).regex.captures(&name).unwrap();
                            // A simple if-else since we only have two possible events by now
                            let msg = if wparam == dbt::DBT_DEVICEARRIVAL { GcnHotplugEvent::Attached } else { GcnHotplugEvent::Detached };

                            // Filter for our specific device
                            if let (Some(vid), Some(pid)) = (captures.get(1), captures.get(2)) {
                                if vid.as_str() == "057E" && pid.as_str() == "2009" {
                                    // If this disconnects, the subclass doesn't need to exist anymore
                                    if let Err(SendError(_)) = (*subclass_ptr).tx.send(msg) {
                                        remove_subclass(hwnd);
                                    }
                                }
                            }
                        },
                        // this shouldn't happen, so we'd be interested to know if it does
                        e => { debug!("excuse me? {}", e) }
                    }

                    0
                },
                _ => commctrl::DefSubclassProc(hwnd, msg, wparam, lparam),
            }
        }
        winuser::WM_NCDESTROY => {
            remove_subclass(hwnd);
            (*subclass_ptr).subclass_removed = true;

            0
        }
        _ => commctrl::DefSubclassProc(hwnd, msg, wparam, lparam),
    };

    (*subclass_ptr).recurse_depth -= 1;

    if let GcnSubclass { subclass_removed: true, recurse_depth: 0, .. } = *subclass_ptr {
        Box::from_raw(subclass_ptr);
    }

    ret
}

// Safety: must pass a valid hwnd and be on the main thread
unsafe fn remove_subclass(hwnd: HWND) {
    let res = commctrl::RemoveWindowSubclass(
        hwnd,
        Some(subclass_callback),
        SUBCLASS_ID,
    );
    assert_eq!(res, 1);
}