Upload FastSeekWpf/Core/MftReader.cs

FastSeekWpf/Core/MftReader.cs

@@ -0,0 +1,350 @@
+using System;
+using System.Buffers.Binary;
+using System.Collections.Generic;
+using System.IO;
+using System.Runtime.InteropServices;
+using FastSeekWpf.NativeInterop;
+
+namespace FastSeekWpf.Core;
+
+public class MftReader : IDisposable
+{
+    private readonly IntPtr _handle;
+    private readonly NtfsDrive _drive;
+    private bool _disposed;
+
+    private const uint FallbackBuf = 4 * 1024 * 1024;
+    private const uint DirectBuf = 4 * 1024 * 1024;
+
+    public NtfsDrive Drive => _drive;
+
+    public MftReader(NtfsDrive drive)
+    {
+        _drive = drive;
+        _handle = 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_BACKUP_SEMANTICS,
+            IntPtr.Zero);
+
+        if (_handle == new IntPtr(-1))
+            throw new IOException($"Failed to open drive {drive.Letter}:");
+    }
+
+    public ScanResult? ScanDirect()
+    {
+        int? recordSize = ReadMftRecordSize();
+        if (recordSize == null) return null;
+
+        string mftPath = $"{_drive.Root}$MFT";
+        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);
+
+        if (mftHandle == new IntPtr(-1))
+            return null;
+
+        try
+        {
+            var records = new List<CompactRecord>();
+            var nameData = new List<char>();
+            byte[] buffer = new byte[DirectBuf];
+            ulong mftIndex = 0;
+            int leftover = 0;
+
+            while (true)
+            {
+                bool ok = Win32Api.ReadFile(mftHandle, buffer, (uint)(buffer.Length - leftover), out uint bytesRead, IntPtr.Zero);
+                if (!ok || bytesRead == 0) break;
+
+                int total = leftover + (int)bytesRead;
+                int offset = 0;
+
+                while (offset + recordSize.Value <= total)
+                {
+                    var span = new ReadOnlySpan<byte>(buffer, offset, recordSize.Value);
+                    if (ApplyFixup(span, recordSize.Value))
+                    {
+                        ParseFileRecord(span, mftIndex, records, nameData);
+                    }
+                    mftIndex++;
+                    offset += recordSize.Value;
+                }
+
+                leftover = total - offset;
+                if (leftover > 0)
+                    Array.Copy(buffer, offset, buffer, 0, leftover);
+            }
+
+            return new ScanResult(records, nameData);
+        }
+        finally
+        {
+            Win32Api.CloseHandle(mftHandle);
+        }
+    }
+
+    public ScanResult Scan()
+    {
+        var records = new List<CompactRecord>();
+        var nameData = new List<char>();
+        byte[] buffer = new byte[FallbackBuf];
+
+        var enumData = new MFT_ENUM_DATA_V0
+        {
+            StartFileReferenceNumber = 0,
+            LowUsn = 0,
+            HighUsn = long.MaxValue
+        };
+
+        int enumDataSize = Marshal.SizeOf<MFT_ENUM_DATA_V0>();
+        IntPtr enumDataPtr = Marshal.AllocHGlobal(enumDataSize);
+        Marshal.StructureToPtr(enumData, enumDataPtr, false);
+
+        try
+        {
+            while (true)
+            {
+                IntPtr bufferPtr = Marshal.AllocHGlobal(buffer.Length);
+                try
+                {
+                    bool ok = Win32Api.DeviceIoControl(
+                        _handle, Win32Api.FSCTL_ENUM_USN_DATA,
+                        enumDataPtr, (uint)enumDataSize,
+                        bufferPtr, (uint)buffer.Length,
+                        out uint bytesReturned,
+                        IntPtr.Zero);
+
+                    if (!ok)
+                    {
+                        int error = Marshal.GetLastWin32Error();
+                        if ((uint)error == 0x26) break;
+                        break;
+                    }
+
+                    if (bytesReturned <= 8) break;
+
+                    ulong nextRef = (ulong)Marshal.ReadInt64(bufferPtr);
+                    enumData.StartFileReferenceNumber = nextRef;
+                    Marshal.StructureToPtr(enumData, enumDataPtr, false);
+
+                    int offset = 8;
+                    while (offset + Marshal.SizeOf<USN_RECORD_V2>() <= (int)bytesReturned)
+                    {
+                        IntPtr recordPtr = IntPtr.Add(bufferPtr, offset);
+                        var record = Marshal.PtrToStructure<USN_RECORD_V2>(recordPtr);
+
+                        if (record.RecordLength == 0) break;
+
+                        int nameOffset = offset + (int)record.FileNameOffset;
+                        int nameLen = record.FileNameLength / 2;
+
+                        var nameChars = new char[nameLen];
+                        for (int i = 0; i < nameLen; i++)
+                            nameChars[i] = (char)Marshal.ReadInt16(bufferPtr, nameOffset + i * 2);
+
+                        int arenaOff = nameData.Count;
+                        nameData.AddRange(nameChars);
+
+                        records.Add(new CompactRecord
+                        {
+                            FileRef = record.FileReferenceNumber,
+                            ParentRef = record.ParentFileReferenceNumber,
+                            NameOff = (uint)arenaOff,
+                            NameLen = (ushort)nameLen,
+                            IsDir = (record.FileAttributes & 0x10) != 0
+                        });
+
+                        offset += (int)record.RecordLength;
+                    }
+                }
+                finally
+                {
+                    Marshal.FreeHGlobal(bufferPtr);
+                }
+            }
+        }
+        finally
+        {
+            Marshal.FreeHGlobal(enumDataPtr);
+        }
+
+        return new ScanResult(records, nameData);
+    }
+
+    private int? ReadMftRecordSize()
+    {
+        bool ok = Win32Api.SetFilePointerEx(_handle, 0, out _, 0);
+        if (!ok) return null;
+
+        byte[] boot = new byte[512];
+        ok = Win32Api.ReadFile(_handle, boot, 512, out uint br, IntPtr.Zero);
+        if (!ok || br < 512) return null;
+
+        if (boot[3] != 'N' || boot[4] != 'T' || boot[5] != 'F' || boot[6] != 'S')
+            return null;
+
+        ushort bytesPerSector = BinaryPrimitives.ReadUInt16LittleEndian(new ReadOnlySpan<byte>(boot, 0x0B, 2));
+        byte sectorsPerCluster = boot[0x0D];
+        int clusterSize = bytesPerSector * sectorsPerCluster;
+
+        sbyte raw = (sbyte)boot[0x40];
+        if (raw > 0)
+            return raw * clusterSize;
+        else
+            return 1 << (-raw);
+    }
+
+    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')
+            return false;
+
+        ushort fixupOff = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(4, 2));
+        ushort fixupCnt = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(6, 2));
+
+        if (fixupCnt < 2 || fixupOff + fixupCnt * 2 > recordSize)
+            return false;
+
+        byte check0 = record[fixupOff];
+        byte check1 = record[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)
+                return false;
+        }
+
+        return true;
+    }
+
+    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;
+
+        bool isDir = (flags & 0x02) != 0;
+        ushort seq = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(0x10, 2));
+        ulong fileRef = mftIndex | ((ulong)seq << 48);
+
+        ushort firstAttr = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(0x14, 2));
+        int aoff = firstAttr;
+
+        byte bestNs = 255;
+        (int pos, int len, ulong parent)? bestName = null;
+
+        while (aoff + 8 <= 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;
+
+            if (atype == 0x30 && record[aoff + 8] == 0)
+            {
+                uint vlen = BinaryPrimitives.ReadUInt32LittleEndian(record.Slice(aoff + 16, 4));
+                ushort voff = BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(aoff + 20, 2));
+                int vs = aoff + voff;
+
+                if (vs + 66 <= record.Length && vlen >= 66)
+                {
+                    ulong parent = BinaryPrimitives.ReadUInt64LittleEndian(record.Slice(vs, 8));
+                    int nlen = record[vs + 64];
+                    byte ns = record[vs + 65];
+
+                    if (vs + 66 + nlen * 2 <= record.Length)
+                    {
+                        if (ns == 2)
+                        {
+                            aoff += (int)alen;
+                            continue;
+                        }
+
+                        byte priority = ns switch
+                        {
+                            1 => 0,
+                            3 => 1,
+                            0 => 2,
+                            _ => 3,
+                        };
+
+                        if (priority < bestNs)
+                        {
+                            bestNs = priority;
+                            bestName = (vs + 66, nlen, parent);
+                            if (priority == 0) break;
+                        }
+                    }
+                }
+            }
+
+            aoff += (int)alen;
+        }
+
+        if (bestName.HasValue)
+        {
+            var (namePos, nlen, parent) = bestName.Value;
+            int arenaOff = nameData.Count;
+
+            for (int i = 0; i < nlen; i++)
+            {
+                int p = namePos + i * 2;
+                nameData.Add((char)BinaryPrimitives.ReadUInt16LittleEndian(record.Slice(p, 2)));
+            }
+
+            records.Add(new CompactRecord
+            {
+                FileRef = fileRef,
+                ParentRef = parent,
+                NameOff = (uint)arenaOff,
+                NameLen = (ushort)nlen,
+                IsDir = isDir
+            });
+        }
+    }
+
+    public void Dispose()
+    {
+        if (!_disposed)
+        {
+            Win32Api.CloseHandle(_handle);
+            _disposed = true;
+        }
+        GC.SuppressFinalize(this);
+    }
+
+    ~MftReader() => Dispose();
+}
+
+public class CompactRecord
+{
+    public ulong FileRef { get; set; }
+    public ulong ParentRef { get; set; }
+    public uint NameOff { get; set; }
+    public ushort NameLen { get; set; }
+    public bool IsDir { get; set; }
+}
+
+public class ScanResult
+{
+    public List<CompactRecord> Records { get; }
+    public List<char> NameData { get; }
+
+    public ScanResult(List<CompactRecord> records, List<char> nameData)
+    {
+        Records = records;
+        NameData = nameData;
+    }
+}