use std::ffi::{OsString, c_void};
use std::io;
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use windows_sys::Win32::Foundation::{CloseHandle, HANDLE, INVALID_HANDLE_VALUE};
use windows_sys::Win32::Storage::FileSystem::{
CreateFileW, FILE_BEGIN, FILE_FLAG_BACKUP_SEMANTICS, FILE_FLAG_NO_BUFFERING,
FILE_FLAG_SEQUENTIAL_SCAN, FILE_READ_DATA, FILE_SHARE_READ, FILE_SHARE_WRITE,
GetVolumeInformationW, GetVolumePathNameW, OPEN_EXISTING, ReadFile, SetFilePointerEx,
};
use windows_sys::Win32::System::IO::DeviceIoControl;
use windows_sys::Win32::System::Ioctl::{
FSCTL_GET_NTFS_VOLUME_DATA, FSCTL_GET_RETRIEVAL_POINTERS, NTFS_VOLUME_DATA_BUFFER,
};
use crate::index::{Node, ScanProgress, Tree};
use super::{Record, RecordError, apply_fixups, build_tree, merge_extension_records, parse_record};
const ERROR_ACCESS_DENIED: i32 = 5;
const ERROR_MORE_DATA: i32 = 234;
const CHUNK: usize = 16 * 1024 * 1024;
pub const DEFAULT_READERS: usize = 8;
pub enum Capability {
Ready(Box<Volume>),
NeedsElevation {
letter: char,
mft_bytes: u64,
},
NotNtfs {
filesystem: String,
},
Unsupported(String),
}
#[derive(Debug, Clone, Copy)]
struct Extent {
lcn: i64,
clusters: i64,
}
pub struct Volume {
device: String,
mount: PathBuf,
bytes_per_sector: usize,
bytes_per_cluster: u64,
record_size: usize,
mft_valid_len: u64,
extents: Vec<Extent>,
}
impl Volume {
pub fn mount(&self) -> &Path {
&self.mount
}
pub fn record_slots(&self) -> u64 {
self.mft_valid_len / self.record_size as u64
}
pub fn mft_bytes(&self) -> u64 {
self.mft_valid_len
}
pub fn fragments(&self) -> usize {
self.extents.len()
}
}
fn wide(s: &str) -> Vec<u16> {
OsString::from(s)
.encode_wide()
.chain(std::iter::once(0))
.collect()
}
fn last_err() -> i32 {
io::Error::last_os_error().raw_os_error().unwrap_or(-1)
}
fn open(path: &str, access: u32, flags: u32) -> Result<Handle, i32> {
let w = wide(path);
let h = unsafe {
CreateFileW(
w.as_ptr(),
access,
FILE_SHARE_READ | FILE_SHARE_WRITE,
std::ptr::null(),
OPEN_EXISTING,
flags,
std::ptr::null_mut(),
)
};
if h == INVALID_HANDLE_VALUE || h.is_null() {
Err(last_err())
} else {
Ok(Handle(h))
}
}
struct Handle(HANDLE);
impl Drop for Handle {
fn drop(&mut self) {
unsafe { CloseHandle(self.0) };
}
}
struct AlignedBuf {
ptr: *mut u8,
len: usize,
layout: std::alloc::Layout,
}
impl AlignedBuf {
fn new(len: usize, align: usize) -> Self {
let layout = std::alloc::Layout::from_size_align(len, align.max(8)).expect("valid layout");
let ptr = unsafe { std::alloc::alloc(layout) };
assert!(!ptr.is_null(), "failed to allocate {len} bytes");
AlignedBuf { ptr, len, layout }
}
fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
}
}
impl Drop for AlignedBuf {
fn drop(&mut self) {
unsafe { std::alloc::dealloc(self.ptr, self.layout) }
}
}
fn mount_point(path: &Path) -> Result<PathBuf, i32> {
let w: Vec<u16> = path
.as_os_str()
.encode_wide()
.chain(std::iter::once(0))
.collect();
let mut out = vec![0u16; 260];
let ok = unsafe { GetVolumePathNameW(w.as_ptr(), out.as_mut_ptr(), out.len() as u32) };
if ok == 0 {
return Err(last_err());
}
let n = out.iter().position(|&c| c == 0).unwrap_or(out.len());
Ok(PathBuf::from(OsString::from_wide(&out[..n])))
}
fn filesystem_name(mount: &Path) -> Result<String, i32> {
let w: Vec<u16> = mount
.as_os_str()
.encode_wide()
.chain(std::iter::once(0))
.collect();
let mut fs = vec![0u16; 32];
let ok = unsafe {
GetVolumeInformationW(
w.as_ptr(),
std::ptr::null_mut(),
0,
std::ptr::null_mut(),
std::ptr::null_mut(),
std::ptr::null_mut(),
fs.as_mut_ptr(),
fs.len() as u32,
)
};
if ok == 0 {
return Err(last_err());
}
let n = fs.iter().position(|&c| c == 0).unwrap_or(fs.len());
Ok(String::from_utf16_lossy(&fs[..n]))
}
fn mft_extents(h: &Handle) -> Result<Vec<Extent>, i32> {
let mut out = Vec::new();
let mut start_vcn: i64 = 0;
let mut buf = vec![0u8; 64 * 1024];
loop {
let input = start_vcn;
let mut returned = 0u32;
let ok = unsafe {
DeviceIoControl(
h.0,
FSCTL_GET_RETRIEVAL_POINTERS,
&input as *const i64 as *const c_void,
size_of::<i64>() as u32,
buf.as_mut_ptr() as *mut c_void,
buf.len() as u32,
&mut returned,
std::ptr::null_mut(),
)
};
let err = if ok == 0 { last_err() } else { 0 };
if ok == 0 && err != ERROR_MORE_DATA {
return Err(err);
}
let count = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize;
let mut vcn = i64::from_le_bytes(buf[8..16].try_into().unwrap());
let before = vcn;
for i in 0..count {
let base = 16 + i * 16;
if base + 16 > buf.len() {
break;
}
let next_vcn = i64::from_le_bytes(buf[base..base + 8].try_into().unwrap());
let lcn = i64::from_le_bytes(buf[base + 8..base + 16].try_into().unwrap());
let clusters = next_vcn - vcn;
if lcn >= 0 && clusters > 0 {
out.push(Extent { lcn, clusters });
}
vcn = next_vcn;
}
if ok != 0 || vcn == before {
return Ok(out);
}
start_vcn = vcn;
}
}
pub fn probe(path: &Path) -> Capability {
let mount = match mount_point(path) {
Ok(m) => m,
Err(e) => return Capability::Unsupported(format!("cannot resolve volume (os error {e})")),
};
match filesystem_name(&mount) {
Ok(fs) if fs == "NTFS" => {}
Ok(fs) => return Capability::NotNtfs { filesystem: fs },
Err(e) => {
return Capability::Unsupported(format!("cannot identify filesystem (os error {e})"));
}
}
let letter = mount
.to_string_lossy()
.chars()
.next()
.unwrap_or('C')
.to_ascii_uppercase();
let root_h = match open(&mount.to_string_lossy(), 0, FILE_FLAG_BACKUP_SEMANTICS) {
Ok(h) => h,
Err(e) => {
return Capability::Unsupported(format!("cannot open volume root (os error {e})"));
}
};
let mut vd = NTFS_VOLUME_DATA_BUFFER::default();
let mut returned = 0u32;
let ok = unsafe {
DeviceIoControl(
root_h.0,
FSCTL_GET_NTFS_VOLUME_DATA,
std::ptr::null(),
0,
&mut vd as *mut _ as *mut c_void,
size_of::<NTFS_VOLUME_DATA_BUFFER>() as u32,
&mut returned,
std::ptr::null_mut(),
)
};
drop(root_h);
if ok == 0 {
return Capability::Unsupported(format!(
"volume data query failed (os error {})",
last_err()
));
}
let record_size = vd.BytesPerFileRecordSegment as usize;
let bytes_per_sector = vd.BytesPerSector as usize;
let bytes_per_cluster = vd.BytesPerCluster as u64;
let mft_valid_len = vd.MftValidDataLength.max(0) as u64;
if record_size == 0 || !record_size.is_power_of_two() || bytes_per_sector == 0 {
return Capability::Unsupported(format!("implausible geometry: record {record_size}"));
}
if mft_valid_len / record_size as u64 > u32::MAX as u64 {
return Capability::Unsupported("MFT has more slots than a u32 can index".to_owned());
}
let device = format!("\\\\.\\{letter}:");
let vol_h = match open(
&device,
FILE_READ_DATA,
FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN,
) {
Ok(h) => h,
Err(ERROR_ACCESS_DENIED) => {
return Capability::NeedsElevation {
letter,
mft_bytes: mft_valid_len,
};
}
Err(e) => return Capability::Unsupported(format!("cannot open {device} (os error {e})")),
};
drop(vol_h);
let mft_path = format!("{letter}:\\$MFT");
let extents = match open(&mft_path, 0, 0) {
Ok(h) => match mft_extents(&h) {
Ok(e) if !e.is_empty() => e,
Ok(_) => return Capability::Unsupported("$MFT reported no extents".to_owned()),
Err(e) => {
return Capability::Unsupported(format!("cannot map $MFT (os error {e})"));
}
},
Err(e) => return Capability::Unsupported(format!("cannot open $MFT (os error {e})")),
};
let mapped: u64 = extents.iter().map(|e| e.clusters as u64).sum::<u64>() * bytes_per_cluster;
if mapped < mft_valid_len {
return Capability::Unsupported(format!(
"$MFT extents cover {mapped} of {mft_valid_len} bytes"
));
}
Capability::Ready(Box::new(Volume {
device,
mount,
bytes_per_sector,
bytes_per_cluster,
record_size,
mft_valid_len,
extents,
}))
}
#[derive(Clone, Copy)]
struct Job {
device_offset: i64,
len: usize,
first_record: u32,
}
fn plan_jobs(vol: &Volume) -> Vec<Job> {
let mut jobs = Vec::new();
let mut mft_off: u64 = 0;
for e in &vol.extents {
if mft_off >= vol.mft_valid_len {
break;
}
let extent_bytes = e.clusters as u64 * vol.bytes_per_cluster;
let usable = extent_bytes.min(vol.mft_valid_len - mft_off);
let mut done: u64 = 0;
while done < usable {
let len = CHUNK.min((usable - done) as usize);
jobs.push(Job {
device_offset: (e.lcn as u64 * vol.bytes_per_cluster + done) as i64,
len,
first_record: ((mft_off + done) / vol.record_size as u64) as u32,
});
done += len as u64;
}
mft_off += extent_bytes;
}
jobs
}
struct JobTiming {
read: u128,
parse: u128,
}
fn run_job(
h: &Handle,
buf: &mut AlignedBuf,
job: Job,
vol: &Volume,
out: &mut Vec<(u32, Record)>,
torn: &mut u64,
timing: &mut JobTiming,
) -> Result<(), i32> {
let read_start = Instant::now();
let mut pos = 0i64;
if unsafe { SetFilePointerEx(h.0, job.device_offset, &mut pos, FILE_BEGIN) } == 0 {
return Err(last_err());
}
let want = job.len;
let mut filled = 0usize;
while filled < want {
let mut got = 0u32;
let ok = unsafe {
ReadFile(
h.0,
buf.as_mut_slice()[filled..].as_mut_ptr(),
(want - filled) as u32,
&mut got,
std::ptr::null_mut(),
)
};
if ok == 0 {
return Err(last_err());
}
if got == 0 {
break;
}
filled += got as usize;
}
timing.read += read_start.elapsed().as_nanos();
let parse_start = Instant::now();
let rec_size = vol.record_size;
let data = &mut buf.as_mut_slice()[..filled];
for (i, rec) in data.chunks_exact_mut(rec_size).enumerate() {
let number = job.first_record + i as u32;
match apply_fixups(rec, vol.bytes_per_sector) {
Ok(()) => {}
Err(RecordError::TornRead) => {
*torn += 1;
continue;
}
Err(_) => continue,
}
if let Ok(r) = parse_record(rec, number) {
out.push((number, r));
}
}
timing.parse += parse_start.elapsed().as_nanos();
Ok(())
}
pub fn scan(vol: &Volume, progress: &ScanProgress, readers: Option<usize>) -> io::Result<Tree> {
let start = Instant::now();
let jobs = plan_jobs(vol);
let n_readers = readers
.unwrap_or(DEFAULT_READERS)
.clamp(1, 64)
.min(jobs.len().max(1));
let slots = (vol.mft_valid_len / vol.record_size as u64) as usize + 1;
if let Ok(mut cur) = progress.current.try_lock() {
*cur = format!("reading $MFT ({} fragments)", vol.extents.len());
}
let next = AtomicUsize::new(0);
let torn_total = AtomicUsize::new(0);
let failed = AtomicUsize::new(0);
let read_nanos = AtomicUsize::new(0);
let parse_nanos = AtomicUsize::new(0);
let collected: Vec<Vec<(u32, Record)>> = std::thread::scope(|s| {
let handles: Vec<_> = (0..n_readers)
.map(|_| {
s.spawn(|| {
let mut mine: Vec<(u32, Record)> = Vec::new();
let Ok(h) = open(
&vol.device,
FILE_READ_DATA,
FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN,
) else {
failed.fetch_add(1, Ordering::Relaxed);
return mine;
};
let mut buf = AlignedBuf::new(CHUNK, vol.bytes_per_sector.max(4096));
let mut torn = 0u64;
let mut timing = JobTiming { read: 0, parse: 0 };
loop {
let i = next.fetch_add(1, Ordering::Relaxed);
let Some(&job) = jobs.get(i) else { break };
if run_job(&h, &mut buf, job, vol, &mut mine, &mut torn, &mut timing)
.is_err()
{
failed.fetch_add(1, Ordering::Relaxed);
}
progress.bytes.fetch_add(job.len as u64, Ordering::Relaxed);
progress
.dirs
.fetch_add((job.len / vol.record_size) as u64, Ordering::Relaxed);
}
torn_total.fetch_add(torn as usize, Ordering::Relaxed);
read_nanos.fetch_add(timing.read as usize, Ordering::Relaxed);
parse_nanos.fetch_add(timing.parse as usize, Ordering::Relaxed);
mine
})
})
.collect();
handles
.into_iter()
.map(|h| h.join().unwrap_or_default())
.collect()
});
let mut records: Vec<Option<Record>> = vec![None; slots];
for batch in collected {
for (n, r) in batch {
if let Some(slot) = records.get_mut(n as usize) {
*slot = Some(r);
}
}
}
let merged = merge_extension_records(&mut records);
let nodes = build_tree(&records);
drop(records);
let n_files: u64 = nodes.iter().filter(|n| !n.is_dir).count() as u64;
let tree = to_tree(
nodes,
vol.mount.clone(),
torn_total.load(Ordering::Relaxed) as u64 + failed.load(Ordering::Relaxed) as u64,
start.elapsed().as_secs_f64(),
);
progress.files.store(n_files, Ordering::Relaxed);
if let Ok(mut cur) = progress.current.try_lock() {
*cur = format!("merged {merged} extension records");
}
if std::env::var_os("DISKTREE_MFT_TIMING").is_some() {
let wall = start.elapsed().as_secs_f64();
let read = read_nanos.load(Ordering::Relaxed) as f64 / 1e9;
let parse = parse_nanos.load(Ordering::Relaxed) as f64 / 1e9;
eprintln!(
" [mft] wall {wall:.3}s | readers {n_readers} | summed read {read:.3}s \
parse {parse:.3}s | overlap {:.2}x | read-only throughput {:.0} MB/s | \
merged {merged} | jobs {}",
(read + parse) / wall.max(1e-9),
vol.mft_valid_len as f64 / 1e6 / read.max(1e-9) * n_readers as f64,
jobs.len(),
);
}
Ok(tree)
}
fn to_tree(nodes: Vec<super::TreeNode>, mount: PathBuf, errors: u64, secs: f64) -> Tree {
let mut out: Vec<Node> = Vec::with_capacity(nodes.len());
for (i, n) in nodes.iter().enumerate() {
out.push(Node {
name: if i == 0 {
OsString::from(mount.display().to_string())
} else {
n.name.clone()
},
is_dir: n.is_dir,
size: n.size,
n_files: u64::from(!n.is_dir),
parent: if i == 0 { None } else { Some(n.parent) },
children: Vec::new(),
deleted: false,
});
}
for i in (1..out.len()).rev() {
let p = nodes[i].parent;
let (size, files) = (out[i].size, out[i].n_files);
out[p].size += size;
out[p].n_files += files;
out[p].children.push(i);
}
for n in out.iter_mut() {
n.children.reverse();
}
Tree {
nodes: out,
root: 0,
root_path: mount,
errors,
scan_secs: secs,
}
}
pub fn calibrate_readers(file: &Path, levels: &[usize]) -> io::Result<Vec<(usize, f64)>> {
let path = file.to_string_lossy().into_owned();
let total = std::fs::metadata(file)?.len() as usize;
let usable = total.min(512 * 1024 * 1024) / CHUNK * CHUNK;
if usable < CHUNK {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("need at least {CHUNK} bytes to calibrate, got {total}"),
));
}
let mut results = Vec::new();
for &readers in levels {
let chunks = usable / CHUNK;
if chunks < readers {
continue;
}
let per = chunks / readers;
let next = &AtomicUsize::new(0);
let start = Instant::now();
let read: usize = std::thread::scope(|s| {
let hs: Vec<_> = (0..readers)
.map(|_| {
let path = path.clone();
s.spawn(move || {
let Ok(h) = open(
&path,
FILE_READ_DATA,
FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN,
) else {
return 0usize;
};
let mut buf = AlignedBuf::new(CHUNK, 4096);
let mut got_total = 0usize;
loop {
let i = next.fetch_add(1, Ordering::Relaxed);
if i >= per * readers {
break;
}
let mut pos = 0i64;
unsafe {
SetFilePointerEx(h.0, (i * CHUNK) as i64, &mut pos, FILE_BEGIN)
};
let mut got = 0u32;
let ok = unsafe {
ReadFile(
h.0,
buf.as_mut_slice().as_mut_ptr(),
CHUNK as u32,
&mut got,
std::ptr::null_mut(),
)
};
if ok == 0 || got == 0 {
break;
}
got_total += got as usize;
}
got_total
})
})
.collect();
hs.into_iter().map(|h| h.join().unwrap_or(0)).sum()
});
let secs = start.elapsed().as_secs_f64().max(1e-9);
results.push((readers, read as f64 / secs));
}
Ok(results)
}