liblitho 0.2.0

//! Windows privileges and volume dismount helpers for raw physical-drive I/O.

use log::{debug, info, warn};
use std::ffi::OsStr;
use std::os::windows::ffi::OsStrExt;
use std::ptr;

use winapi::shared::minwindef::DWORD;
use winapi::shared::winerror::{ERROR_FILE_NOT_FOUND, ERROR_NOT_READY, ERROR_PATH_NOT_FOUND};
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::fileapi::{
    CreateFileW, FindFirstVolumeW, FindNextVolumeW, FindVolumeClose, OPEN_EXISTING,
};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::processthreadsapi::{GetCurrentProcess, OpenProcessToken};
use winapi::um::securitybaseapi::AdjustTokenPrivileges;
use winapi::um::winbase::LookupPrivilegeValueW;
use winapi::um::winioctl::{
    DISK_EXTENT, FSCTL_DISMOUNT_VOLUME, FSCTL_LOCK_VOLUME, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
    IOCTL_VOLUME_OFFLINE,
};
use winapi::um::winnt::HANDLE;
use winapi::um::winnt::{
    FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE,
    LUID_AND_ATTRIBUTES, SE_BACKUP_NAME, SE_MANAGE_VOLUME_NAME, SE_PRIVILEGE_ENABLED,
    SE_RESTORE_NAME, TOKEN_ADJUST_PRIVILEGES, TOKEN_QUERY,
};

const VOLUME_NAME_BUFFER_CHARS: usize = 256;
const MAX_VOLUME_DISK_EXTENTS: usize = 16;
const PRIVILEGE_BUFFER_SIZE: usize = 3;

fn wide_path(path: &str) -> Vec<u16> {
    OsStr::new(path).encode_wide().chain(Some(0)).collect()
}

fn wide_const(value: &str) -> Vec<u16> {
    wide_path(value)
}

fn format_windows_io_error(context: &str) -> String {
    let err = std::io::Error::last_os_error();
    let code = err.raw_os_error().unwrap_or(0);
    format!("{context}: {err} (Windows error {code:#x})")
}

fn last_windows_error_code() -> DWORD {
    // SAFETY: GetLastError is a pure TLS read of the last error for this thread.
    unsafe { GetLastError() }
}

fn volume_unavailable_error(code: DWORD) -> bool {
    matches!(
        code,
        ERROR_NOT_READY | ERROR_FILE_NOT_FOUND | ERROR_PATH_NOT_FOUND
    )
}

/// Enable privileges required to lock/dismount volumes and access the raw disk.
pub fn enable_raw_disk_privileges() -> Result<(), String> {
    let privilege_names = [SE_MANAGE_VOLUME_NAME, SE_BACKUP_NAME, SE_RESTORE_NAME];
    let mut token = ptr::null_mut();
    // SAFETY: process handle is current process; token handle is closed on all paths; privilege buffers are correctly sized.
    let token_ok = unsafe {
        OpenProcessToken(
            GetCurrentProcess(),
            TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
            &mut token,
        )
    };
    if token_ok == 0 {
        return Err(format_windows_io_error("Failed to open process token"));
    }

    let mut luid_and_attributes = [LUID_AND_ATTRIBUTES {
        // SAFETY: Zero-initialized POD/C struct for WinAPI out-parameter; all fields overwritten before use or treated as opaque buffer.
        Luid: unsafe { std::mem::zeroed() },
        Attributes: SE_PRIVILEGE_ENABLED,
    }; PRIVILEGE_BUFFER_SIZE];

    let mut enabled: usize = 0;
    for name in privilege_names {
        // SAFETY: Zero-initialized POD/C struct for WinAPI out-parameter; all fields overwritten before use or treated as opaque buffer.
        let mut luid = unsafe { std::mem::zeroed() };
        let lookup_ok =
            // SAFETY: process handle is current process; token handle is closed on all paths; privilege buffers are correctly sized.
            unsafe { LookupPrivilegeValueW(ptr::null(), wide_const(name).as_ptr(), &mut luid) };
        if lookup_ok == 0 {
            warn!(
                "Could not look up privilege {name}: {}",
                format_windows_io_error("lookup")
            );
            continue;
        }
        luid_and_attributes[enabled].Luid = luid;
        luid_and_attributes[enabled].Attributes = SE_PRIVILEGE_ENABLED;
        enabled += 1;
    }

    if enabled == 0 {
        // SAFETY: handle was obtained from a successful WinAPI open and is not used after close; ownership transferred to this call.
        unsafe {
            CloseHandle(token);
        }
        return Err("Failed to resolve any raw-disk privileges".into());
    }

    #[repr(C)]
    struct TokenPrivilegesBuf {
        privilege_count: DWORD,
        privileges: [LUID_AND_ATTRIBUTES; PRIVILEGE_BUFFER_SIZE],
    }

    let mut privileges = TokenPrivilegesBuf {
        privilege_count: enabled as DWORD,
        privileges: luid_and_attributes,
    };

    // SAFETY: process handle is current process; token handle is closed on all paths; privilege buffers are correctly sized.
    let adjust_ok = unsafe {
        AdjustTokenPrivileges(
            token,
            0,
            &mut privileges as *mut _ as *mut _,
            0,
            ptr::null_mut(),
            ptr::null_mut(),
        )
    };
    // SAFETY: handle was obtained from a successful WinAPI open and is not used after close; ownership transferred to this call.
    unsafe {
        CloseHandle(token);
    }

    if adjust_ok == 0 {
        return Err(format_windows_io_error(
            "Failed to enable raw-disk privileges",
        ));
    }

    let code = last_windows_error_code();
    if code != 0 {
        // ERROR_NOT_ALL_ASSIGNED (1300) is acceptable if some privileges were enabled.
        debug!("AdjustTokenPrivileges completed with code {code:#x}");
    }

    debug!("Raw-disk privileges enabled");
    Ok(())
}

fn volume_name_to_device_path(name: &str) -> String {
    let trimmed = name.trim_end_matches('\0').trim();
    if let Some(rest) = trimmed.strip_prefix(r"\\?\") {
        format!(r"\\.\{rest}")
    } else if trimmed.starts_with(r"\\.\") {
        trimmed.to_string()
    } else {
        format!(r"\\.\{trimmed}")
    }
}

#[repr(C)]
struct VolumeDiskExtentsBuf {
    number_of_disk_extents: DWORD,
    extents: [DISK_EXTENT; MAX_VOLUME_DISK_EXTENTS],
}

fn volume_disk_number(volume_device_path: &str) -> Option<u32> {
    let wide = wide_path(volume_device_path);
    // SAFETY: path is a NUL-terminated wide string; on success ownership of the HANDLE is transferred to Rust File/LockedVolume.
    let handle = unsafe {
        CreateFileW(
            wide.as_ptr(),
            GENERIC_READ,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            ptr::null_mut(),
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            ptr::null_mut(),
        )
    };
    if handle == INVALID_HANDLE_VALUE {
        debug!(
            "Could not open {volume_device_path} for disk-extent query: {}",
            format_windows_io_error("open")
        );
        return None;
    }

    let mut info = VolumeDiskExtentsBuf {
        number_of_disk_extents: 0,
        // SAFETY: Zero-initialized POD/C struct for WinAPI out-parameter; all fields overwritten before use or treated as opaque buffer.
        extents: unsafe { std::mem::zeroed() },
    };
    let mut bytes_returned: DWORD = 0;
    // SAFETY: handle is a valid open device handle; buffers are sized for the ioctl/write contract.
    let ok = unsafe {
        DeviceIoControl(
            handle,
            IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
            ptr::null_mut(),
            0,
            &mut info as *mut _ as *mut _,
            std::mem::size_of::<VolumeDiskExtentsBuf>() as DWORD,
            &mut bytes_returned,
            ptr::null_mut(),
        )
    };
    // SAFETY: handle was obtained from a successful WinAPI open and is not used after close; ownership transferred to this call.
    unsafe {
        CloseHandle(handle);
    }

    if ok == 0 || info.number_of_disk_extents == 0 {
        return None;
    }

    Some(info.extents[0].DiskNumber)
}

struct LockedVolume {
    label: String,
    handle: HANDLE,
}

// SAFETY: `handle` is an exclusively owned Windows HANDLE (kernel object reference).
// Ownership may move across threads; we do not implement Sync because concurrent
// use of the same handle is not required and is not synchronized here.
unsafe impl Send for LockedVolume {}

impl Drop for LockedVolume {
    fn drop(&mut self) {
        // SAFETY: handle was obtained from a successful WinAPI open and is not used after close; ownership transferred to this call.
        unsafe {
            CloseHandle(self.handle);
        }
    }
}

/// Keeps volume handles locked for the lifetime of a physical-drive write session.
///
/// Closing the lock immediately after dismount allows Windows to re-mount the volume
/// when the partition table is rewritten mid-flash.
pub struct PhysicalDriveIoSession {
    _locks: Vec<LockedVolume>,
}

fn lock_and_dismount_volume(
    volume_device_path: &str,
    label: &str,
) -> Result<Option<LockedVolume>, String> {
    let wide = wide_path(volume_device_path);
    // SAFETY: path is a NUL-terminated wide string; on success ownership of the HANDLE is transferred to Rust File/LockedVolume.
    let handle = unsafe {
        CreateFileW(
            wide.as_ptr(),
            GENERIC_READ | GENERIC_WRITE,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            ptr::null_mut(),
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            ptr::null_mut(),
        )
    };

    if handle == INVALID_HANDLE_VALUE {
        let code = last_windows_error_code();
        if volume_unavailable_error(code) {
            debug!("Volume {label} is not open (error {code:#x}); skip lock");
            return Ok(None);
        }
        return Err(format!(
            "{}. Close File Explorer and any apps using {label}, then retry.",
            format_windows_io_error(&format!(
                "Could not open volume {label} at {volume_device_path}"
            ))
        ));
    }

    let mut bytes_returned: DWORD = 0;
    // SAFETY: handle is a valid open device handle; buffers are sized for the ioctl/write contract.
    let lock_ok = unsafe {
        DeviceIoControl(
            handle,
            FSCTL_LOCK_VOLUME,
            ptr::null_mut(),
            0,
            ptr::null_mut(),
            0,
            &mut bytes_returned,
            ptr::null_mut(),
        )
    };
    if lock_ok == 0 {
        let err = format_windows_io_error(&format!(
            "Could not lock volume {label} at {volume_device_path}"
        ));
        // SAFETY: handle was obtained from a successful WinAPI open and is not used after close; ownership transferred to this call.
        unsafe {
            CloseHandle(handle);
        }
        return Err(format!(
            "{err}. Close File Explorer and any apps using {label}, then retry."
        ));
    }

    // SAFETY: handle is a valid open device handle; buffers are sized for the ioctl/write contract.
    let dismount_ok = unsafe {
        DeviceIoControl(
            handle,
            FSCTL_DISMOUNT_VOLUME,
            ptr::null_mut(),
            0,
            ptr::null_mut(),
            0,
            &mut bytes_returned,
            ptr::null_mut(),
        )
    };
    if dismount_ok == 0 {
        // SAFETY: handle is a valid open device handle; buffers are sized for the ioctl/write contract.
        let offline_ok = unsafe {
            DeviceIoControl(
                handle,
                IOCTL_VOLUME_OFFLINE,
                ptr::null_mut(),
                0,
                ptr::null_mut(),
                0,
                &mut bytes_returned,
                ptr::null_mut(),
            )
        };
        if offline_ok == 0 {
            let err = format_windows_io_error(&format!(
                "Could not dismount or offline volume {label} at {volume_device_path}"
            ));
            // SAFETY: handle was obtained from a successful WinAPI open and is not used after close; ownership transferred to this call.
            unsafe {
                CloseHandle(handle);
            }
            return Err(format!(
                "{err}. Close File Explorer and any apps using {label}, then retry."
            ));
        }
        debug!("Volume {label} taken offline via IOCTL_VOLUME_OFFLINE");
    }

    info!("Locked and dismounted volume {label} ({volume_device_path})");
    Ok(Some(LockedVolume {
        label: label.to_string(),
        handle,
    }))
}

pub fn dismount_volume_at_path(volume_device_path: &str, label: &str) -> Result<(), String> {
    let _ = lock_and_dismount_volume(volume_device_path, label)?;
    Ok(())
}

fn dismount_logical_volume(drive_letter: &str) -> Result<(), String> {
    let normalized = drive_letter
        .trim()
        .trim_end_matches('\\')
        .trim_end_matches(':')
        .to_ascii_uppercase();
    if normalized.len() != 1 || !normalized.chars().all(|c| c.is_ascii_alphabetic()) {
        return Err(format!("Invalid drive letter: {drive_letter}"));
    }

    let volume_path = format!(r"\\.\{normalized}:");
    dismount_volume_at_path(&volume_path, &format!("{normalized}:"))
}

fn dismount_volumes_enumerated(disk_index: u32) -> (Vec<String>, Vec<String>) {
    let mut dismounted = Vec::new();
    let mut failures = Vec::new();

    let mut name_buf = vec![0u16; VOLUME_NAME_BUFFER_CHARS];
    let find =
        // SAFETY: volume name buffer is sized to VOLUME_NAME_BUFFER_CHARS; find handle is closed with FindVolumeClose.
        unsafe { FindFirstVolumeW(name_buf.as_mut_ptr(), VOLUME_NAME_BUFFER_CHARS as DWORD) };
    if find == INVALID_HANDLE_VALUE {
        debug!(
            "FindFirstVolumeW failed: {}",
            format_windows_io_error("enumerate")
        );
        return (dismounted, failures);
    }

    loop {
        let len = name_buf
            .iter()
            .position(|&c| c == 0)
            .unwrap_or(name_buf.len());
        let volume_name = String::from_utf16_lossy(&name_buf[..len]);
        let device_path = volume_name_to_device_path(&volume_name);
        let label = volume_name.clone();

        if volume_disk_number(&device_path) == Some(disk_index) {
            match dismount_volume_at_path(&device_path, &label) {
                Ok(()) => dismounted.push(label),
                Err(err) => failures.push(err),
            }
        }

        // SAFETY: volume name buffer is sized to VOLUME_NAME_BUFFER_CHARS; find handle is closed with FindVolumeClose.
        let has_next = unsafe {
            FindNextVolumeW(
                find,
                name_buf.as_mut_ptr(),
                VOLUME_NAME_BUFFER_CHARS as DWORD,
            )
        };
        if has_next == 0 {
            break;
        }
    }

    // SAFETY: volume name buffer is sized to VOLUME_NAME_BUFFER_CHARS; find handle is closed with FindVolumeClose.
    unsafe {
        FindVolumeClose(find);
    }

    (dismounted, failures)
}

fn dismount_known_targets(
    targets: Vec<(String, String)>,
    dismounted: &mut Vec<String>,
    failures: &mut Vec<String>,
) {
    for (device_path, label) in targets {
        if dismounted.iter().any(|entry| entry == &label) {
            continue;
        }
        match dismount_volume_at_path(&device_path, &label) {
            Ok(()) => dismounted.push(label),
            Err(err) => failures.push(err),
        }
    }
}

fn lock_known_targets(
    targets: Vec<(String, String)>,
    locked: &mut Vec<LockedVolume>,
    failures: &mut Vec<String>,
) {
    for (device_path, label) in targets {
        if locked.iter().any(|entry| entry.label == label) {
            continue;
        }
        match lock_and_dismount_volume(&device_path, &label) {
            Ok(Some(volume)) => locked.push(volume),
            Ok(None) => {}
            Err(err) => failures.push(err),
        }
    }
}

fn lock_volumes_enumerated(disk_index: u32) -> (Vec<LockedVolume>, Vec<String>) {
    let mut locked = Vec::new();
    let mut failures = Vec::new();

    let mut name_buf = vec![0u16; VOLUME_NAME_BUFFER_CHARS];
    let find =
        // SAFETY: volume name buffer is sized to VOLUME_NAME_BUFFER_CHARS; find handle is closed with FindVolumeClose.
        unsafe { FindFirstVolumeW(name_buf.as_mut_ptr(), VOLUME_NAME_BUFFER_CHARS as DWORD) };
    if find == INVALID_HANDLE_VALUE {
        debug!(
            "FindFirstVolumeW failed: {}",
            format_windows_io_error("enumerate")
        );
        return (locked, failures);
    }

    loop {
        let len = name_buf
            .iter()
            .position(|&c| c == 0)
            .unwrap_or(name_buf.len());
        let volume_name = String::from_utf16_lossy(&name_buf[..len]);
        let device_path = volume_name_to_device_path(&volume_name);
        let label = volume_name.clone();

        if volume_disk_number(&device_path) == Some(disk_index) {
            if locked.iter().any(|entry| entry.label == label) {
                continue;
            }
            match lock_and_dismount_volume(&device_path, &label) {
                Ok(Some(volume)) => locked.push(volume),
                Ok(None) => {}
                Err(err) => failures.push(err),
            }
        }

        // SAFETY: volume name buffer is sized to VOLUME_NAME_BUFFER_CHARS; find handle is closed with FindVolumeClose.
        let has_next = unsafe {
            FindNextVolumeW(
                find,
                name_buf.as_mut_ptr(),
                VOLUME_NAME_BUFFER_CHARS as DWORD,
            )
        };
        if has_next == 0 {
            break;
        }
    }

    // SAFETY: volume name buffer is sized to VOLUME_NAME_BUFFER_CHARS; find handle is closed with FindVolumeClose.
    unsafe {
        FindVolumeClose(find);
    }

    (locked, failures)
}

fn prepare_volumes_inner(
    physical_path: &str,
    keep_locks: bool,
) -> Result<(Vec<String>, Vec<LockedVolume>), String> {
    if let Err(err) = enable_raw_disk_privileges() {
        warn!("Could not enable all raw-disk privileges: {err}");
    }

    let disk_index = crate::platform::windows::devices::parse_physical_drive_index(physical_path)
        .ok_or_else(|| format!("Invalid physical drive path: {physical_path}"))?;

    info!("Preparing {physical_path} (disk {disk_index}) for raw I/O");

    let mut dismounted = Vec::new();
    let mut locked = Vec::new();
    let mut failures = Vec::new();

    match crate::platform::windows::devices::volume_dismount_targets_for_disk_index(disk_index) {
        Ok(targets) => {
            if targets.is_empty() {
                debug!("No WMI volume targets found for disk {disk_index}");
            } else {
                info!(
                    "Dismounting {} volume(s) on disk {disk_index}",
                    targets.len()
                );
            }
            if keep_locks {
                lock_known_targets(targets, &mut locked, &mut failures);
                dismounted.extend(locked.iter().map(|entry| entry.label.clone()));
            } else {
                dismount_known_targets(targets, &mut dismounted, &mut failures);
            }
        }
        Err(err) => {
            warn!("Volume target lookup failed for disk {disk_index}: {err}");
        }
    }

    if keep_locks {
        let (enum_locked, enum_failures) = lock_volumes_enumerated(disk_index);
        for volume in enum_locked {
            if !locked.iter().any(|entry| entry.label == volume.label) {
                dismounted.push(volume.label.clone());
                locked.push(volume);
            }
        }
        failures.extend(enum_failures);
    } else {
        let (enum_dismounted, enum_failures) = dismount_volumes_enumerated(disk_index);
        for label in enum_dismounted {
            if !dismounted.iter().any(|entry| entry == &label) {
                dismounted.push(label);
            }
        }
        failures.extend(enum_failures);
    }

    if let Ok(letters) =
        crate::platform::windows::devices::mounted_drive_letters_for_disk_index(disk_index)
    {
        for letter in letters {
            let normalized = format!("{letter}:");
            if dismounted.iter().any(|entry| entry.contains(&letter)) {
                continue;
            }
            if keep_locks {
                let device_path = format!(
                    r"\\.\{}:",
                    letter.trim().trim_end_matches(':').to_ascii_uppercase()
                );
                match lock_and_dismount_volume(&device_path, &normalized) {
                    Ok(Some(volume)) => {
                        dismounted.push(normalized);
                        locked.push(volume);
                    }
                    Ok(None) => {}
                    Err(err) => failures.push(err),
                }
            } else {
                match dismount_logical_volume(&letter) {
                    Ok(()) => dismounted.push(normalized),
                    Err(err) => failures.push(err),
                }
            }
        }
    }

    if failures.is_empty() {
        if !dismounted.is_empty() {
            info!(
                "Dismounted volumes on {physical_path}: {}",
                dismounted.join(", ")
            );
        } else {
            debug!("No mounted volumes required dismount on {physical_path}");
        }
        return Ok((dismounted, locked));
    }

    let mut message = format!("Could not dismount all volumes on {physical_path}");
    if !dismounted.is_empty() {
        message.push_str(&format!(
            ". Dismounted {} but still failed: {}",
            dismounted.join(", "),
            failures.join("; ")
        ));
    } else {
        message.push_str(&format!(". Failures: {}", failures.join("; ")));
    }
    message.push_str(
        ". Close File Explorer and any apps using the disk, then retry from an elevated session.",
    );
    Err(message)
}

/// Lock and dismount volumes on a physical drive so raw I/O can proceed.
pub fn prepare_physical_drive_for_io(physical_path: &str) -> Result<(), String> {
    prepare_volumes_inner(physical_path, false).map(|_| ())
}

/// Prepare for read-only access. Dismount is best-effort — reads can proceed if it fails.
pub fn prepare_physical_drive_for_read(physical_path: &str) -> Result<(), String> {
    if let Err(err) = enable_raw_disk_privileges() {
        warn!("Could not enable all raw-disk privileges: {err}");
    }

    match prepare_volumes_inner(physical_path, false) {
        Ok(_) => Ok(()),
        Err(err) => {
            warn!("Could not dismount volumes before read on {physical_path} (continuing): {err}");
            Ok(())
        }
    }
}

/// Like [`prepare_physical_drive_for_io`], but keeps volume lock handles open until dropped.
pub fn prepare_physical_drive_for_write(
    physical_path: &str,
) -> Result<PhysicalDriveIoSession, String> {
    let (_, locks) = prepare_volumes_inner(physical_path, true)?;
    Ok(PhysicalDriveIoSession { _locks: locks })
}