win-drives 0.1.0

use core::ops::Deref;
use heapless::{String, format};
use winapi::shared::ntdef::HANDLE;
use crate::{DeviceNumber, DiskGeometry, DriveLayout, DriverError, api::*};

#[derive(Debug)]
pub struct PhysicalDrive {
    handle: HANDLE,
    geometry: DiskGeometry,
    position: u64,
}

impl Deref for PhysicalDrive {
    type Target = DiskGeometry;

    fn deref(&self) -> &Self::Target {
        &self.geometry
    }
}

impl PhysicalDrive {
    pub fn enumerate() -> PhysicalDriveIter {
        PhysicalDriveIter {
            next_num: 0,
            consecutive_misses: 0,
            done: false,
        }
    }

    pub fn open(drive_num: u8) -> Result<Self, DriverError> {
        let obj_name: String<_> = format!(20; r#"\??\PhysicalDrive{}"#, drive_num).unwrap();
        let handle = open_handle(obj_name)?;
        let geometry = geometry(handle)?.into();

        Ok(PhysicalDrive {
            handle,
            geometry,
            position: 0
        })
    }

    pub fn device_number(&self) -> Result<DeviceNumber, DriverError> {
        Ok(device_number(self.handle)?.into())
    }

    pub fn layout(&self) -> Result<DriveLayout, DriverError> {
        Ok(drive_layout(self.handle)?.into())
    }

    pub fn read_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<usize, DriverError> {
        let sector_size = self.geometry.bytes_per_sector() as usize;
        if buf.len() < sector_size || buf.len() % sector_size != 0 {
            return Err(DriverError::InvalidParameter);
        }

        let read = read_file(self.handle, buf, offset)?;

        Ok(read)
    }
}

#[cfg(feature = "std")]
impl std::io::Read for PhysicalDrive {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
   
        let sector = self.geometry.bytes_per_sector() as usize;
        if buf.len() < sector {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                DriverError::BufferTooSmall {
                    needed: sector as u32,
                    got: buf.len(),
                },
            ));
        }
        let aligned_len = (buf.len() / sector) * sector;

        let n = self.read_at(self.position, &mut buf[..aligned_len])
            .map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))?;

        self.position += n as u64;
        Ok(n)
    }
}

#[cfg(feature = "std")]
impl std::io::Seek for PhysicalDrive {
    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
        use std::io::SeekFrom;
        let sector = self.geometry.bytes_per_sector() as u64;

        let new_pos = match pos {
            SeekFrom::Start(p) => p,
            SeekFrom::Current(p) => (self.position as i128 + p as i128) as u64,
            SeekFrom::End(p) => {
                let end = self.geometry.size() as i128;
                (end + p as i128) as u64
            }
        };

        if new_pos % sector != 0 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                DriverError::Unaligned {
                    offset: new_pos,
                    sector: sector as u32,
                },
            ));
        }

        self.position = new_pos;
        Ok(new_pos)
    }
}

#[cfg(feature = "no-std")]
impl embedded_io::Error for DriverError {
    fn kind(&self) -> embedded_io::ErrorKind {
        embedded_io::ErrorKind::Other
    }
}

#[cfg(feature = "no-std")]
impl embedded_io::ErrorType for PhysicalDrive {
    type Error = DriverError;
}

#[cfg(feature = "no-std")]
impl embedded_io::Read for PhysicalDrive {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {

        let sector = self.geometry.bytes_per_sector() as usize;
        if buf.len() < sector {
            return Err(DriverError::BufferTooSmall {
                needed: sector as u32,
                got: buf.len(),
            });
        }
        let aligned_len = (buf.len() / sector) * sector;
        let n = self.read_at(self.position, &mut buf[..aligned_len])?;
        self.position += n as u64;
        Ok(n)
    }
}

#[cfg(feature = "no-std")]
impl embedded_io::Seek for PhysicalDrive {
    fn seek(&mut self, pos: embedded_io::SeekFrom) -> Result<u64, Self::Error> {
        use embedded_io::SeekFrom;
        let sector = self.geometry.bytes_per_sector() as u64;

        let new_pos = match pos {
            SeekFrom::Start(p) => p,
            SeekFrom::Current(p) => (self.position as i128 + p as i128) as u64,
            SeekFrom::End(p) => {
                let end = self.geometry.size() as i128;
                (end + p as i128) as u64
            }
        };

        if new_pos % sector != 0 {
            return Err(DriverError::Unaligned {
                offset: new_pos,
                sector: sector as u32,
            });
        }

        self.position = new_pos;
        Ok(new_pos)
    }
}

pub struct PhysicalDriveIter {
    next_num: u8,
    consecutive_misses: u8,
    done: bool,
}

impl PhysicalDriveIter {
    const MAX_DRIVES: u8 = 64;
    const MISS_THRESHOLD: u8 = 8;
}

impl Iterator for PhysicalDriveIter {
    type Item = PhysicalDrive;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }

        while self.next_num < Self::MAX_DRIVES {
            let num = self.next_num;
            self.next_num += 1;

            match PhysicalDrive::open(num) {
                Ok(drive) => {
                    self.consecutive_misses = 0;
                    return Some(drive);
                }
                Err(DriverError::NotFound) | Err(DriverError::PathNotFound) => {
                    self.consecutive_misses += 1;
                    if self.consecutive_misses >= Self::MISS_THRESHOLD {
                        self.done = true;
                        return None;
                    }
                }
                Err(DriverError::Permission) => {
                    self.consecutive_misses = 0;
                    continue;
                }
                Err(_) => {
                    self.done = true;
                    return None;
                }
            }
        }

        self.done = true;
        None
    }
}