Skip to main content

heim_disk/sys/windows/bindings/
drive.rs

1use std::ffi::OsString;
2use std::os::windows::ffi::OsStringExt;
3use std::path::PathBuf;
4use std::ptr;
5use std::str::FromStr;
6
7use winapi::ctypes::wchar_t;
8use winapi::shared::minwindef::{DWORD, MAX_PATH};
9use winapi::um::{errhandlingapi, fileapi, winbase};
10
11use crate::os::windows::DriveType;
12use crate::os::windows::Flags;
13use crate::FileSystem;
14use heim_common::prelude::{Error, Result};
15
16// According to winapi docs 50 is a reasonable length to accomodate the volume path
17// https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-getvolumenameforvolumemountpointw
18const VOLUME_MAX_LEN: DWORD = 50;
19
20/// Represents the logical drive.
21///
22/// Vec inside should contain path like `C:\\\0`,
23/// four `wchar_t`s: drive letter, colon, backslash and a zero char.
24#[derive(Debug)]
25pub struct Drive([wchar_t; 4]);
26
27impl Drive {
28    /// Returns the inner buffer as a `PathBuf` (ex. `C:\`) withour trailing zero char.
29    pub fn to_path_buf(&self) -> PathBuf {
30        // Skipping trailing zero char also
31        PathBuf::from(OsString::from_wide(&self.0[..(self.0.len() - 1)]))
32    }
33
34    pub fn volume_name(&self) -> Result<OsString> {
35        let mut volume: Vec<wchar_t> = Vec::with_capacity(VOLUME_MAX_LEN as usize);
36        let result = unsafe {
37            fileapi::GetVolumeNameForVolumeMountPointW(
38                self.0.as_ptr(),
39                volume.as_mut_ptr(),
40                VOLUME_MAX_LEN,
41            )
42        };
43
44        if result == 0 {
45            Err(Error::last_os_error())
46        } else {
47            unsafe {
48                volume.set_len(VOLUME_MAX_LEN as usize);
49            }
50
51            let str_end = volume.iter().position(|chr| *chr == 0x00).unwrap_or(0);
52            volume.truncate(str_end);
53            let volume_path = OsString::from_wide(&volume);
54
55            Ok(volume_path)
56        }
57    }
58
59    /// ## Returns
60    ///
61    /// `Ok(Some(..))` - successful info fetch
62    /// `Ok(None)` - disk should be ignored
63    /// `Err(..)` - whoops
64    pub fn information(&self) -> Result<Option<(Option<DriveType>, Flags, FileSystem)>> {
65        let drive_type = DriveType::from_slice(&self.0);
66
67        let mut flags: DWORD = 0;
68        let mut fs_type: Vec<wchar_t> = Vec::with_capacity(MAX_PATH + 1);
69        let mut old_mode: DWORD = 0;
70
71        let err_mode_result = unsafe {
72            errhandlingapi::SetThreadErrorMode(winbase::SEM_FAILCRITICALERRORS, &mut old_mode)
73        };
74        if err_mode_result == 0 {
75            return Err(Error::last_os_error());
76        }
77
78        let result = unsafe {
79            fileapi::GetVolumeInformationW(
80                self.0.as_ptr(),
81                ptr::null_mut(),
82                // Originally Windows `ARRAYSIZE` macro is used here,
83                // and we need to calculate length in bytes (while having u16 chars)
84                (self.0.len() * 2) as DWORD,
85                ptr::null_mut(),
86                ptr::null_mut(),
87                &mut flags,
88                fs_type.as_mut_ptr(),
89                fs_type.capacity() as DWORD,
90            )
91        };
92
93        let err_mode_result =
94            unsafe { errhandlingapi::SetThreadErrorMode(old_mode, ptr::null_mut()) };
95        if err_mode_result == 0 {
96            return Err(Error::last_os_error());
97        }
98
99        match result {
100            // Same to `psutil` and `gopsutil` we are going to ignore any errors
101            // from the CDRoms, floppies and other removable disks
102            0 if drive_type == Some(DriveType::CdRom) => Ok(None),
103            0 if drive_type == Some(DriveType::Removable) => Ok(None),
104            0 => Err(Error::last_os_error()),
105            _ => {
106                let flags = Flags::from_bits_truncate(flags);
107
108                // So, since `GetVolumeInformationW` does not returns how much bytes
109                // it wrote into a passed buffer, we need to find that manually.
110                //
111                // Quite unsafe, because we are going to poke some random memory here.
112                // It is still in the range of pre-allocated capacity buffer (see above),
113                // but it might be filled with a garbage (let's hope it is not).
114                unsafe {
115                    fs_type.set_len(MAX_PATH + 1);
116                }
117                let str_end = fs_type.iter().position(|chr| *chr == 0x00).unwrap_or(0);
118                fs_type.truncate(str_end);
119
120                let fs_type_str = OsString::from_wide(&fs_type);
121                let fs = FileSystem::from_str(&fs_type_str.to_string_lossy())?;
122
123                Ok(Some((drive_type, flags, fs)))
124            }
125        }
126    }
127}
128
129impl<T> From<T> for Drive
130where
131    T: AsRef<[u16]>,
132{
133    fn from(data: T) -> Drive {
134        let buffer = data.as_ref();
135        debug_assert!(buffer.len() == 4);
136        debug_assert!(buffer[0] >= 0x0041 && buffer[0] <= 0x005a);
137        debug_assert!(buffer[1] == 0x003a);
138        debug_assert!(buffer[2] == 0x005c);
139        debug_assert!(buffer[3] == 0x0000);
140
141        let mut inner = [0; 4];
142        inner.copy_from_slice(buffer);
143        Drive(inner)
144    }
145}