liblitho 0.2.0

cli tool to flash/clone the images to storage devices
Documentation
use crate::platform::traits::{DeviceReader, DeviceWriter};
use anyhow::{Context, Result};
use log::{debug, info};
use std::ffi::OsStr;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::os::windows::ffi::OsStrExt;
use std::os::windows::io::{AsRawHandle, FromRawHandle};
use std::ptr;

use winapi::shared::minwindef::DWORD;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::fileapi::{CreateFileW, WriteFile, OPEN_EXISTING};
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::winioctl::{FSCTL_ALLOW_EXTENDED_DASD_IO, IOCTL_DISK_IS_WRITABLE};
use winapi::um::winnt::{
    FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE,
};

fn normalize_physical_drive_path(device_path: &str) -> String {
    crate::platform::windows::devices::canonical_physical_drive_path(device_path).unwrap_or_else(
        |_| {
            let trimmed = device_path.trim();
            if trimmed.starts_with(r"\\.\") {
                trimmed.to_string()
            } else {
                format!(r"\\.\{trimmed}")
            }
        },
    )
}

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

/// Cap each `WriteFile` syscall — large single writes can fail on some USB stacks.
const MAX_WRITE_CHUNK: usize = 1_048_576;

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 create_physical_drive_file(path: &str, write: bool) -> Result<File> {
    if write {
        let _ = crate::platform::windows::access::enable_raw_disk_privileges();
    }

    let access = if write {
        GENERIC_READ | GENERIC_WRITE
    } else {
        GENERIC_READ
    };

    let wide = wide_path(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(),
            access,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            ptr::null_mut(),
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            ptr::null_mut(),
        )
    };

    if handle == INVALID_HANDLE_VALUE {
        // SAFETY: GetLastError is a pure TLS read of the last error for this thread.
        let code = unsafe { GetLastError() };
        let mut message = format_windows_io_error(&format!("Failed to open device {path}"));
        if code == 5 {
            message.push_str(
                ". Ensure litho is running elevated (UAC), all volumes on the disk are dismounted, \
                 and no other disk tools have the device open.",
            );
        }
        anyhow::bail!(message);
    }

    // SAFETY: HANDLE is non-INVALID and exclusive ownership transfers into File; no other owner remains.
    Ok(unsafe { File::from_raw_handle(handle as _) })
}

/// Allow raw sector writes to the full physical drive surface (including partition gaps).
fn allow_extended_dasd_io(file: &File) {
    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(
            file.as_raw_handle() as *mut _,
            FSCTL_ALLOW_EXTENDED_DASD_IO,
            ptr::null_mut(),
            0,
            ptr::null_mut(),
            0,
            &mut bytes_returned,
            ptr::null_mut(),
        )
    };
    if ok == 0 {
        debug!(
            "FSCTL_ALLOW_EXTENDED_DASD_IO not enabled: {}",
            format_windows_io_error("ioctl")
        );
    }
}

fn verify_disk_writable(file: &File, path: &str) -> Result<()> {
    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(
            file.as_raw_handle() as *mut _,
            IOCTL_DISK_IS_WRITABLE,
            ptr::null_mut(),
            0,
            ptr::null_mut(),
            0,
            &mut bytes_returned,
            ptr::null_mut(),
        )
    };
    if ok != 0 {
        return Ok(());
    }

    // SAFETY: GetLastError is a pure TLS read of the last error for this thread.
    let code = unsafe { GetLastError() };
    let mut message = format_windows_io_error(&format!("Disk {path} is not writable"));
    if code == 5 {
        message.push_str(
            ". Volumes on this disk may still be mounted — close File Explorer, confirm the \
             dismount prompt, and run litho-tui elevated (UAC).",
        );
    } else if code == 19 {
        message.push_str(". The device may be hardware write-protected.");
    }
    anyhow::bail!(message);
}

/// Buffered device reader for clone and post-write verification.
pub struct WindowsBufferedDeviceReader {
    file: File,
}

impl DeviceReader for WindowsBufferedDeviceReader {
    fn open(device_path: &str) -> Result<Self> {
        let windows_device = normalize_physical_drive_path(device_path);
        debug!("Opening Windows device for buffered read: {windows_device}");
        crate::platform::windows::access::prepare_physical_drive_for_read(&windows_device)
            .map_err(|e| anyhow::anyhow!(e))?;
        let file = create_physical_drive_file(&windows_device, false)?;
        Ok(Self { file })
    }

    fn device_size(&self) -> Result<u64> {
        device_size_from_file(&self.file)
    }
}

impl Read for WindowsBufferedDeviceReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.file.read(buf)
    }
}

pub struct WindowsDeviceWriter {
    file: File,
    /// Exclusive volume locks; type is [`super::WindowsVolumeGuard`].
    _volume_guard: super::WindowsVolumeGuard,
}

impl DeviceWriter for WindowsDeviceWriter {
    fn open(device_path: &str) -> Result<Self> {
        use crate::platform::traits::VolumeOps;
        use crate::platform::windows::WindowsPlatform;

        let windows_device = normalize_physical_drive_path(device_path);
        debug!("Opening Windows device for writing: {windows_device}");

        // Type-enforced exclusive access: Guard must outlive the write handle.
        let volume_guard = WindowsPlatform::prepare_for_io(&windows_device, true)
            .map_err(|e| anyhow::anyhow!(e))?;

        let file = create_physical_drive_file(&windows_device, true)?;
        allow_extended_dasd_io(&file);
        verify_disk_writable(&file, &windows_device)?;

        info!("Opened {windows_device} for writing");
        Ok(Self {
            file,
            _volume_guard: volume_guard,
        })
    }

    fn write_at(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
        self.file
            .seek(SeekFrom::Start(offset))
            .context(format!("Failed to seek device to offset {offset}"))?;
        self.write_all(buf)
            .context(format!("Failed to write {offset} bytes at offset {offset}"))?;
        Ok(())
    }

    fn flush_and_sync(&mut self) -> Result<()> {
        self.file.flush().context("Failed to flush device")?;
        self.file.sync_all().context("Failed to sync device")?;
        debug!("Device flushed and synced");
        Ok(())
    }

    fn device_size(&self) -> Result<u64> {
        device_size_from_file(&self.file)
    }

    fn supports_inline_verify(&self) -> bool {
        true
    }

    fn rewind_for_verify(&mut self) -> Result<()> {
        self.file
            .seek(SeekFrom::Start(0))
            .context("Failed to rewind device for verification")?;
        Ok(())
    }

    fn read_for_verify(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.file.read(buf)
    }
}

impl Write for WindowsDeviceWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }

        let mut total_written = 0usize;
        for chunk in buf.chunks(MAX_WRITE_CHUNK) {
            let mut bytes_written: DWORD = 0;
            // SAFETY: handle is a valid open device handle; buffers are sized for the ioctl/write contract.
            let ok = unsafe {
                WriteFile(
                    self.file.as_raw_handle() as *mut _,
                    chunk.as_ptr() as *const _,
                    chunk.len() as DWORD,
                    &mut bytes_written,
                    ptr::null_mut(),
                )
            };

            if ok == 0 {
                return Err(std::io::Error::last_os_error());
            }

            total_written += bytes_written as usize;
        }

        Ok(total_written)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.file.flush()
    }
}

fn device_size_from_file(file: &File) -> Result<u64> {
    let metadata = file.metadata().context("Failed to get device metadata")?;
    Ok(metadata.len())
}