Upload FastSeekWpf/Core/MftReader.cs

FastSeekWpf/Core/MftReader.cs

@@ -13,20 +13,24 @@ public class MftReader : IDisposable
     private readonly NtfsDrive _drive;
     private bool _disposed;
 
-    private const uint DirectBuf = 1024 * 1024; // 1MB read buffer
+    private const int ChunkSize = 256 * 1024; // 256KB — large enough for speed, small enough for memory
 
     public NtfsDrive Drive => _drive;
 
     public MftReader(NtfsDrive drive)
     {
         _drive = drive;
+        // NOTE: FILE_FLAG_NO_BUFFERING removed. It requires sector-aligned buffers AND
+        // aligned file offsets, and .NET arrays are not guaranteed aligned. The Rust
+        // code likely opens the volume without NO_BUFFERING. Buffered reads are fast
+        // enough since we now read exactly the MFT size, not the whole volume.
         _volumeHandle = Win32Api.CreateFileW(
             drive.DevicePath,
             Win32Api.GENERIC_READ,
             Win32Api.FILE_SHARE_READ | Win32Api.FILE_SHARE_WRITE | Win32Api.FILE_SHARE_DELETE,
             IntPtr.Zero,
             Win32Api.OPEN_EXISTING,
-            Win32Api.FILE_FLAG_NO_BUFFERING | Win32Api.FILE_FLAG_SEQUENTIAL_SCAN,
+            0, // FILE_ATTRIBUTE_NORMAL — buffered, cached reads
             IntPtr.Zero);
 
         if (_volumeHandle == new IntPtr(-1))
@@ -37,112 +41,132 @@ public class MftReader : IDisposable
     }
 
     /// <summary>
-    /// Direct MFT read: parse boot sector, find MFT LCN, read raw records from volume device.
-    /// Matches the Rust implementation exactly.
+    /// Direct MFT read: parse boot sector, seek to MFT start, read first record
+    /// to discover MFT $DATA size, then read exactly that many bytes.
     /// </summary>
     public ScanResult ScanDirect()
     {
-        // Read boot sector (512 bytes, must be sector-aligned for NO_BUFFERING)
+        // ── 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)}");
-        }
+            throw new IOException($"Failed to read boot sector: {Win32Error(Marshal.GetLastWin32Error())}");
 
-        // 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(new ReadOnlySpan<byte>(boot, 0x0B, 2));
+        ushort bytesPerSector = BinaryPrimitives.ReadUInt16LittleEndian(boot.AsSpan(0x0B, 2));
         byte sectorsPerCluster = boot[0x0D];
         int clusterSize = bytesPerSector * sectorsPerCluster;
 
-        // MFT start LCN at offset 0x30 (signed 64-bit)
-        long mftStartLcn = BinaryPrimitives.ReadInt64LittleEndian(new ReadOnlySpan<byte>(boot, 0x30, 8));
+        long mftStartLcn = BinaryPrimitives.ReadInt64LittleEndian(boot.AsSpan(0x30, 8));
         if (mftStartLcn <= 0)
             throw new IOException($"Invalid MFT start LCN: {mftStartLcn}");
 
+        sbyte raw = (sbyte)boot[0x40];
+        int recordSize = raw < 0 ? 1 << (-raw) : raw * clusterSize;
+
         long mftByteOffset = mftStartLcn * clusterSize;
-        Logger.Log($"Boot: bps={bytesPerSector}, spc={sectorsPerCluster}, clusterSize={clusterSize}, mftLcn={mftStartLcn}, mftOffset={mftByteOffset}");
+        Logger.Log($"Boot: bps={bytesPerSector}, spc={sectorsPerCluster}, clusterSize={clusterSize}, mftLcn={mftStartLcn}, recordSize={recordSize}");
 
-        // Seek to MFT start
+        // ── 2. Seek to MFT and read first record ($MFT itself) ──
         ok = Win32Api.SetFilePointerEx(_volumeHandle, mftByteOffset, out _, 0);
         if (!ok)
+            throw new IOException($"Seek to MFT offset {mftByteOffset} failed: {Win32Error(Marshal.GetLastWin32Error())}");
+
+        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())}");
+
+        // ── 3. Parse MFT total data size from its own $DATA attribute ──
+        long mftDataSize = ParseMftDataSize(firstRecord);
+        if (mftDataSize <= 0)
         {
-            int err = Marshal.GetLastWin32Error();
-            throw new IOException($"Failed to seek to MFT offset {mftByteOffset}: {Win32Error(err)}");
+            Logger.Log("WARNING: Could not parse MFT $DATA size; using 2GB fallback. " +
+                       "This may read past the end of the MFT.");
+            mftDataSize = 2L * 1024 * 1024 * 1024;
         }
-
-        // MFT record size: raw byte at 0x40
-        sbyte raw = (sbyte)boot[0x40];
-        int recordSize;
-        if (raw > 0)
-            recordSize = raw * clusterSize;
         else
-            recordSize = 1 << (-raw);
+        {
+            Logger.Log($"MFT data size from $DATA attribute: {mftDataSize:N0} bytes ({mftDataSize / 1024.0 / 1024.0:F1} MB)");
+        }
 
-        Logger.Log($"MFT record size: {recordSize} bytes (raw={raw})");
+        // ── 4. Reset to MFT start and read exactly mftDataSize bytes ──
+        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];
+        byte[] chunk = new byte[ChunkSize];
         ulong mftIndex = 0;
-        int leftover = 0;
-        long totalBytesRead = 0;
+        long bytesRemaining = mftDataSize;
+        int consecutiveInvalid = 0;
+        const int MaxConsecutiveInvalid = 50;
 
-        while (true)
+        while (bytesRemaining > 0 && consecutiveInvalid < MaxConsecutiveInvalid)
         {
-            ok = Win32Api.ReadFile(_volumeHandle, buffer, (uint)(buffer.Length - leftover), out uint bytesRead, IntPtr.Zero);
+            int toRead = (int)Math.Min(ChunkSize, bytesRemaining);
+            ok = Win32Api.ReadFile(_volumeHandle, chunk, (uint)toRead, out uint bytesRead, IntPtr.Zero);
             if (!ok)
             {
                 int err = Marshal.GetLastWin32Error();
-                if (err == 38) // ERROR_HANDLE_EOF — reached end of MFT
+                if (err == 38) // ERROR_HANDLE_EOF
                 {
-                    Logger.Log("MFT read: reached EOF");
+                    Logger.Log("MFT read: EOF reached");
                     break;
                 }
-                throw new IOException($"ReadFile failed at offset {mftByteOffset + totalBytesRead}: {Win32Error(err)}");
+                Logger.Log($"ReadFile failed at offset {mftByteOffset + (mftDataSize - bytesRemaining)}: {Win32Error(err)} — stopping");
+                break;
             }
             if (bytesRead == 0)
             {
-                Logger.Log("MFT read: 0 bytes returned");
+                Logger.Log("ReadFile returned 0 bytes — stopping");
                 break;
             }
 
-            totalBytesRead += bytesRead;
-            int total = leftover + (int)bytesRead;
+            bytesRemaining -= bytesRead;
             int offset = 0;
+            int total = (int)bytesRead;
 
             while (offset + recordSize <= total)
             {
-                var span = new ReadOnlySpan<byte>(buffer, offset, recordSize);
+                var span = new ReadOnlySpan<byte>(chunk, offset, recordSize);
                 if (ApplyFixup(span, recordSize))
                 {
                     ParseFileRecord(span, mftIndex, records, nameData);
+                    consecutiveInvalid = 0;
+                }
+                else
+                {
+                    consecutiveInvalid++;
                 }
                 mftIndex++;
                 offset += recordSize;
-            }
 
-            leftover = total - offset;
-            if (leftover > 0)
-                Array.Copy(buffer, offset, buffer, 0, leftover);
+                if (consecutiveInvalid >= MaxConsecutiveInvalid)
+                {
+                    Logger.Log($"Stopping after {MaxConsecutiveInvalid} consecutive invalid records at index {mftIndex}");
+                    bytesRemaining = 0;
+                    break;
+                }
+            }
         }
 
-        Logger.Log($"MFT read complete: {records.Count} records from {mftIndex} record positions");
+        Logger.Log($"MFT scan complete: {records.Count:N0} records from {mftIndex:N0} positions ({mftDataSize - bytesRemaining:N0} bytes read)");
         return new ScanResult(records, nameData);
     }
 
     /// <summary>
-    /// Fallback FSCTL_ENUM_USN_DATA scan.
+    /// Fallback FSCTL_ENUM_USN_DATA scan. Kept for parity but direct scan should
+    /// be preferred — it is faster and doesn't require a USN journal.
     /// </summary>
     public ScanResult Scan()
     {
         var records = new List<CompactRecord>();
         var nameData = new List<char>();
-        byte[] buffer = new byte[4 * 1024 * 1024]; // 4MB
+        byte[] buffer = new byte[4 * 1024 * 1024];
 
         var enumData = new MFT_ENUM_DATA_V0
         {
@@ -179,15 +203,12 @@ public class MftReader : IDisposable
                             Logger.Log($"FSCTL_ENUM_USN_DATA: EOF after {iterations} iterations");
                             break;
                         }
-
-                        // Check if it's a buffer-too-small error
                         if (error == 122) // ERROR_INSUFFICIENT_BUFFER
                         {
                             Logger.Log("FSCTL_ENUM_USN_DATA: insufficient buffer, retrying with larger buffer");
                             buffer = new byte[buffer.Length * 2];
                             continue;
                         }
-
                         throw new IOException(
                             $"FSCTL_ENUM_USN_DATA failed (iter {iterations}): {Win32Error(error)}. " +
                             $"StartFileReferenceNumber={enumData.StartFileReferenceNumber}");
@@ -248,6 +269,46 @@ public class MftReader : IDisposable
         return new ScanResult(records, nameData);
     }
 
+    // ── Helpers ──
+
+    /// <summary>
+    /// Parses the $DATA attribute (0x80) from the $MFT record itself (record 0)
+    /// to discover the total size of the MFT data on disk.
+    /// </summary>
+    private static long ParseMftDataSize(ReadOnlySpan<byte> record)
+    {
+        if (record.Length < 48 || record[0] != 'F' || record[1] != 'I' || record[2] != 'L' || record[3] != 'E')
+            return -1;
+
+        // Apply fixup first so the record is valid
+        if (!ApplyFixup(record, record.Length))
+            return -1;
+
+        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;
+
+            uint alen = BinaryPrimitives.ReadUInt32LittleEndian(record.Slice(aoff + 4, 4));
+            if (alen == 0 || aoff + alen > record.Length) break;
+
+            // 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;
+            }
+
+            aoff += (int)alen;
+        }
+
+        return -1;
+    }
+
     private static bool ApplyFixup(ReadOnlySpan<byte> record, int recordSize)
     {
         if (record.Length < 48 || record[0] != 'F' || record[1] != 'I' || record[2] != 'L' || record[3] != 'E')