use crate::devices::DeviceInfo;
use anyhow::Result;
use log::debug;
use std::collections::HashMap;
use std::fs::OpenOptions;
use std::mem;
use std::sync::Mutex;
use std::thread;
use wmi::{COMLibrary, Variant, WMIConnection};
static WMI_MUTEX: Mutex<()> = Mutex::new(());
#[cfg(target_os = "windows")]
use std::os::windows::io::AsRawHandle;
#[cfg(target_os = "windows")]
use winapi::shared::minwindef::DWORD;
#[cfg(target_os = "windows")]
use winapi::um::ioapiset::DeviceIoControl;
#[cfg(target_os = "windows")]
use winapi::um::winioctl::IOCTL_DISK_GET_LENGTH_INFO;
#[derive(Debug, Clone)]
struct RawDiskDrive {
device_id: String,
model: String,
manufacturer: String,
size: Option<u64>,
media_type: Option<String>,
interface_type: Option<String>,
index: Option<u32>,
}
#[derive(Debug, Clone)]
struct RawPartition {
device_id: String,
disk_index: Option<u32>,
size: Option<u64>,
}
fn run_wmi_thread<T: Send + 'static>(
label: &str,
f: impl FnOnce() -> Result<T, String> + Send + 'static,
) -> Result<T, String> {
let thread_label = label.to_string();
let lock_label = thread_label.clone();
let spawn_label = thread_label.clone();
let panic_label = thread_label;
thread::Builder::new()
.name("litho-wmi".into())
.spawn(move || {
let _guard = WMI_MUTEX
.lock()
.map_err(|_| format!("WMI lock poisoned during {lock_label}"))?;
f()
})
.map_err(|e| format!("Failed to spawn WMI thread for {spawn_label}: {e}"))?
.join()
.map_err(|_| format!("WMI thread for {panic_label} panicked"))?
}
fn is_removable(media_type: &Option<String>, interface_type: &Option<String>) -> u8 {
if interface_type
.as_ref()
.is_some_and(|iface| iface.eq_ignore_ascii_case("USB"))
{
return 1;
}
media_type
.as_ref()
.map(|media| {
let lower = media.to_lowercase();
lower.contains("removable media") || lower.contains("external hard disk media")
})
.map(u8::from)
.unwrap_or(0)
}
fn variant_to_u64(value: &Variant) -> Option<u64> {
match value {
Variant::UI8(n) => Some(*n),
Variant::UI4(n) => Some(u64::from(*n)),
Variant::UI2(n) => Some(u64::from(*n)),
Variant::UI1(n) => Some(u64::from(*n)),
Variant::I8(n) if *n >= 0 => Some(*n as u64),
Variant::I4(n) if *n >= 0 => Some(*n as u64),
Variant::String(s) => s.trim().parse().ok(),
Variant::Null | Variant::Empty => None,
_ => None,
}
}
fn variant_to_string(value: &Variant) -> Option<String> {
match value {
Variant::String(s) => {
let trimmed = s.trim();
if trimmed.is_empty() {
None
} else {
Some(trimmed.to_string())
}
}
_ => None,
}
}
fn variant_to_u32(value: &Variant) -> Option<u32> {
variant_to_u64(value).and_then(|n| u32::try_from(n).ok())
}
fn extract_quoted_value(path: &str) -> Option<String> {
let start = path.find('"')? + 1;
let end = path.rfind('"')?;
if end <= start {
return None;
}
let value = path[start..end].trim().to_string();
if value.is_empty() {
None
} else {
Some(value)
}
}
fn parse_disk_drive_row(row: HashMap<String, Variant>) -> Option<RawDiskDrive> {
Some(RawDiskDrive {
device_id: variant_to_string(row.get("DeviceID")?)?,
model: row
.get("Model")
.and_then(variant_to_string)
.unwrap_or_default(),
manufacturer: row
.get("Manufacturer")
.and_then(variant_to_string)
.unwrap_or_default(),
size: row.get("Size").and_then(variant_to_u64),
media_type: row.get("MediaType").and_then(variant_to_string),
interface_type: row.get("InterfaceType").and_then(variant_to_string),
index: row.get("Index").and_then(variant_to_u32),
})
}
fn parse_partition_row(row: HashMap<String, Variant>) -> Option<RawPartition> {
Some(RawPartition {
device_id: variant_to_string(row.get("DeviceID")?)?,
disk_index: row.get("DiskIndex").and_then(variant_to_u32),
size: row.get("Size").and_then(variant_to_u64),
})
}
fn query_disk_drives(wmi: &WMIConnection) -> Result<Vec<RawDiskDrive>, String> {
let rows: Vec<HashMap<String, Variant>> = wmi
.raw_query(
"SELECT DeviceID, Model, Manufacturer, Size, MediaType, InterfaceType, Index \
FROM Win32_DiskDrive",
)
.map_err(|e| format!("Failed to query Win32_DiskDrive: {e}"))?;
Ok(rows.into_iter().filter_map(parse_disk_drive_row).collect())
}
fn query_partitions(wmi: &WMIConnection) -> Vec<RawPartition> {
let rows: Vec<HashMap<String, Variant>> =
match wmi.raw_query("SELECT DeviceID, DiskIndex, Size FROM Win32_DiskPartition") {
Ok(rows) => rows,
Err(e) => {
debug!("Win32_DiskPartition query unavailable: {e}");
return Vec::new();
}
};
rows.into_iter().filter_map(parse_partition_row).collect()
}
fn build_partition_to_drive_map(wmi: &WMIConnection) -> HashMap<String, Vec<String>> {
let rows: Vec<HashMap<String, Variant>> =
match wmi.raw_query("SELECT Antecedent, Dependent FROM Win32_LogicalDiskToPartition") {
Ok(rows) => rows,
Err(e) => {
debug!("Win32_LogicalDiskToPartition query unavailable: {e}");
return HashMap::new();
}
};
let mut map: HashMap<String, Vec<String>> = HashMap::new();
for row in rows {
let Some(antecedent) = row.get("Antecedent").and_then(variant_to_string) else {
continue;
};
let Some(dependent) = row.get("Dependent").and_then(variant_to_string) else {
continue;
};
let Some(partition) = extract_quoted_value(&antecedent) else {
continue;
};
let Some(drive_letter) = extract_quoted_value(&dependent)
.map(|letter| letter.trim_end_matches('\\').to_string())
else {
continue;
};
map.entry(partition).or_default().push(drive_letter);
}
for letters in map.values_mut() {
letters.sort();
letters.dedup();
}
map
}
fn build_disk_index_to_drive_letters(
partitions: &[RawPartition],
partition_to_drive: &HashMap<String, Vec<String>>,
) -> HashMap<u32, Vec<String>> {
let mut map: HashMap<u32, Vec<String>> = HashMap::new();
for partition in partitions {
let Some(disk_index) = partition.disk_index else {
continue;
};
let Some(mut letters) = partition_to_drive.get(&partition.device_id).cloned() else {
continue;
};
map.entry(disk_index).or_default().append(&mut letters);
}
for letters in map.values_mut() {
letters.sort();
letters.dedup();
}
map
}
fn bytes_to_sectors(size_bytes: u64) -> u64 {
size_bytes / 512
}
#[repr(C)]
struct GetLengthInformation {
length: i64,
}
pub fn physical_drive_size_bytes(device_path: &str) -> Option<u64> {
let path = canonical_physical_drive_path(device_path).ok()?;
let file = OpenOptions::new().read(true).open(&path).ok()?;
let mut info = GetLengthInformation { length: 0 };
let mut bytes_returned: DWORD = 0;
let ok = unsafe {
DeviceIoControl(
file.as_raw_handle() as *mut _,
IOCTL_DISK_GET_LENGTH_INFO,
std::ptr::null_mut(),
0,
&mut info as *mut _ as *mut _,
mem::size_of::<GetLengthInformation>() as DWORD,
&mut bytes_returned,
std::ptr::null_mut(),
)
};
if ok == 0 {
debug!("IOCTL_DISK_GET_LENGTH_INFO failed for {path}");
return None;
}
let bytes = info.length as u64;
if bytes == 0 {
None
} else {
Some(bytes)
}
}
fn load_msft_physical_disk_sizes() -> HashMap<u32, u64> {
let com = match COMLibrary::new() {
Ok(com) => com,
Err(e) => {
debug!("Failed to initialize COM for Storage WMI: {e}");
return HashMap::new();
}
};
let storage_wmi =
match WMIConnection::with_namespace_path("ROOT\\Microsoft\\Windows\\Storage", com) {
Ok(wmi) => wmi,
Err(e) => {
debug!("Storage WMI namespace unavailable: {e}");
return HashMap::new();
}
};
let rows: Vec<HashMap<String, Variant>> =
match storage_wmi.raw_query("SELECT DeviceId, Size FROM MSFT_PhysicalDisk") {
Ok(rows) => rows,
Err(e) => {
debug!("MSFT_PhysicalDisk query unavailable: {e}");
return HashMap::new();
}
};
let mut map = HashMap::new();
for row in rows {
let Some(device_id) = row.get("DeviceId").and_then(variant_to_string) else {
continue;
};
let Ok(index) = device_id.trim().parse::<u32>() else {
continue;
};
let Some(bytes) = row.get("Size").and_then(variant_to_u64) else {
continue;
};
if bytes > 0 {
map.insert(index, bytes);
}
}
map
}
fn resolve_disk_size_bytes(
raw: &RawDiskDrive,
partition_bytes: Option<u64>,
msft_physical_bytes: Option<u64>,
) -> u64 {
if let Some(size) = raw.size.filter(|size| *size > 0) {
return size;
}
if let Some(bytes) = msft_physical_bytes.filter(|bytes| *bytes > 0) {
return bytes;
}
if let Some(bytes) = partition_bytes.filter(|bytes| *bytes > 0) {
return bytes;
}
physical_drive_size_bytes(&raw.device_id).unwrap_or(0)
}
fn build_disk_index_to_partition_bytes(partitions: &[RawPartition]) -> HashMap<u32, u64> {
let mut map: HashMap<u32, u64> = HashMap::new();
for partition in partitions {
let Some(disk_index) = partition.disk_index else {
continue;
};
let Some(bytes) = partition.size.filter(|bytes| *bytes > 0) else {
continue;
};
*map.entry(disk_index).or_insert(0) += bytes;
}
map
}
fn get_storage_devices_inner() -> Result<Vec<DeviceInfo>, String> {
let com = COMLibrary::new().map_err(|e| format!("Failed to initialize COM for WMI: {e}"))?;
let wmi = WMIConnection::new(com).map_err(|e| format!("Failed to connect to WMI: {e}"))?;
let raw_drives = query_disk_drives(&wmi)?;
let partitions = query_partitions(&wmi);
let partition_to_drive = build_partition_to_drive_map(&wmi);
let disk_index_to_letters = build_disk_index_to_drive_letters(&partitions, &partition_to_drive);
let disk_index_to_partition_bytes = build_disk_index_to_partition_bytes(&partitions);
drop(wmi);
let msft_physical_sizes = load_msft_physical_disk_sizes();
let mut devices = Vec::with_capacity(raw_drives.len());
for raw in raw_drives {
let device_name = raw.device_id.trim().to_string();
let _drive_letters = raw
.index
.and_then(|index| disk_index_to_letters.get(&index).cloned())
.unwrap_or_default();
let partition_bytes = raw
.index
.and_then(|index| disk_index_to_partition_bytes.get(&index).copied());
let msft_physical_bytes = raw
.index
.and_then(|index| msft_physical_sizes.get(&index).copied());
let size_bytes = resolve_disk_size_bytes(&raw, partition_bytes, msft_physical_bytes);
devices.push(DeviceInfo {
device_name,
vendor_name: raw.manufacturer.clone(),
model_name: raw.model.clone(),
removable: is_removable(&raw.media_type, &raw.interface_type),
size: bytes_to_sectors(size_bytes),
});
}
devices
.sort_by_key(|device| parse_physical_drive_index(&device.device_name).unwrap_or(u32::MAX));
Ok(devices)
}
pub fn get_storage_devices() -> Result<Vec<DeviceInfo>> {
run_wmi_thread("device enumeration", get_storage_devices_inner).map_err(|e| anyhow::anyhow!(e))
}
pub fn parse_physical_drive_index(device_name: &str) -> Option<u32> {
let upper = device_name.to_ascii_uppercase();
let suffix = upper
.strip_prefix(r"\\.\PHYSICALDRIVE")
.or_else(|| upper.strip_prefix("PHYSICALDRIVE"))?;
suffix.parse().ok()
}
pub fn canonical_physical_drive_path(path: &str) -> Result<String, String> {
let trimmed = path.trim();
if trimmed.is_empty() {
return Err("Device path is empty.".into());
}
let index = parse_physical_drive_index(trimmed).ok_or_else(|| {
format!("Device must be a physical drive path (e.g. \\\\.\\PHYSICALDRIVE0), got: {trimmed}")
})?;
Ok(format!(r"\\.\PHYSICALDRIVE{index}"))
}
pub fn validate_block_device_path(path: &str) -> Result<(), String> {
canonical_physical_drive_path(path).map(|_| ())
}
pub fn validate_device_not_system_disk(path: &str) -> Result<(), String> {
let target = parse_physical_drive_index(path)
.ok_or_else(|| format!("Invalid physical drive path: {path}"))?;
let system_disk = system_physical_drive_index()?;
if system_disk == Some(target) {
return Err(format!(
"Refusing {path}: it is the system disk (hosts the Windows boot volume)"
));
}
Ok(())
}
pub fn validate_device_not_busy(path: &str) -> Result<(), String> {
let _ = path;
Ok(())
}
pub fn list_mounted_drive_letters(path: &str) -> Result<Vec<String>, String> {
let target = parse_physical_drive_index(path)
.ok_or_else(|| format!("Invalid physical drive path: {path}"))?;
drive_letters_for_disk_index(target)
}
fn boot_partition_disk_index(wmi: &WMIConnection) -> Option<u32> {
let rows: Vec<HashMap<String, Variant>> = wmi
.raw_query("SELECT DiskIndex FROM Win32_DiskPartition WHERE BootPartition=TRUE")
.ok()?;
rows.into_iter()
.find_map(|row| row.get("DiskIndex").and_then(variant_to_u32))
}
fn system_drive_from_wmi(wmi: &WMIConnection) -> Option<String> {
let rows: Vec<HashMap<String, Variant>> = wmi
.raw_query("SELECT SystemDrive FROM Win32_OperatingSystem")
.ok()?;
rows.into_iter()
.find_map(|row| row.get("SystemDrive").and_then(variant_to_string))
.map(|drive| drive.trim_end_matches('\\').to_ascii_uppercase())
.filter(|drive| !drive.is_empty())
}
fn system_physical_drive_index_inner() -> Result<Option<u32>, String> {
let com = COMLibrary::new().map_err(|e| format!("Failed to initialize COM for WMI: {e}"))?;
let wmi = WMIConnection::new(com).map_err(|e| format!("Failed to connect to WMI: {e}"))?;
if let Some(index) = boot_partition_disk_index(&wmi) {
return Ok(Some(index));
}
let system_drive = system_drive_from_wmi(&wmi).or_else(system_drive_from_env);
let Some(system_drive) = system_drive else {
debug!("Could not determine Windows system drive; skipping system-disk check");
return Ok(None);
};
let partitions = query_partitions(&wmi);
let partition_to_drive = build_partition_to_drive_map(&wmi);
for partition in partitions {
let Some(disk_index) = partition.disk_index else {
continue;
};
let Some(letters) = partition_to_drive.get(&partition.device_id) else {
continue;
};
if letters
.iter()
.any(|letter| letter.eq_ignore_ascii_case(&system_drive))
{
return Ok(Some(disk_index));
}
}
Ok(None)
}
pub fn system_physical_drive_index() -> Result<Option<u32>, String> {
run_wmi_thread("system disk lookup", system_physical_drive_index_inner)
}
fn drive_letters_for_disk_index_inner(disk_index: u32) -> Result<Vec<String>, String> {
let com = COMLibrary::new().map_err(|e| format!("Failed to initialize COM for WMI: {e}"))?;
let wmi = WMIConnection::new(com).map_err(|e| format!("Failed to connect to WMI: {e}"))?;
let partitions = query_partitions(&wmi);
let partition_to_drive = build_partition_to_drive_map(&wmi);
let mut letters = Vec::new();
for partition in partitions {
if partition.disk_index != Some(disk_index) {
continue;
}
if let Some(mut mapped) = partition_to_drive.get(&partition.device_id).cloned() {
letters.append(&mut mapped);
}
}
letters.sort();
letters.dedup();
Ok(letters)
}
fn drive_letters_for_disk_index(disk_index: u32) -> Result<Vec<String>, String> {
run_wmi_thread("drive letter lookup", move || {
drive_letters_for_disk_index_inner(disk_index)
})
}
pub fn mounted_drive_letters_for_disk_index(disk_index: u32) -> Result<Vec<String>, String> {
drive_letters_for_disk_index(disk_index)
}
fn wmi_path_to_volume_device_path(path: &str) -> String {
let trimmed = path.trim();
if let Some(rest) = trimmed.strip_prefix(r"\\?\") {
format!(r"\\.\{rest}")
} else if trimmed.starts_with(r"\\.\") {
trimmed.to_string()
} else if trimmed.len() == 1 && trimmed.chars().all(|c| c.is_ascii_alphabetic()) {
format!(r"\\.\{}:", trimmed.to_ascii_uppercase())
} else if trimmed.ends_with(':') {
format!(r"\\.\{trimmed}")
} else {
format!(r"\\.\{trimmed}")
}
}
fn msft_volume_dismount_targets_inner(disk_index: u32) -> Vec<(String, String)> {
let com = match COMLibrary::new() {
Ok(com) => com,
Err(e) => {
debug!("MSFT_Volume query skipped (COM init failed): {e}");
return Vec::new();
}
};
let storage_wmi =
match WMIConnection::with_namespace_path("ROOT\\Microsoft\\Windows\\Storage", com) {
Ok(wmi) => wmi,
Err(e) => {
debug!("MSFT_Volume query skipped (storage WMI failed): {e}");
return Vec::new();
}
};
let rows: Vec<HashMap<String, Variant>> = match storage_wmi
.raw_query("SELECT DeviceId, DriveLetter, Path, DiskNumber FROM MSFT_Volume")
{
Ok(rows) => rows,
Err(e) => {
debug!("MSFT_Volume query failed: {e}");
return Vec::new();
}
};
let mut targets = Vec::new();
for row in rows {
let Some(volume_disk) = row.get("DiskNumber").and_then(variant_to_u32) else {
continue;
};
if volume_disk != disk_index {
continue;
}
let label = row
.get("DriveLetter")
.and_then(variant_to_string)
.filter(|value| !value.is_empty())
.or_else(|| row.get("DeviceId").and_then(variant_to_string))
.or_else(|| row.get("Path").and_then(variant_to_string))
.unwrap_or_else(|| format!("disk {disk_index} volume"));
let device_path = if let Some(letter) = row.get("DriveLetter").and_then(variant_to_string) {
let normalized = letter.trim().trim_end_matches(':').to_ascii_uppercase();
format!(r"\\.\{normalized}:")
} else if let Some(path) = row.get("Path").and_then(variant_to_string) {
wmi_path_to_volume_device_path(&path)
} else if let Some(device_id) = row.get("DeviceId").and_then(variant_to_string) {
wmi_path_to_volume_device_path(&device_id)
} else {
continue;
};
targets.push((device_path, label));
}
targets
}
pub fn volume_dismount_targets_for_disk_index(
disk_index: u32,
) -> Result<Vec<(String, String)>, String> {
run_wmi_thread("volume target lookup", move || {
let mut targets = msft_volume_dismount_targets_inner(disk_index);
if let Ok(letters) = drive_letters_for_disk_index_inner(disk_index) {
for letter in letters {
let normalized = letter.trim().trim_end_matches(':').to_ascii_uppercase();
let device_path = format!(r"\\.\{normalized}:");
if targets
.iter()
.any(|(path, _)| path.eq_ignore_ascii_case(&device_path))
{
continue;
}
targets.push((device_path, format!("{normalized}:")));
}
}
Ok(targets)
})
}
fn system_drive_from_env() -> Option<String> {
std::env::var_os("SystemDrive").map(|value| {
value
.to_string_lossy()
.trim_end_matches('\\')
.to_ascii_uppercase()
})
}
pub fn device_path_matches(device_name: &str, query_path: &str) -> bool {
fn normalize(path: &str) -> String {
let trimmed = path.trim();
let without_prefix = trimmed
.strip_prefix(r"\\.\")
.or_else(|| trimmed.strip_prefix(r"\\.\"))
.unwrap_or(trimmed);
without_prefix.to_ascii_uppercase()
}
normalize(device_name) == normalize(query_path)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn device_path_matches_physical_drive_aliases() {
assert!(device_path_matches(
r"\\.\PHYSICALDRIVE0",
r"\\.\PhysicalDrive0"
));
assert!(device_path_matches("PhysicalDrive1", r"\\.\PHYSICALDRIVE1"));
assert!(!device_path_matches(
r"\\.\PHYSICALDRIVE0",
r"\\.\PHYSICALDRIVE1"
));
}
#[test]
fn parse_physical_drive_index_handles_prefixes() {
assert_eq!(parse_physical_drive_index(r"\\.\PHYSICALDRIVE2"), Some(2));
assert_eq!(parse_physical_drive_index("PhysicalDrive3"), Some(3));
}
#[test]
fn canonical_physical_drive_path_normalizes_input() {
assert_eq!(
canonical_physical_drive_path(r"\\.\PhysicalDrive1").unwrap(),
r"\\.\PHYSICALDRIVE1"
);
assert!(canonical_physical_drive_path(r"\\.\C:").is_err());
}
#[test]
fn variant_to_u64_parses_wmi_integer_shapes() {
assert_eq!(
variant_to_u64(&Variant::UI8(1_603_901_849_6)),
Some(1_603_901_849_6)
);
assert_eq!(variant_to_u64(&Variant::UI4(512)), Some(512));
assert_eq!(variant_to_u64(&Variant::Null), None);
}
}