punktfunk/crates/punktfunk-host/src/inject/windows/gamepad_windows.rs

//! Windows virtual Xbox 360 gamepad via the punktfunk **XUSB companion** UMDF driver
//! (`packaging/windows/xusb-driver`) — the in-tree replacement for ViGEmBus. One virtual Xbox 360
//! controller per client pad index, visible to classic **XInput** (`XInputGetState`) with no kernel
//! bus driver: each pad `SwDeviceCreate`s a `pf_xusb_<index>` devnode (the driver loads on it and
//! registers `GUID_DEVINTERFACE_XUSB`) and the host pushes the XInput state into the shared section
//! `Global\pfxusb-shm-<index>`. GameStream/Moonlight already speak the XInput conventions (low-16
//! button bits, sticks −32768..32767 +Y up, triggers 0..255), so the state copy is ~1:1.
//!
//! Rumble flows back the other way: a game writes force-feedback via `XInputSetState`, the driver
//! parses the `SET_STATE` packet into the shared section, and [`GamepadManager::pump_rumble`] relays
//! level changes to the client (the universal 0xCA plane), mirroring the Linux `EV_FF` read path.
//!
//! NB: the driver currently maps `Global\pfxusb-shm-0` (hardcoded), so a single pad (index 0) is
//! fully correct; mixed multi-pad needs the driver to read its own index first (same limitation as
//! the DualSense backend).

use crate::gamestream::gamepad::{GamepadEvent, MAX_PADS};
use anyhow::{anyhow, Result};
use std::ffi::c_void;
use windows::core::{w, GUID, HRESULT, HSTRING, PCWSTR};
use windows::Win32::Devices::Enumeration::Pnp::{
    SwDeviceClose, SwDeviceCreate, HSWDEVICE, SW_DEVICE_CREATE_INFO,
};
use windows::Win32::Foundation::{CloseHandle, E_FAIL, HANDLE, WAIT_OBJECT_0};
use windows::Win32::System::Threading::{CreateEventW, SetEvent, WaitForSingleObject};

// Shared-section layout — the single source of truth is `pf_driver_proto::gamepad::XusbShm` (offset
// asserts pin every field; the `pf_xusb` driver maps the same struct). Derive the size/offsets/magic from
// it so a layout change is a compile error, not a hand-synced literal (audit §6.1).
use pf_driver_proto::gamepad::XusbShm;
const SHM_SIZE: usize = core::mem::size_of::<XusbShm>();
const SHM_MAGIC: u32 = pf_driver_proto::gamepad::XUSB_MAGIC; // "PFXU"
const OFF_PACKET: usize = core::mem::offset_of!(XusbShm, packet);
const OFF_BUTTONS: usize = core::mem::offset_of!(XusbShm, buttons);
const OFF_LT: usize = core::mem::offset_of!(XusbShm, left_trigger);
const OFF_RT: usize = core::mem::offset_of!(XusbShm, right_trigger);
const OFF_LX: usize = core::mem::offset_of!(XusbShm, thumb_lx);
const OFF_LY: usize = core::mem::offset_of!(XusbShm, thumb_ly);
const OFF_RX: usize = core::mem::offset_of!(XusbShm, thumb_rx);
const OFF_RY: usize = core::mem::offset_of!(XusbShm, thumb_ry);
const OFF_RUMBLE_SEQ: usize = core::mem::offset_of!(XusbShm, rumble_seq);
const OFF_RUMBLE: usize = core::mem::offset_of!(XusbShm, rumble_large); // large @28, small @29
const OFF_DRIVER_PROTO: usize = core::mem::offset_of!(XusbShm, driver_proto);

/// Context for the `SwDeviceCreate` completion callback: an event to signal, the HRESULT it reports,
/// and the PnP instance id PnP assigned (captured for devnode health diagnostics).
#[repr(C)]
struct SwCreateCtx {
    event: HANDLE,
    result: HRESULT,
    instance_id: [u16; 128],
}

/// `SwDeviceCreate` fires this once PnP has enumerated the device; stash the result + wake the creator.
unsafe extern "system" fn sw_create_cb(
    _dev: HSWDEVICE,
    result: HRESULT,
    ctx: *const c_void,
    id: PCWSTR,
) {
    if !ctx.is_null() {
        // SAFETY: ctx is the &mut SwCreateCtx the creator passed; it outlives this callback (the
        // creator blocks on the event). `id` is a NUL-terminated string for the callback's duration.
        unsafe {
            let c = ctx as *mut SwCreateCtx;
            (*c).result = result;
            if !id.is_null() {
                for i in 0..(*c).instance_id.len() - 1 {
                    let ch = *id.0.add(i);
                    (*c).instance_id[i] = ch;
                    if ch == 0 {
                        break;
                    }
                }
            }
            let _ = SetEvent((*c).event);
        }
    }
}

impl SwCreateCtx {
    fn instance_id(&self) -> Option<String> {
        let len = self.instance_id.iter().position(|&c| c == 0)?;
        (len > 0).then(|| String::from_utf16_lossy(&self.instance_id[..len]))
    }
}

/// Spawn the `pf_xusb_<index>` companion devnode (hardware id `pf_xusb`, enumerator `punktfunk`). The
/// INF (System class) binds our UMDF driver, which registers the XUSB interface. Unlike the HID pads,
/// no USB compatible-ids are needed — XInput finds the device by the interface GUID, not VID/PID — but
/// we still pass a deterministic non-null `pContainerId` (the null-sentinel trips an `xinput1_4`
/// slot-skip bug). `SwDeviceClose` removes it on drop.
fn create_swdevice(index: u8) -> Result<(HSWDEVICE, Option<String>)> {
    let hwids: Vec<u16> = "pf_xusb".encode_utf16().chain([0u16, 0u16]).collect();
    let instid: Vec<u16> = format!("pf_xusb_{index}")
        .encode_utf16()
        .chain(std::iter::once(0))
        .collect();
    let desc: Vec<u16> = "punktfunk Virtual Xbox 360 (XUSB)"
        .encode_utf16()
        .chain(std::iter::once(0))
        .collect();
    // The pad index, stamped into the device Location — the driver reads it to map `pfxusb-shm-<index>`
    // (multi-pad). The buffer must outlive the SwDeviceCreate call (it does; we wait on the event).
    let loc: Vec<u16> = format!("{index}")
        .encode_utf16()
        .chain(std::iter::once(0))
        .collect();
    let container = GUID::from_values(0x5046_5855, 0x0000, 0x0000, [0, 0, 0, 0, 0, 0, 0, index]);

    // SAFETY: zeroed then the fields we use are set; the buffers + container outlive the call.
    let mut info: SW_DEVICE_CREATE_INFO = unsafe { std::mem::zeroed() };
    info.cbSize = std::mem::size_of::<SW_DEVICE_CREATE_INFO>() as u32;
    info.pszInstanceId = PCWSTR(instid.as_ptr());
    info.pszzHardwareIds = PCWSTR(hwids.as_ptr());
    info.pContainerId = &container;
    info.pszDeviceDescription = PCWSTR(desc.as_ptr());
    info.pszDeviceLocation = PCWSTR(loc.as_ptr());
    info.CapabilityFlags = 0x0000_000B; // DriverRequired | SilentInstall | Removable

    // SAFETY: a manual-reset, initially-unsignaled, unnamed event.
    let event = unsafe { CreateEventW(None, true, false, PCWSTR::null())? };
    // `result` starts as E_FAIL, NOT S_OK: if the wait below times out, a zero-initialised HRESULT
    // would read as success and mask the failure (found by the 2026-07 driver-health audit).
    let mut ctx = SwCreateCtx {
        event,
        result: E_FAIL,
        instance_id: [0; 128],
    };
    // SAFETY: info + buffers + ctx outlive the call (we wait on the event before returning).
    let hsw = match unsafe {
        SwDeviceCreate(
            w!("punktfunk"),
            w!("HTREE\\ROOT\\0"),
            &info,
            None,
            Some(sw_create_cb),
            Some(&mut ctx as *mut SwCreateCtx as *const c_void),
        )
    } {
        Ok(h) => h,
        Err(e) => {
            // SAFETY: event is valid.
            unsafe {
                let _ = CloseHandle(event);
            }
            return Err(anyhow!("SwDeviceCreate(pf_xusb) failed: {e}"));
        }
    };
    // SAFETY: event valid; block until PnP finishes enumerating, then check the callback result.
    let wait = unsafe { WaitForSingleObject(event, 10_000) };
    // SAFETY: event is valid.
    unsafe {
        let _ = CloseHandle(event);
    }
    if wait != WAIT_OBJECT_0 {
        // SAFETY: hsw is the handle SwDeviceCreate returned.
        unsafe { SwDeviceClose(hsw) };
        return Err(anyhow!(
            "SwDeviceCreate(pf_xusb) enumeration callback never fired (10s) — PnP may be wedged"
        ));
    }
    if ctx.result.is_err() {
        // SAFETY: hsw is the handle SwDeviceCreate returned.
        unsafe { SwDeviceClose(hsw) };
        return Err(anyhow!(
            "SwDeviceCreate(pf_xusb) enumeration failed: {:?}",
            ctx.result
        ));
    }
    Ok((hsw, ctx.instance_id()))
}

/// A single virtual Xbox 360 pad: the `pf_xusb_<index>` devnode plus the mapped shared section.
struct XusbWinPad {
    /// Owns the `pf_xusb_<index>` devnode (dropped → `SwDeviceClose`). `None` if `SwDeviceCreate` failed.
    _sw: Option<super::gamepad_raii::SwDevice>,
    /// Owns `Global\pfxusb-shm-<index>` (the section + its mapped view; drop unmaps + closes).
    shm: super::gamepad_raii::Shm,
    /// Watches the section's `driver_proto` field and logs attach / never-attached diagnosis.
    attach: super::gamepad_raii::DriverAttach,
    packet: u32,
    last_rumble_seq: u32,
}

impl XusbWinPad {
    /// Create + map `Global\pfxusb-shm-<index>`, stamp the magic, then spawn the devnode.
    fn open(index: u8) -> Result<XusbWinPad> {
        // Permissive-DACL named section the WUDFHost (whatever account) can open; `Shm` owns the
        // section handle + its mapped view (zero-filled) and unmaps/closes on drop.
        let shm_name = pf_driver_proto::gamepad::xusb_shm_name(index);
        let shm = super::gamepad_raii::Shm::create(&HSTRING::from(shm_name.as_str()), SHM_SIZE)?;
        let base = shm.base();
        // Zero the section then stamp the magic LAST (the driver only accepts it once magic is set).
        // SAFETY: base points at SHM_SIZE writable bytes.
        unsafe {
            std::ptr::write_bytes(base, 0, SHM_SIZE);
            std::ptr::write_unaligned(base as *mut u32, SHM_MAGIC);
        }
        let (hsw, instance_id) = match create_swdevice(index) {
            Ok((h, id)) => (Some(h), id),
            Err(e) => {
                tracing::warn!(error = %format!("{e:#}"), "SwDeviceCreate failed; XUSB devnode unavailable");
                (None, None)
            }
        };
        let _sw = hsw.map(super::gamepad_raii::SwDevice::new);
        Ok(XusbWinPad {
            _sw,
            shm,
            attach: super::gamepad_raii::DriverAttach::new(
                "pf_xusb",
                "pf_xusb.inf",
                "C:\\Users\\Public\\pfxusb-driver.log",
                shm_name,
                instance_id,
            ),
            packet: 0,
            last_rumble_seq: 0,
        })
    }

    /// Publish the XInput state to the section and bump the packet number (XInput uses it to detect
    /// change). `buttons` is the XINPUT_GAMEPAD_* bitmap; sticks are i16, triggers u8.
    #[allow(clippy::too_many_arguments)]
    fn write_state(&mut self, buttons: u16, lt: u8, rt: u8, lx: i16, ly: i16, rx: i16, ry: i16) {
        self.packet = self.packet.wrapping_add(1);
        let base = self.shm.base();
        // SAFETY: `base` is the start of the mapped section (`SHM_SIZE` bytes, owned by `Shm`); every
        // `OFF_*` is a fixed in-range offset into it and `write_unaligned` handles the unaligned field
        // writes. Single owner (`&mut self`), so no concurrent writer races these stores.
        unsafe {
            std::ptr::write_unaligned(base.add(OFF_BUTTONS) as *mut u16, buttons);
            *base.add(OFF_LT) = lt;
            *base.add(OFF_RT) = rt;
            std::ptr::write_unaligned(base.add(OFF_LX) as *mut i16, lx);
            std::ptr::write_unaligned(base.add(OFF_LY) as *mut i16, ly);
            std::ptr::write_unaligned(base.add(OFF_RX) as *mut i16, rx);
            std::ptr::write_unaligned(base.add(OFF_RY) as *mut i16, ry);
            std::ptr::write_unaligned(base.add(OFF_PACKET) as *mut u32, self.packet);
        }
    }

    /// Poll the section for a game's rumble (the driver bumps `rumble_seq` on each SET_STATE). Returns
    /// `(large, small)` motor levels (0..=255) when a new one arrived. Also feeds the driver-attach
    /// health watcher (the driver stamps `driver_proto` at device add + on every serviced IOCTL).
    fn service(&mut self) -> Option<(u8, u8)> {
        let base = self.shm.base();
        // SAFETY: base points at SHM_SIZE bytes.
        let proto = unsafe { std::ptr::read_unaligned(base.add(OFF_DRIVER_PROTO) as *const u32) };
        self.attach.observe(proto);
        // SAFETY: base points at SHM_SIZE bytes.
        let seq = unsafe { std::ptr::read_unaligned(base.add(OFF_RUMBLE_SEQ) as *const u32) };
        if seq == self.last_rumble_seq {
            return None;
        }
        self.last_rumble_seq = seq;
        // SAFETY: rumble bytes at OFF_RUMBLE / OFF_RUMBLE+1.
        let (large, small) = unsafe { (*base.add(OFF_RUMBLE), *base.add(OFF_RUMBLE + 1)) };
        Some((large, small))
    }
}

/// All virtual Xbox 360 pads of a session — the Windows analogue of the Linux uinput-xpad manager,
/// now backed by the XUSB companion driver. Same method surface (`new`/`handle`/`pump_rumble`) the
/// session input thread already drives.
pub struct GamepadManager {
    pads: Vec<Option<XusbWinPad>>,
    last_rumble: Vec<(u8, u8)>,
    broken: bool,
}

impl Default for GamepadManager {
    fn default() -> GamepadManager {
        GamepadManager::new()
    }
}

impl GamepadManager {
    pub fn new() -> GamepadManager {
        GamepadManager {
            pads: (0..MAX_PADS).map(|_| None).collect(),
            last_rumble: vec![(0, 0); MAX_PADS],
            broken: false,
        }
    }

    fn ensure(&mut self, idx: usize) {
        if idx >= MAX_PADS || self.pads[idx].is_some() || self.broken {
            return;
        }
        match XusbWinPad::open(idx as u8) {
            Ok(p) => {
                tracing::info!(
                    index = idx,
                    "virtual Xbox 360 created (Windows XUSB companion)"
                );
                self.pads[idx] = Some(p);
                self.last_rumble[idx] = (0, 0);
            }
            Err(e) => {
                tracing::error!(error = %format!("{e:#}"), "virtual Xbox 360 creation failed — controller input disabled until the next client connect (install/repair: punktfunk-host.exe driver install --gamepad)");
                self.broken = true;
            }
        }
    }

    pub fn handle(&mut self, ev: &GamepadEvent) {
        let GamepadEvent::State(f) = ev else {
            return; // Arrival metadata — the pad is created lazily on the first State
        };
        let idx = f.index.max(0) as usize;
        if idx >= MAX_PADS {
            return;
        }
        // Unplugs: drop any allocated pad whose mask bit cleared.
        for (i, slot) in self.pads.iter_mut().enumerate() {
            if slot.is_some() && f.active_mask & (1 << i) == 0 {
                tracing::info!(index = i, "controller unplugged (Xbox 360/Windows)");
                *slot = None;
                self.last_rumble[i] = (0, 0);
            }
        }
        if f.active_mask & (1 << idx) == 0 {
            return;
        }
        self.ensure(idx);
        if let Some(pad) = self.pads[idx].as_mut() {
            pad.write_state(
                (f.buttons & 0xffff) as u16,
                f.left_trigger,
                f.right_trigger,
                f.ls_x,
                f.ls_y,
                f.rs_x,
                f.rs_y,
            );
        }
    }

    /// Relay any changed rumble level to the client. XUSB motors are 0..255; the wire carries
    /// 0..65535, so scale by 257. `large` (low-frequency) → the datagram's `low`, `small`
    /// (high-frequency) → `high` — matching the other backends.
    pub fn pump_rumble(&mut self, mut send: impl FnMut(u16, u16, u16)) {
        for i in 0..self.pads.len() {
            let Some(pad) = self.pads[i].as_mut() else {
                continue;
            };
            if let Some((large, small)) = pad.service() {
                if self.last_rumble[i] != (large, small) {
                    self.last_rumble[i] = (large, small);
                    send(i as u16, large as u16 * 257, small as u16 * 257);
                }
            }
        }
    }
}