1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use winapi::um::winioctl::{IOCTL_DISK_GET_DRIVE_GEOMETRY, DISK_GEOMETRY};
use crate::*;

#[derive(Clone)]
#[cfg_attr(any(feature = "std", test), derive(Debug))]
pub struct Disk(Handle);

impl From<Handle> for Disk {
    fn from(handle: Handle) -> Self {
        Self(handle)
    }
}

impl Disk {
    fn index_to_name(index: u32) -> NtString {
        let name = crate::format!("\\??\\PhysicalDrive{}", index);
        NtString::from(name)
    }

    pub fn open_nth(index: u32) -> Result<Self> {
        let name = Self::index_to_name(index);
        Self::open(&name)
    }

    pub fn open(name: &NtString) -> Result<Self> {
        let (handle, _) = NewHandle::device(Access::GENERIC_READ | Access::GENERIC_WRITE).build(name)?;

        Ok(Self(handle))
    }

    pub fn open_nth_readonly(index: u32) -> Result<Self> {
        let name = Self::index_to_name(index);
        Self::open_readonly(&name)
    }

    pub fn open_readonly(name: &NtString) -> Result<Self> {
        let (handle, _) = NewHandle::device(Access::SYNCHRONIZE).build(name)?;

        Ok(Self(handle))
    }

    /// Usefull only to get device_name without admin rights.
    pub fn open_nth_info(index: u32) -> Result<Self> {
        let name = Self::index_to_name(index);
        Self::open_info(&name)
    }

    /// Usefull only to get device_name without admin rights.
    pub fn open_info(name: &NtString) -> Result<Self> {
        let (handle, _) = NewHandle::device(Access::SYNCHRONIZE).build(name)?;

        Ok(Self(handle))
    }

    /// Returns the disk device name: `\Device\Harddisk0\DR0`.
    ///
    /// To pass this name to a Win32 function add `\\?\GLOBALROOT\` prefix: `\\?\GLOBALROOT\ \Device\Harddisk0\DR0`
    pub fn device_name(&self) -> Result<NtString> {
        self.0.object_name()
    }

    pub fn capacity(&self) -> Result<u64> {
        ioctl::length(&self.0)
    }

    pub fn is_readonly(&self) -> Result<bool> {
        ioctl::is_readonly(&self.0)
    }

    pub fn is_offline(&self) -> Result<bool> {
        const DISK_ATTRIBUTE_OFFLINE: u64 = 0x01;
        let attr = self.attributes()?;
        Ok((attr & DISK_ATTRIBUTE_OFFLINE) == DISK_ATTRIBUTE_OFFLINE)
    }

    pub fn is_removable(&self) -> Result<bool> {
        ioctl::is_removable(&self.0)
    }

    pub fn is_trim_enabled(&self) -> Result<bool> {
        ioctl::is_trim_enabled(&self.0)
    }
    pub fn has_seek_penalty(&self) -> Result<bool> {
        ioctl::has_seek_penalty(&self.0)
    }

    pub fn device_number(&self) -> Result<u32> {
        ioctl::device_number(&self.0)
    }

    pub fn attributes(&self) -> Result<u64> {
        ioctl::attributes(&self.0)
    }

    pub fn geometry(&self) -> Result<DISK_GEOMETRY> {
        let mut data = unsafe { StructBuffer::<DISK_GEOMETRY>::new() };
        self.0.ioctl_query(IOCTL_DISK_GET_DRIVE_GEOMETRY, &mut data)?;

        Ok(data.take())
    }
}

impl Flush for Disk {
    fn flush(&self) -> Result<()> {
        self.0.flush()
    }
}

impl Size for Disk {
    fn size(&self) -> Result<u64> {
        Disk::capacity(self)
    }
}

impl ReadAt for Disk {
    fn read_at(&self, offset: u64, buffer: &mut [u8]) -> Result<usize> {
        self.0.read(buffer, Some(offset))
    }
}

impl WriteAt for Disk {
    fn write_at(&self, offset: u64, data: &[u8]) -> Result<usize> {
        self.0.write(data, Some(offset))
    }
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;

    #[test] // needs admin rights
    fn open() {
        if std::env::var("NT_NATIVE_TEST_ADMIN").is_ok() {
            let disk = Disk::open_nth(0).unwrap();
            let mut buffer = vec![0_u8; 512];
            let readed = disk.read_at(0, &mut buffer).unwrap();
            assert_eq!(readed, buffer.len());

            println!("First bytes (of {}):", disk.capacity().unwrap());
            for chunk in buffer.chunks(16).take(4) {
                println!("{:02x?}", chunk)
            }
        } else {
            print!("Non admin, skipped ... ");
        }
    }

    #[test]
    fn open_disk_info() {
        let disk = Disk::open_nth_info(0).unwrap();

        let name = disk.device_name().unwrap();
        println!("DeviceName: {}", name.to_string());
    }
}