punktfunk/packaging/windows/vdisplay-driver/pf-vdisplay/src/control.rs

//! SudoVDA-compatible IOCTL control plane (`EVT_IDD_CX_DEVICE_IO_CONTROL`). The host's
//! `vdisplay/sudovda.rs` drives this unchanged: ADD a monitor at a requested mode → `{LUID, target_id}`,
//! REMOVE by GUID, PING the watchdog, GET_VERSION/GET_WATCHDOG, SET_RENDER_ADAPTER. Struct layouts are
//! byte-identical to `Common/Include/sudovda-ioctl.h`.

use std::ffi::c_void;
use std::mem::size_of;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use std::thread;
use std::time::{Duration, Instant};

use log::{error, info};
use wdf_umdf::{
    IddCxAdapterSetRenderAdapter, IddCxMonitorDeparture, WdfRequestCompleteWithInformation,
    WdfRequestRetrieveInputBuffer, WdfRequestRetrieveOutputBuffer,
};
use wdf_umdf_sys::{IDARG_IN_ADAPTERSETRENDERADAPTER, LUID, NTSTATUS, WDFDEVICE, WDFREQUEST};

use crate::context::{DeviceContext, MonitorContext};
use crate::monitor::{
    default_modes, Mode, MonitorData, MonitorObject, ADAPTER, MONITOR_MODES, NEXT_ID,
    PREFERRED_RENDER_ADAPTER, PROTOCOL_VERSION, WATCHDOG_COUNTDOWN, WATCHDOG_TIMEOUT,
};

// CTL_CODE(FILE_DEVICE_UNKNOWN=0x22, func, METHOD_BUFFERED=0, FILE_ANY_ACCESS=0).
const fn ctl(func: u32) -> u32 {
    (0x22u32 << 16) | (func << 2)
}
const IOCTL_ADD: u32 = ctl(0x800);
const IOCTL_REMOVE: u32 = ctl(0x801);
const IOCTL_SET_RENDER_ADAPTER: u32 = ctl(0x802);
const IOCTL_GET_WATCHDOG: u32 = ctl(0x803);
/// pf-vdisplay extension (NOT in SudoVDA): tear down every monitor. The host issues this on startup to
/// reap monitors orphaned by a crashed/killed previous host instance. SudoVDA returns invalid for it
/// (harmlessly ignored), so the host can send it unconditionally.
const IOCTL_CLEAR_ALL: u32 = ctl(0x804);
const IOCTL_PING: u32 = ctl(0x888);
const IOCTL_GET_VERSION: u32 = ctl(0x8FF);

/// Serializes monitor lifecycle ops — ADD / REMOVE / watchdog-teardown — against each other. Without
/// it, a watchdog expiry can drain an entry out from under an in-flight `do_add` (which releases the
/// `MONITOR_MODES` lock before the slow `create_monitor`), leaving `do_add` to return
/// `STATUS_UNSUCCESSFUL` → the host sees `ERROR_GEN_FAILURE`. This was the reconnect-churn fault.
static MONITOR_OP_LOCK: Mutex<()> = Mutex::new(());
/// A monitor created less than this ago is still in its host-side setup window (CCD commit + GDI-name
/// resolve + topology settle, ~5 s) and is never reaped by the watchdog — only by an explicit
/// CLEAR_ALL. Protects a freshly-born monitor from a transient PING gap during reconnect churn.
const MONITOR_GRACE: Duration = Duration::from_secs(6);

#[repr(C)]
struct AddParams {
    width: u32,
    height: u32,
    refresh: u32,
    guid: [u8; 16],
    device_name: [u8; 14],
    serial: [u8; 14],
}
#[repr(C)]
struct AddOut {
    luid_low: u32,
    luid_high: i32,
    target_id: u32,
}
#[repr(C)]
struct RemoveParams {
    guid: [u8; 16],
}
#[repr(C)]
struct SetRenderAdapterParams {
    luid_low: u32,
    luid_high: i32,
}
#[repr(C)]
struct WatchdogOut {
    timeout: u32,
    countdown: u32,
}

fn guid_key(b: &[u8; 16]) -> u128 {
    u128::from_le_bytes(*b)
}

/// SAFETY: `request` valid; returns a pointer to the request's input buffer of at least `min` bytes.
unsafe fn input_buf(request: WDFREQUEST, min: usize) -> Option<*const u8> {
    let mut p: *mut c_void = std::ptr::null_mut();
    let mut len: usize = 0;
    let r = unsafe { WdfRequestRetrieveInputBuffer(request, min, &mut p, &mut len) };
    if r.is_err() || p.is_null() || len < min {
        return None;
    }
    Some(p.cast::<u8>())
}

/// SAFETY: `request` valid; returns a pointer to the request's output buffer of at least `min` bytes.
unsafe fn output_buf(request: WDFREQUEST, min: usize) -> Option<*mut u8> {
    let mut p: *mut c_void = std::ptr::null_mut();
    let mut len: usize = 0;
    let r = unsafe { WdfRequestRetrieveOutputBuffer(request, min, &mut p, &mut len) };
    if r.is_err() || p.is_null() || len < min {
        return None;
    }
    Some(p.cast::<u8>())
}

/// `EVT_IDD_CX_DEVICE_IO_CONTROL` — IddCx redirects device IOCTLs here. Signature matches SudoVDA's
/// `SudoVDAIoDeviceControl(Device, Request, OutputBufferLength, InputBufferLength, IoControlCode)`.
pub extern "C-unwind" fn device_io_control(
    device: WDFDEVICE,
    request: WDFREQUEST,
    output_len: usize,
    input_len: usize,
    ioctl_code: u32,
) {
    // Reset the watchdog on any IOCTL except the watchdog query (the host PINGs to keep alive).
    if ioctl_code != IOCTL_GET_WATCHDOG {
        WATCHDOG_COUNTDOWN.store(WATCHDOG_TIMEOUT.load(Ordering::Relaxed), Ordering::Relaxed);
    }

    let mut bytes: usize = 0;
    // SAFETY: dispatch reads/writes the request buffers it validated; `device` is the IddCx device.
    let status = unsafe {
        match ioctl_code {
            IOCTL_ADD => do_add(device, request, input_len, output_len, &mut bytes),
            IOCTL_REMOVE => do_remove(request, input_len),
            IOCTL_SET_RENDER_ADAPTER => do_set_render_adapter(request, input_len),
            IOCTL_GET_WATCHDOG => do_get_watchdog(request, output_len, &mut bytes),
            IOCTL_PING => NTSTATUS::STATUS_SUCCESS,
            IOCTL_CLEAR_ALL => {
                disconnect_all_monitors(true);
                NTSTATUS::STATUS_SUCCESS
            }
            IOCTL_GET_VERSION => do_get_version(request, output_len, &mut bytes),
            _ => NTSTATUS::STATUS_INVALID_DEVICE_REQUEST,
        }
    };

    // SAFETY: completing the request we were handed.
    let _ = unsafe { WdfRequestCompleteWithInformation(request, status, bytes as u64) };
}

unsafe fn do_add(
    device: WDFDEVICE,
    request: WDFREQUEST,
    input_len: usize,
    output_len: usize,
    bytes: &mut usize,
) -> NTSTATUS {
    // Serialize the whole ADD (push entry → create_monitor → verify) against the watchdog teardown +
    // REMOVE, so an expiry can never drain this entry mid-flight. `create_monitor` is fast (the slow
    // CCD/GDI work is host-side, after this returns), and PING/GET_WATCHDOG don't take this lock, so
    // the host keeps the watchdog reset while we hold it.
    let _op = MONITOR_OP_LOCK.lock().unwrap();
    if input_len < size_of::<AddParams>() || output_len < size_of::<AddOut>() {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    }
    let (Some(pin), Some(pout)) = (
        unsafe { input_buf(request, size_of::<AddParams>()) },
        unsafe { output_buf(request, size_of::<AddOut>()) },
    ) else {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    };
    let params = unsafe { &*pin.cast::<AddParams>() };
    let guid = guid_key(&params.guid);

    // Dedup: an existing GUID returns its LUID + target id (the host may re-ADD on reconnect).
    {
        let lock = MONITOR_MODES.lock().unwrap();
        if let Some(mon) = lock.iter().find(|m| m.guid == guid) {
            let out = AddOut {
                luid_low: mon.adapter_luid_low,
                luid_high: mon.adapter_luid_high,
                target_id: mon.target_id,
            };
            unsafe { pout.cast::<AddOut>().write_unaligned(out) };
            *bytes = size_of::<AddOut>();
            return NTSTATUS::STATUS_SUCCESS;
        }
    }

    if params.width == 0 || params.height == 0 || params.refresh == 0 {
        return NTSTATUS::STATUS_INVALID_PARAMETER;
    }

    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    // Requested mode first (preferred), then fallbacks.
    let mut modes = vec![Mode {
        width: params.width,
        height: params.height,
        refresh_rates: vec![params.refresh],
    }];
    modes.extend(default_modes());
    MONITOR_MODES.lock().unwrap().push(MonitorObject {
        object: None,
        data: MonitorData { id, modes },
        guid,
        target_id: 0,
        adapter_luid_low: 0,
        adapter_luid_high: 0,
        created_at: Instant::now(),
    });

    // Create the IddCx monitor via the device context (captures target id + LUID into the entry).
    let created = unsafe {
        DeviceContext::get_mut(device.cast(), |ctx| {
            if let Err(e) = ctx.create_monitor(id) {
                error!("ADD: create_monitor failed: {e:?}");
            }
        })
    };

    let lock = MONITOR_MODES.lock().unwrap();
    let mon = lock.iter().find(|m| m.data.id == id);
    if created.is_err() || mon.map_or(true, |m| m.object.is_none()) {
        drop(lock);
        MONITOR_MODES.lock().unwrap().retain(|m| m.data.id != id);
        error!("ADD: monitor {id} failed to arrive");
        return NTSTATUS::STATUS_UNSUCCESSFUL;
    }
    let mon = mon.unwrap();
    let out = AddOut {
        luid_low: mon.adapter_luid_low,
        luid_high: mon.adapter_luid_high,
        target_id: mon.target_id,
    };
    unsafe { pout.cast::<AddOut>().write_unaligned(out) };
    *bytes = size_of::<AddOut>();
    info!(
        "ADD {}x{}@{} -> target_id={} luid={:08x}:{:08x}",
        params.width, params.height, params.refresh, mon.target_id, mon.adapter_luid_high, mon.adapter_luid_low
    );
    NTSTATUS::STATUS_SUCCESS
}

unsafe fn do_remove(request: WDFREQUEST, input_len: usize) -> NTSTATUS {
    if input_len < size_of::<RemoveParams>() {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    }
    let Some(pin) = (unsafe { input_buf(request, size_of::<RemoveParams>()) }) else {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    };
    let params = unsafe { &*pin.cast::<RemoveParams>() };
    let guid = guid_key(&params.guid);

    // Serialize against ADD + watchdog teardown (lock order: OP_LOCK → MONITOR_MODES).
    let _op = MONITOR_OP_LOCK.lock().unwrap();
    let mon = {
        let mut lock = MONITOR_MODES.lock().unwrap();
        match lock.iter().position(|m| m.guid == guid) {
            Some(pos) => lock.remove(pos),
            None => return NTSTATUS::STATUS_NOT_FOUND,
        }
        // MONITOR_MODES released here — the processor-join + departure below must not hold it.
    };
    if let Some(obj) = mon.object {
        free_swap_chain_processor(obj.as_ptr());
        if let Err(e) = unsafe { IddCxMonitorDeparture(obj.as_ptr()) } {
            error!("REMOVE: departure failed: {e:?}");
        }
    }
    info!("REMOVE target_id={}", mon.target_id);
    NTSTATUS::STATUS_SUCCESS
}

/// Drop a monitor's live swap-chain processor BEFORE departure. The WDF context is an
/// `Arc<RwLock<MonitorContext>>` that WDF frees WITHOUT running Rust `Drop` (no `EvtCleanupCallback`
/// is wired), and the OS does not reliably call UNASSIGN on a host-initiated departure — so the
/// streaming `Direct3DDevice` (its ~dozens of D3D worker threads + tens of MB of VRAM) was orphaned
/// once per session, the dominant reconnect-churn leak. `get_mut` takes the context `RwLock`, so this
/// is safe against a concurrent OS unassign callback (whichever runs second sees `None`).
fn free_swap_chain_processor(monitor: *mut wdf_umdf_sys::IDDCX_MONITOR__) {
    // SAFETY: `monitor` is a live IddCx monitor object whose context was init'd at creation.
    let r = unsafe { MonitorContext::get_mut(monitor.cast(), |ctx| ctx.unassign_swap_chain()) };
    if let Err(e) = r {
        error!("free_swap_chain_processor: get_mut FAILED: {e:?}");
    }
}

unsafe fn do_set_render_adapter(request: WDFREQUEST, input_len: usize) -> NTSTATUS {
    if input_len < size_of::<SetRenderAdapterParams>() {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    }
    let Some(pin) = (unsafe { input_buf(request, size_of::<SetRenderAdapterParams>()) }) else {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    };
    let params = unsafe { &*pin.cast::<SetRenderAdapterParams>() };
    PREFERRED_RENDER_ADAPTER.store(
        ((params.luid_high as u32 as u64) << 32) | u64::from(params.luid_low),
        Ordering::Relaxed,
    );
    if let Some(adapter) = ADAPTER.get() {
        let in_args = IDARG_IN_ADAPTERSETRENDERADAPTER {
            PreferredRenderAdapter: LUID {
                LowPart: params.luid_low,
                HighPart: params.luid_high,
            },
        };
        if let Err(e) = unsafe { IddCxAdapterSetRenderAdapter(adapter.0.as_ptr(), &in_args) } {
            error!("SET_RENDER_ADAPTER failed: {e:?}");
        }
    }
    NTSTATUS::STATUS_SUCCESS
}

unsafe fn do_get_watchdog(request: WDFREQUEST, output_len: usize, bytes: &mut usize) -> NTSTATUS {
    if output_len < size_of::<WatchdogOut>() {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    }
    let Some(pout) = (unsafe { output_buf(request, size_of::<WatchdogOut>()) }) else {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    };
    let out = WatchdogOut {
        timeout: WATCHDOG_TIMEOUT.load(Ordering::Relaxed),
        countdown: WATCHDOG_COUNTDOWN.load(Ordering::Relaxed),
    };
    unsafe { pout.cast::<WatchdogOut>().write_unaligned(out) };
    *bytes = size_of::<WatchdogOut>();
    NTSTATUS::STATUS_SUCCESS
}

unsafe fn do_get_version(request: WDFREQUEST, output_len: usize, bytes: &mut usize) -> NTSTATUS {
    if output_len < PROTOCOL_VERSION.len() {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    }
    let Some(pout) = (unsafe { output_buf(request, PROTOCOL_VERSION.len()) }) else {
        return NTSTATUS::STATUS_BUFFER_TOO_SMALL;
    };
    unsafe { std::ptr::copy_nonoverlapping(PROTOCOL_VERSION.as_ptr(), pout, PROTOCOL_VERSION.len()) };
    *bytes = PROTOCOL_VERSION.len();
    NTSTATUS::STATUS_SUCCESS
}

/// Tear down monitors. `force` (CLEAR_ALL) reaps EVERYTHING — orphans from a crashed previous host;
/// the watchdog passes `false`, which spares any monitor still inside its creation grace
/// (`MONITOR_GRACE`) so a freshly-born monitor is never reaped mid-setup. Caller MUST hold
/// `MONITOR_OP_LOCK` (lock order: OP_LOCK → MONITOR_MODES). Mirrors SudoVDA's DisconnectAllMonitors.
fn disconnect_all_monitors_locked(force: bool) {
    // Drain under the lock (fast); free processors + depart OUTSIDE it (the processor-join blocks).
    let to_depart: Vec<MonitorObject> = {
        let mut lock = MONITOR_MODES.lock().unwrap();
        if lock.is_empty() {
            return;
        }
        let mut keep: Vec<MonitorObject> = Vec::new();
        let mut depart: Vec<MonitorObject> = Vec::new();
        for mon in lock.drain(..) {
            if !force && mon.created_at.elapsed() < MONITOR_GRACE {
                keep.push(mon); // still in its host-side setup window — leave it alone
            } else {
                depart.push(mon);
            }
        }
        *lock = keep;
        depart
    };
    for mon in to_depart {
        if let Some(obj) = mon.object {
            free_swap_chain_processor(obj.as_ptr());
            // SAFETY: `obj` is a live IddCx monitor object.
            if let Err(e) = unsafe { IddCxMonitorDeparture(obj.as_ptr()) } {
                error!("teardown: monitor departure failed: {e:?}");
            }
        }
    }
}

/// Public entry: takes `MONITOR_OP_LOCK`, then tears down. Used by CLEAR_ALL (`force = true`).
fn disconnect_all_monitors(force: bool) {
    let _op = MONITOR_OP_LOCK.lock().unwrap();
    disconnect_all_monitors_locked(force);
}

/// Start the watchdog thread (once). The host reads the timeout via GET_WATCHDOG and PINGs every
/// timeout/3; if it stops, the countdown reaches 0 and every monitor is torn down — so a crashed/gone
/// host never leaves a phantom display. Mirrors SudoVDA's RunWatchdog.
pub fn start_watchdog() {
    static STARTED: AtomicBool = AtomicBool::new(false);
    if STARTED.swap(true, Ordering::Relaxed) {
        return;
    }
    let timeout = WATCHDOG_TIMEOUT.load(Ordering::Relaxed);
    if timeout == 0 {
        return;
    }
    WATCHDOG_COUNTDOWN.store(timeout, Ordering::Relaxed);
    thread::spawn(|| loop {
        thread::sleep(Duration::from_secs(1));
        // Nothing to guard while there are no monitors.
        if MONITOR_MODES.lock().unwrap().is_empty() {
            continue;
        }
        let prev = WATCHDOG_COUNTDOWN.load(Ordering::Relaxed);
        if prev == 0 {
            continue;
        }
        // Decrement without clobbering a concurrent IOCTL reset (CAS).
        if WATCHDOG_COUNTDOWN
            .compare_exchange(prev, prev - 1, Ordering::Relaxed, Ordering::Relaxed)
            .is_ok()
            && prev - 1 == 0
        {
            // About to fire. Serialize against do_add/do_remove (so we never tear an entry out from
            // under an in-flight ADD), then RE-CHECK the countdown under the lock: if a concurrent
            // IOCTL (PING/ADD) reset it while we were acquiring the lock, the host is alive — abort.
            let _op = MONITOR_OP_LOCK.lock().unwrap();
            if WATCHDOG_COUNTDOWN.load(Ordering::Relaxed) == 0 {
                error!("watchdog expired (host stopped pinging) — tearing down stale monitors");
                disconnect_all_monitors_locked(false);
            }
        }
    });
}