use crate::{DiskExt, DiskType};
use std::ffi::{OsStr, OsString};
use std::mem::size_of;
use std::path::Path;
use winapi::ctypes::c_void;
use winapi::shared::minwindef::{DWORD, MAX_PATH};
use winapi::um::fileapi::{
CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW,
OPEN_EXISTING,
};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::winbase::{DRIVE_FIXED, DRIVE_REMOVABLE};
use winapi::um::winioctl::{
PropertyStandardQuery, StorageDeviceSeekPenaltyProperty, IOCTL_STORAGE_QUERY_PROPERTY,
STORAGE_PROPERTY_QUERY,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE, ULARGE_INTEGER};
#[doc = include_str!("../../md_doc/disk.md")]
pub struct Disk {
type_: DiskType,
name: OsString,
file_system: Vec<u8>,
mount_point: Vec<u16>,
s_mount_point: String,
total_space: u64,
available_space: u64,
is_removable: bool,
}
impl DiskExt for Disk {
fn type_(&self) -> DiskType {
self.type_
}
fn name(&self) -> &OsStr {
&self.name
}
fn file_system(&self) -> &[u8] {
&self.file_system
}
fn mount_point(&self) -> &Path {
Path::new(&self.s_mount_point)
}
fn total_space(&self) -> u64 {
self.total_space
}
fn available_space(&self) -> u64 {
self.available_space
}
fn is_removable(&self) -> bool {
self.is_removable
}
fn refresh(&mut self) -> bool {
if self.total_space != 0 {
unsafe {
let mut tmp: ULARGE_INTEGER = std::mem::zeroed();
if GetDiskFreeSpaceExW(
self.mount_point.as_ptr(),
std::ptr::null_mut(),
std::ptr::null_mut(),
&mut tmp,
) != 0
{
self.available_space = *tmp.QuadPart();
return true;
}
}
}
false
}
}
struct HandleWrapper(HANDLE);
impl HandleWrapper {
unsafe fn new(drive_name: &[u16], open_rights: DWORD) -> Option<Self> {
let handle = CreateFileW(
drive_name.as_ptr(),
open_rights,
FILE_SHARE_READ | FILE_SHARE_WRITE,
std::ptr::null_mut(),
OPEN_EXISTING,
0,
std::ptr::null_mut(),
);
if handle == INVALID_HANDLE_VALUE {
CloseHandle(handle);
None
} else {
Some(Self(handle))
}
}
}
impl Drop for HandleWrapper {
fn drop(&mut self) {
unsafe {
CloseHandle(self.0);
}
}
}
unsafe fn get_drive_size(mount_point: &[u16]) -> Option<(u64, u64)> {
let mut total_size: ULARGE_INTEGER = std::mem::zeroed();
let mut available_space: ULARGE_INTEGER = std::mem::zeroed();
if GetDiskFreeSpaceExW(
mount_point.as_ptr(),
std::ptr::null_mut(),
&mut total_size,
&mut available_space,
) != 0
{
Some((
*total_size.QuadPart() as _,
*available_space.QuadPart() as _,
))
} else {
None
}
}
#[allow(non_snake_case)]
#[repr(C)]
struct DEVICE_SEEK_PENALTY_DESCRIPTOR {
Version: DWORD,
Size: DWORD,
IncursSeekPenalty: BOOLEAN,
}
pub(crate) unsafe fn get_disks() -> Vec<Disk> {
let drives = GetLogicalDrives();
if drives == 0 {
return Vec::new();
}
#[cfg(feature = "multithread")]
use rayon::iter::ParallelIterator;
crate::utils::into_iter(0..DWORD::BITS)
.filter_map(|x| {
if (drives >> x) & 1 == 0 {
return None;
}
let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];
let drive_type = GetDriveTypeW(mount_point.as_ptr());
let is_removable = drive_type == DRIVE_REMOVABLE;
if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
return None;
}
let mut name = [0u16; MAX_PATH + 1];
let mut file_system = [0u16; 32];
if GetVolumeInformationW(
mount_point.as_ptr(),
name.as_mut_ptr(),
name.len() as DWORD,
std::ptr::null_mut(),
std::ptr::null_mut(),
std::ptr::null_mut(),
file_system.as_mut_ptr(),
file_system.len() as DWORD,
) == 0
{
return None;
}
let mut pos = 0;
for x in name.iter() {
if *x == 0 {
break;
}
pos += 1;
}
let name = String::from_utf16_lossy(&name[..pos]);
let name = OsStr::new(&name);
pos = 0;
for x in file_system.iter() {
if *x == 0 {
break;
}
pos += 1;
}
let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();
let drive_name = [
b'\\' as u16,
b'\\' as u16,
b'.' as u16,
b'\\' as u16,
b'A' as u16 + x as u16,
b':' as u16,
0,
];
let handle = HandleWrapper::new(&drive_name, 0)?;
let (total_space, available_space) = get_drive_size(&mount_point)?;
if total_space == 0 {
return None;
}
let mut spq_trim = STORAGE_PROPERTY_QUERY {
PropertyId: StorageDeviceSeekPenaltyProperty,
QueryType: PropertyStandardQuery,
AdditionalParameters: [0],
};
let mut result: DEVICE_SEEK_PENALTY_DESCRIPTOR = std::mem::zeroed();
let mut dw_size = 0;
let type_ = if DeviceIoControl(
handle.0,
IOCTL_STORAGE_QUERY_PROPERTY,
&mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
&mut result as *mut DEVICE_SEEK_PENALTY_DESCRIPTOR as *mut c_void,
size_of::<DEVICE_SEEK_PENALTY_DESCRIPTOR>() as DWORD,
&mut dw_size,
std::ptr::null_mut(),
) == 0
|| dw_size != size_of::<DEVICE_SEEK_PENALTY_DESCRIPTOR>() as DWORD
{
DiskType::Unknown(-1)
} else {
let is_ssd = result.IncursSeekPenalty == 0;
if is_ssd {
DiskType::SSD
} else {
DiskType::HDD
}
};
Some(Disk {
type_,
name: name.to_owned(),
file_system: file_system.to_vec(),
mount_point: mount_point.to_vec(),
s_mount_point: String::from_utf16_lossy(&mount_point[..mount_point.len() - 1]),
total_space,
available_space,
is_removable,
})
})
.collect::<Vec<_>>()
}