Upload FastSeekWpf/Core/MftReader.cs

FastSeekWpf/Core/MftReader.cs

@@ -9,20 +9,19 @@ namespace FastSeekWpf.Core;
 
 public class MftReader : IDisposable
 {
+    // Volume handle — used for boot sector read and FSCTL fallback
     private readonly IntPtr _volumeHandle;
     private readonly NtfsDrive _drive;
     private bool _disposed;
 
+    private const int FallbackBuf = 4 * 1024 * 1024;
+    private const int DirectBuf = 4 * 1024 * 1024;
+
     public NtfsDrive Drive => _drive;
 
     public MftReader(NtfsDrive drive)
     {
         _drive = drive;
-
-        // FILE_FLAG_BACKUP_SEMANTICS is REQUIRED for opening a volume device
-        // for direct MFT reads. FILE_FLAG_NO_BUFFERING is removed because .NET
-        // arrays are not guaranteed sector-aligned, and unaligned reads with
-        // NO_BUFFERING hang forever or return ERROR_INVALID_PARAMETER.
         _volumeHandle = Win32Api.CreateFileW(
             drive.DevicePath,
             Win32Api.GENERIC_READ,
@@ -37,134 +36,92 @@ public class MftReader : IDisposable
             int err = Marshal.GetLastWin32Error();
             throw new IOException($"Failed to open volume {drive.DevicePath}: {Win32Error(err)}");
         }
-
-        Logger.Log($"Volume {drive.DevicePath} opened OK, handle=0x{_volumeHandle:X}");
     }
 
+    // ------------------------------------------------------------------
+    //  Primary: direct $MFT file read  (falls back to FSCTL if fails)
+    // ------------------------------------------------------------------
+
     /// <summary>
-    /// Direct MFT read matching the Rust implementation.
+    /// Read $MFT file directly for maximum speed. Returns null if unavailable.
+    /// Matches Rust scan_direct() exactly.
     /// </summary>
-    public ScanResult ScanDirect()
+    public ScanResult? ScanDirect()
     {
-        Logger.Log($"[{Drive.Letter}] Reading boot sector...");
-
-        // ── 1. Boot sector ──
-        byte[] boot = new byte[512];
-        bool ok = Win32Api.ReadFile(_volumeHandle, boot, 512, out uint br, IntPtr.Zero);
-        if (!ok || br < 512)
-        {
-            int err = Marshal.GetLastWin32Error();
-            throw new IOException($"Failed to read boot sector: {Win32Error(err)}");
-        }
+        int? recordSize = ReadMftRecordSize();
+        if (!recordSize.HasValue) return null;
 
-        // Validate NTFS signature at offset 3: "NTFS    "
-        if (boot[3] != 'N' || boot[4] != 'T' || boot[5] != 'F' || boot[6] != 'S')
-            throw new IOException("Not an NTFS volume (missing NTFS signature)");
-
-        ushort bytesPerSector = BinaryPrimitives.ReadUInt16LittleEndian(boot.AsSpan(0x0B, 2));
-        byte sectorsPerCluster = boot[0x0D];
-        int clusterSize = bytesPerSector * sectorsPerCluster;
-
-        long mftStartLcn = BinaryPrimitives.ReadInt64LittleEndian(boot.AsSpan(0x30, 8));
-        if (mftStartLcn <= 0)
-            throw new IOException($"Invalid MFT start LCN: {mftStartLcn}");
+        string mftPath = $"{_drive.Root}$MFT";
 
-        sbyte raw = (sbyte)boot[0x40];
-        int recordSize = raw < 0 ? 1 << (-raw) : raw * clusterSize;
-
-        long mftByteOffset = mftStartLcn * clusterSize;
-        Logger.Log($"[{Drive.Letter}] Boot: bps={bytesPerSector}, spc={sectorsPerCluster}, clusterSize={clusterSize}, mftLcn={mftStartLcn}, recordSize={recordSize}, offset={mftByteOffset}");
-
-        // ── 2. Seek to MFT ──
-        ok = Win32Api.SetFilePointerEx(_volumeHandle, mftByteOffset, out _, 0);
-        if (!ok)
-            throw new IOException($"Seek to MFT offset {mftByteOffset} failed: {Win32Error(Marshal.GetLastWin32Error())}");
+        // Open the $MFT file as a normal file (not the volume device)
+        IntPtr mftHandle = Win32Api.CreateFileW(
+            mftPath,
+            Win32Api.GENERIC_READ,
+            Win32Api.FILE_SHARE_READ | Win32Api.FILE_SHARE_WRITE | Win32Api.FILE_SHARE_DELETE,
+            IntPtr.Zero,
+            Win32Api.OPEN_EXISTING,
+            Win32Api.FILE_FLAG_BACKUP_SEMANTICS | Win32Api.FILE_FLAG_SEQUENTIAL_SCAN,
+            IntPtr.Zero);
 
-        // ── 3. Read first record ($MFT itself) to get data size ──
-        Logger.Log($"[{Drive.Letter}] Reading first MFT record to find $DATA size...");
-        byte[] firstRecord = new byte[recordSize];
-        ok = Win32Api.ReadFile(_volumeHandle, firstRecord, (uint)recordSize, out br, IntPtr.Zero);
-        if (!ok || br < recordSize)
-            throw new IOException($"Failed to read first MFT record: {Win32Error(Marshal.GetLastWin32Error())}");
+        if (mftHandle == new IntPtr(-1))
+            return null;
 
-        long mftDataSize = ParseMftDataSize(firstRecord);
-        if (mftDataSize <= 0)
-        {
-            Logger.Log($"[{Drive.Letter}] WARNING: Could not parse MFT $DATA size; using 2GB fallback.");
-            mftDataSize = 2L * 1024 * 1024 * 1024;
-        }
-        else
+        try
         {
-            Logger.Log($"[{Drive.Letter}] MFT data size: {mftDataSize:N0} bytes ({mftDataSize / 1024.0 / 1024.0:F1} MB)");
-        }
-
-        // ── 4. Seek back to MFT start and read all records ──
-        ok = Win32Api.SetFilePointerEx(_volumeHandle, mftByteOffset, out _, 0);
-        if (!ok)
-            throw new IOException($"Seek back to MFT start failed: {Win32Error(Marshal.GetLastWin32Error())}");
+            var records = new List<CompactRecord>();
+            var nameData = new List<char>();
+            byte[] buffer = new byte[DirectBuf];
+            ulong mftIndex = 0;
+            int leftover = 0;
 
-        var records = new List<CompactRecord>();
-        var nameData = new List<char>();
-        byte[] buffer = new byte[recordSize]; // read one record at a time for safety
-        ulong mftIndex = 0;
-        long bytesRemaining = mftDataSize;
-        int consecutiveInvalid = 0;
-        const int MaxConsecutiveInvalid = 50;
+            while (true)
+            {
+                bool ok = Win32Api.ReadFile(
+                    mftHandle, buffer, (uint)(buffer.Length - leftover), out uint bytesRead, IntPtr.Zero);
+                if (!ok || bytesRead == 0)
+                    break;
 
-        Logger.Log($"[{Drive.Letter}] Starting MFT read loop...");
+                int total = leftover + (int)bytesRead;
+                int offset = 0;
 
-        while (bytesRemaining > 0 && consecutiveInvalid < MaxConsecutiveInvalid)
-        {
-            ok = Win32Api.ReadFile(_volumeHandle, buffer, (uint)recordSize, out br, IntPtr.Zero);
-            if (!ok)
-            {
-                int err = Marshal.GetLastWin32Error();
-                if (err == 38) // ERROR_HANDLE_EOF
+                while (offset + recordSize.Value <= total)
                 {
-                    Logger.Log($"[{Drive.Letter}] MFT read: EOF reached");
-                    break;
+                    bool applied = ApplyFixup(buffer, offset, recordSize.Value);
+                    if (applied)
+                    {
+                        ParseFileRecord(buffer.AsSpan(offset, recordSize.Value), mftIndex, records, nameData);
+                    }
+
+                    mftIndex++;
+                    offset += recordSize.Value;
                 }
-                Logger.Log($"[{Drive.Letter}] ReadFile failed at offset {mftByteOffset + (mftDataSize - bytesRemaining)}: {Win32Error(err)} — stopping");
-                break;
-            }
-            if (br == 0)
-            {
-                Logger.Log($"[{Drive.Letter}] ReadFile returned 0 bytes — stopping");
-                break;
-            }
 
-            bytesRemaining -= br;
+                // Align offset down to record boundary
+                offset = total - (total % recordSize.Value);
 
-            if (ApplyFixup(buffer, recordSize))
-            {
-                ParseFileRecord(buffer.AsSpan(), mftIndex, records, nameData);
-                consecutiveInvalid = 0;
-            }
-            else
-            {
-                consecutiveInvalid++;
+                leftover = total - offset;
+                if (leftover > 0)
+                    Array.Copy(buffer, offset, buffer, 0, leftover);
             }
 
-            mftIndex++;
-
-            if (mftIndex % 10000 == 0)
-                Logger.Log($"[{Drive.Letter}] Progress: {mftIndex:N0} records, {records.Count:N0} valid, {bytesRemaining:N0} bytes remaining");
+            Logger.Log($"[{Drive.Letter}] Direct scan: {records.Count} records from {mftIndex} positions");
+            return new ScanResult(records, nameData);
+        }
+        finally
+        {
+            Win32Api.CloseHandle(mftHandle);
         }
-
-        Logger.Log($"[{Drive.Letter}] MFT scan complete: {records.Count:N0} records from {mftIndex:N0} positions");
-        return new ScanResult(records, nameData);
     }
 
-    /// <summary>
-    /// Fallback FSCTL_ENUM_USN_DATA scan.
-    /// </summary>
+    // ------------------------------------------------------------------
+    //  Fallback: FSCTL_ENUM_USN_DATA (4 MB buffer)
+    // ------------------------------------------------------------------
+
     public ScanResult Scan()
     {
-        Logger.Log($"[{Drive.Letter}] Starting FSCTL fallback scan...");
-
         var records = new List<CompactRecord>();
         var nameData = new List<char>();
-        byte[] buffer = new byte[4 * 1024 * 1024];
+        byte[] buffer = new byte[FallbackBuf];
 
         var enumData = new MFT_ENUM_DATA_V0
         {
@@ -203,18 +160,17 @@ public class MftReader : IDisposable
                         }
                         if (error == 122) // ERROR_INSUFFICIENT_BUFFER
                         {
-                            Logger.Log($"[{Drive.Letter}] FSCTL_ENUM_USN_DATA: insufficient buffer, retrying with larger buffer");
+                            Logger.Log($"[{Drive.Letter}] FSCTL_ENUM_USN_DATA: buffer too small, retrying");
                             buffer = new byte[buffer.Length * 2];
                             continue;
                         }
-                        throw new IOException(
-                            $"FSCTL_ENUM_USN_DATA failed (iter {iterations}): {Win32Error(error)}. " +
-                            $"StartFileReferenceNumber={enumData.StartFileReferenceNumber}");
+                        Logger.Log($"[{Drive.Letter}] FSCTL_ENUM_USN_DATA failed: {Win32Error(error)}");
+                        break;
                     }
 
                     if (bytesReturned <= 8)
                     {
-                        Logger.Log($"[{Drive.Letter}] FSCTL_ENUM_USN_DATA: {bytesReturned} bytes returned, stopping");
+                        Logger.Log($"[{Drive.Letter}] FSCTL: {bytesReturned} bytes, stopping");
                         break;
                     }
 
@@ -228,7 +184,8 @@ public class MftReader : IDisposable
                         IntPtr recordPtr = IntPtr.Add(bufferPtr, offset);
                         var record = Marshal.PtrToStructure<USN_RECORD_V2>(recordPtr);
 
-                        if (record.RecordLength == 0) break;
+                        int recLen = (int)record.RecordLength;
+                        if (recLen == 0 || offset + recLen > (int)bytesReturned) break;
 
                         int nameOffset = offset + (int)record.FileNameOffset;
                         int nameLen = record.FileNameLength / 2;
@@ -249,7 +206,7 @@ public class MftReader : IDisposable
                             IsDir = (record.FileAttributes & 0x10) != 0
                         });
 
-                        offset += (int)record.RecordLength;
+                        offset += recLen;
                     }
                 }
                 finally
@@ -263,72 +220,87 @@ public class MftReader : IDisposable
             Marshal.FreeHGlobal(enumDataPtr);
         }
 
-        Logger.Log($"[{Drive.Letter}] FSCTL scan complete: {records.Count} records");
+        Logger.Log($"[{Drive.Letter}] FSCTL scan: {records.Count} records");
         return new ScanResult(records, nameData);
     }
 
-    // ── Helpers ──
+    // ------------------------------------------------------------------
+    //  NTFS helpers
+    // ------------------------------------------------------------------
 
-    private static long ParseMftDataSize(ReadOnlySpan<byte> record)
+    /// <summary>
+    /// Read MFT record size from NTFS boot sector via volume handle.
+    /// Matches Rust read_mft_record_size().
+    /// </summary>
+    private int? ReadMftRecordSize()
     {
-        if (record.Length < 48 || record[0] != 'F' || record[1] != 'I' || record[2] != 'L' || record[3] != 'E')
-            return -1;
-
-        if (!ApplyFixup(record, record.Length))
-            return -1;
+        // Seek to beginning of volume
+        Win32Api.SetFilePointerEx(_volumeHandle, 0, out _, 0);
 
-        ushort firstAttr = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(0x14, 2));
-        int aoff = firstAttr;
-
-        while (aoff + 0x40 <= record.Length)
-        {
-            uint atype = BinaryPrimitives.ReadUInt32LittleEndian(record.Slice(aoff, 4));
-            if (atype == 0xFFFFFFFF) break;
+        byte[] boot = new byte[512];
+        bool ok = Win32Api.ReadFile(_volumeHandle, boot, 512, out uint br, IntPtr.Zero);
+        if (!ok || br < 512)
+            return null;
 
-            uint alen = BinaryPrimitives.ReadUInt32LittleEndian(record.Slice(aoff + 4, 4));
-            if (alen == 0 || aoff + alen > record.Length) break;
+        if (boot[3] != 'N' || boot[4] != 'T' || boot[5] != 'F' || boot[6] != 'S')
+            return null;
 
-            // 0x80 = $DATA, non-resident flag = 1
-            if (atype == 0x80 && record[aoff + 8] == 1)
-            {
-                // Non-resident header: DataSize (RealSize) is at offset 0x30 from attr start
-                long dataSize = BinaryPrimitives.ReadInt64LittleEndian(record.Slice(aoff + 0x30, 8));
-                return dataSize;
-            }
+        ushort bytesPerSector = BinaryPrimitives.ReadUInt16LittleEndian(boot.AsSpan(0x0B, 2));
+        byte sectorsPerCluster = boot[0x0D];
+        int clusterSize = bytesPerSector * sectorsPerCluster;
 
-            aoff += (int)alen;
-        }
+        sbyte raw = (sbyte)boot[0x40];
+        int recordSize = raw > 0 ? raw * clusterSize : 1 << (-raw);
 
-        return -1;
+        Logger.Log($"[{Drive.Letter}] Boot: bps={bytesPerSector}, spc={sectorsPerCluster}, clusterSize={clusterSize}, recordSize={recordSize}");
+        return recordSize;
     }
 
-    private static bool ApplyFixup(ReadOnlySpan<byte> record, int recordSize)
+    /// <summary>
+    /// Apply NTFS multi-sector fixup. Returns false if record is invalid.
+    /// CRITICAL: Restores the original bytes from the fixup array back to sector ends.
+    /// Matches Rust apply_fixup() exactly.
+    /// </summary>
+    private static bool ApplyFixup(byte[] record, int offset, int recordSize)
     {
-        if (record.Length < 48 || record[0] != 'F' || record[1] != 'I' || record[2] != 'L' || record[3] != 'E')
+        if (recordSize < 48 || record[offset] != 'F' || record[offset + 1] != 'I'
+            || record[offset + 2] != 'L' || record[offset + 3] != 'E')
             return false;
 
-        ushort fixupOff = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(4, 2));
-        ushort fixupCnt = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(6, 2));
+        ushort fixupOff = BinaryPrimitives.ReadUInt16LittleEndian(record.AsSpan(offset + 4, 2));
+        ushort fixupCnt = BinaryPrimitives.ReadUInt16LittleEndian(record.AsSpan(offset + 6, 2));
 
         if (fixupCnt < 2 || fixupOff + fixupCnt * 2 > recordSize)
             return false;
 
-        byte check0 = record[fixupOff];
-        byte check1 = record[fixupOff + 1];
+        byte check0 = record[offset + fixupOff];
+        byte check1 = record[offset + fixupOff + 1];
 
         for (int i = 1; i < fixupCnt; i++)
         {
             int end = i * 512 - 2;
             if (end + 1 >= recordSize) break;
-            if (record[end] != check0 || record[end + 1] != check1)
+
+            if (record[offset + end] != check0 || record[offset + end + 1] != check1)
                 return false;
+
+            // RESTORE the real bytes from the fixup array
+            record[offset + end] = record[offset + fixupOff + i * 2];
+            record[offset + end + 1] = record[offset + fixupOff + i * 2 + 1];
         }
 
         return true;
     }
 
-    private static void ParseFileRecord(ReadOnlySpan<byte> record, ulong mftIndex,
-        List<CompactRecord> records, List<char> nameData)
+    /// <summary>
+    /// Parse one MFT FILE record, extracting name + parent into the arena.
+    /// Matches Rust parse_file_record() exactly.
+    /// </summary>
+    private static void ParseFileRecord(
+        ReadOnlySpan<byte> record,
+        ulong mftIndex,
+        List<CompactRecord> records,
+        List<char> nameData)
     {
         ushort flags = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(0x16, 2));
         if ((flags & 0x01) == 0) return; // Record not in use
@@ -365,7 +337,7 @@ public class MftReader : IDisposable
 
                     if (vs + 66 + nlen * 2 <= record.Length)
                     {
-                        if (ns == 2) // DOS 8.3 name, skip
+                        if (ns == 2) // DOS 8.3 name — skip
                         {
                             aoff += (int)alen;
                             continue;
@@ -373,9 +345,9 @@ public class MftReader : IDisposable
 
                         byte priority = ns switch
                         {
-                            1 => 0, // POSIX
-                            3 => 1, // Win32 & DOS
-                            0 => 2, // Win32
+                            1 => 0, // Win32
+                            3 => 1, // Win32 + DOS
+                            0 => 2, // POSIX
                             _ => 3,
                         };