From d583029506d52ec32d46b819b704f991b62b7bb2 Mon Sep 17 00:00:00 2001 From: kostassoid Date: Mon, 23 Sep 2019 02:22:23 +0300 Subject: [PATCH] Use aligned buffer for verification --- src/actions/wipe.rs | 11 +++++----- src/main.rs | 14 ++++++++++++- src/sanitization/mem.rs | 44 +++++++++++++++++++++++++++++++++++++++ src/sanitization/mod.rs | 4 +++- src/sanitization/stage.rs | 17 ++++----------- src/storage/mod.rs | 2 ++ src/storage/nix/macos.rs | 29 +++++++++++++++++++++++--- src/storage/nix/mod.rs | 26 +---------------------- src/storage/windows.rs | 14 ++----------- 9 files changed, 101 insertions(+), 60 deletions(-) create mode 100644 src/sanitization/mem.rs diff --git a/src/actions/wipe.rs b/src/actions/wipe.rs index ff690ef..89b96bb 100644 --- a/src/actions/wipe.rs +++ b/src/actions/wipe.rs @@ -1,4 +1,5 @@ use crate::sanitization::*; +use crate::sanitization::mem::*; use crate::storage::StorageAccess; use std::io::{Error, ErrorKind}; use std::rc::Rc; @@ -60,7 +61,7 @@ pub trait WipeEventReceiver { } impl WipeTask { - pub fn run(self, access: &mut StorageAccess, state: &mut WipeState, frontend: &mut WipeEventReceiver) -> bool { + pub fn run(self, access: &mut dyn StorageAccess, state: &mut WipeState, frontend: &mut dyn WipeEventReceiver) -> bool { frontend.handle(&self, state, WipeEvent::Started); @@ -117,7 +118,7 @@ impl WipeTask { } } -fn fill(access: &mut StorageAccess, task: &WipeTask, state: &mut WipeState, stage: &Stage, frontend: &mut WipeEventReceiver) -> Option> { +fn fill(access: &mut dyn StorageAccess, task: &WipeTask, state: &mut WipeState, stage: &Stage, frontend: &mut dyn WipeEventReceiver) -> Option> { let mut stream = stage.stream( task.total_size, @@ -156,7 +157,7 @@ fn fill(access: &mut StorageAccess, task: &WipeTask, state: &mut WipeState, stag None } -fn verify(access: &mut StorageAccess, task: &WipeTask, state: &mut WipeState, stage: &Stage, frontend: &mut WipeEventReceiver) -> Option> { +fn verify(access: &mut dyn StorageAccess, task: &WipeTask, state: &mut WipeState, stage: &Stage, frontend: &mut dyn WipeEventReceiver) -> Option> { frontend.handle(task, state, WipeEvent::StageStarted); @@ -172,7 +173,7 @@ fn verify(access: &mut StorageAccess, task: &WipeTask, state: &mut WipeState, st state.position ); - let mut buf: Vec = vec![0; task.block_size]; + let mut buf = alloc_aligned_byte_vec(task.block_size, task.block_size); while let Some(chunk) = stream.next() { let b = &mut buf[..chunk.len()]; @@ -387,4 +388,4 @@ mod test { Ok(()) } } -} \ No newline at end of file +} diff --git a/src/main.rs b/src/main.rs index d478c81..385616f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -90,8 +90,20 @@ fn main() { ) .get_matches(); - let storage_devices = System::get_storage_devices().unwrap(); //todo: handle errors + let storage_devices = System::get_storage_devices() + .unwrap_or_else(|err| { + eprintln!("Unable to enumerate storage devices. {}", err); + let is_wsl = std::fs::read_to_string("/proc/version") + .map(|v| v.contains("Microsoft")) + .unwrap_or(false); + + if is_wsl { + eprintln!("WSL is not supported at the moment as it doesn't provide direct storage device access."); + } + + std::process::exit(1); + }); let frontend = cli::ConsoleFrontend::new(); match app.subcommand() { diff --git a/src/sanitization/mem.rs b/src/sanitization/mem.rs new file mode 100644 index 0000000..4cde428 --- /dev/null +++ b/src/sanitization/mem.rs @@ -0,0 +1,44 @@ +pub fn alloc_aligned_byte_vec(size: usize, align: usize) -> Vec { + unsafe { + let buf_layout = std::alloc::Layout::from_size_align_unchecked(size, align); + let buf_ptr = std::alloc::alloc(buf_layout); + Vec::from_raw_parts(buf_ptr, size, size) + } +} + +pub fn fill_byte_slice(buf: &mut [u8], value: u8) { + // unsafe { + // std::libc::memset( + // buf.as_mut_ptr() as _, + // 0, + // buf.len() + // ); + // }; + buf.iter_mut().map(|x| *x = value).count(); //todo: rewrite +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_aligned_allocation() { + let size = 65536; + let align = 4096; + + let buf = alloc_aligned_byte_vec(size, align); + + assert_eq!(buf.len(), size); + assert_eq!(buf.as_ptr() as usize % align, 0); + } + + #[test] + fn test_vec_fill() { + let mut buf = vec![0x00; 1024]; + + fill_byte_slice(&mut buf, 0xff); + + assert_eq!(buf.iter().filter(|x| **x != 0xff).count(), 0); + } + +} diff --git a/src/sanitization/mod.rs b/src/sanitization/mod.rs index 89738ef..ae167f4 100644 --- a/src/sanitization/mod.rs +++ b/src/sanitization/mod.rs @@ -1,7 +1,9 @@ pub mod stage; +pub use stage::*; + +pub mod mem; use std::collections::BTreeMap; -pub use stage::*; #[derive(Debug, Clone)] pub struct Scheme { diff --git a/src/sanitization/stage.rs b/src/sanitization/stage.rs index 3bda4dc..cf82474 100644 --- a/src/sanitization/stage.rs +++ b/src/sanitization/stage.rs @@ -3,6 +3,8 @@ use rand::SeedableRng; use rand::RngCore; pub use streaming_iterator::StreamingIterator; +use super::mem::*; + const RANDOM_SEED_SIZE: usize = 32; type RandomGenerator = rand_chacha::ChaCha8Rng; @@ -55,22 +57,11 @@ impl Stage { pub fn stream(&self, total_size: u64, block_size: usize, start_from: u64) -> SanitizationStream { - let mut buf = unsafe { - let buf_layout = std::alloc::Layout::from_size_align_unchecked(block_size, block_size); - let buf_ptr = std::alloc::alloc(buf_layout); - Vec::from_raw_parts(buf_ptr, block_size, block_size) - }; + let mut buf = alloc_aligned_byte_vec(block_size, block_size); let kind = match self { Stage::Fill { value } => { - // unsafe { - // std::libc::memset( - // buf.as_mut_ptr() as _, - // 0, - // buf.len() - // ); - // }; - buf.iter_mut().map(|x| *x = *value).count(); //todo: rewrite + fill_byte_slice(&mut buf, *value); StreamKind::Fill }, Stage::Random { seed } => { diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 7434cd4..7096f6c 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -1,9 +1,11 @@ #[cfg(unix)] pub use self::nix::*; +#[cfg(unix)] pub mod nix; #[cfg(windows)] pub use windows::*; +#[cfg(windows)] pub mod windows; pub type IoResult = std::io::Result; diff --git a/src/storage/nix/macos.rs b/src/storage/nix/macos.rs index fa9776f..f6f8ea3 100644 --- a/src/storage/nix/macos.rs +++ b/src/storage/nix/macos.rs @@ -1,10 +1,12 @@ //extern crate IOKit_sys as iokit; use std::fs::{File, OpenOptions}; -use std::path::Path; -use crate::storage::*; +use std::path::{Path, PathBuf}; +use std::fs::read_dir; use std::os::unix::io::*; use ::nix::*; +use crate::storage::*; + pub fn open_file_direct>(file_path: P, write_access: bool) -> IoResult { let file = OpenOptions::new() @@ -48,13 +50,34 @@ pub fn is_trim_supported(fd: RawFd) -> bool { } pub fn get_storage_devices() -> IoResult> { - super::discover_file_based_devices( + discover_file_based_devices( "/dev", |p| p.to_str().unwrap().contains("/dev/rdisk"), |_m| true ) } +fn discover_file_based_devices>( + root: P, + path_filter: fn(&PathBuf) -> bool, + meta_filter: fn(&StorageDetails) -> bool +) -> IoResult> { + let rd = read_dir(&root)?; + let mut refs = rd.filter_map(std::io::Result::ok) + .map(|de| de.path()) + .filter(|path| + (path_filter)(&path.to_path_buf()) + ) + .flat_map(FileRef::new) + .filter(|r| + (meta_filter)(&r.details) + ) + .collect::>(); + + refs.sort_by(|a, b| a.path.to_str().cmp(&b.path.to_str())); + Ok(refs) +} + /* fn get_mounts() -> IoResult<()> { unsafe { diff --git a/src/storage/nix/mod.rs b/src/storage/nix/mod.rs index dfe12eb..7b84580 100644 --- a/src/storage/nix/mod.rs +++ b/src/storage/nix/mod.rs @@ -1,8 +1,5 @@ #![cfg(unix)] -extern crate nix; - use std::fs::File; -use std::fs::read_dir; use std::io::prelude::*; use std::path::{Path, PathBuf}; use std::io::SeekFrom; @@ -35,27 +32,6 @@ fn resolve_storage_size(file_type: &FileType, stat: &libc::stat, fd: RawFd) -> u } } -fn discover_file_based_devices>( - root: P, - path_filter: fn(&PathBuf) -> bool, - meta_filter: fn(&StorageDetails) -> bool -) -> IoResult> { - let rd = read_dir(&root)?; - let mut refs = rd.filter_map(std::io::Result::ok) - .map(|de| de.path()) - .filter(|path| - (path_filter)(&path.to_path_buf()) - ) - .flat_map(FileRef::new) - .filter(|r| - (meta_filter)(&r.details) - ) - .collect::>(); - - refs.sort_by(|a, b| a.path.to_str().cmp(&b.path.to_str())); - Ok(refs) -} - #[derive(Debug)] pub struct FileAccess { file: File @@ -157,4 +133,4 @@ impl System { pub fn get_storage_devices() -> IoResult> { os::get_storage_devices() } -} \ No newline at end of file +} diff --git a/src/storage/windows.rs b/src/storage/windows.rs index d67a763..2e5e98b 100644 --- a/src/storage/windows.rs +++ b/src/storage/windows.rs @@ -7,12 +7,9 @@ use winapi::shared::minwindef::FALSE; use winapi::um::fileapi::{GetDiskFreeSpaceExW, GetLogicalDriveStringsW, GetVolumeInformationW}; use winapi::um::winnt::ULARGE_INTEGER; -pub struct StorageDeviceEnumerator { -} - impl System { - pub fn system_drives() -> Box { - Box::new(StorageDeviceEnumerator{}) + pub fn get_storage_devices() -> IoResult> { + Ok(vec![DeviceRef{}]) } } @@ -69,10 +66,3 @@ impl StorageRef for DeviceRef { } } -impl StorageEnumerator for StorageDeviceEnumerator { - type Ref = DeviceRef; - - fn list(&self) -> IoResult> { - Ok(Vec::default()) - } -}