FastSeek.Core/MftReader.cs

using System.Runtime.InteropServices;
using Microsoft.Win32.SafeHandles;

namespace FastSeek.Core.Mft;

internal sealed class MftReader : IDisposable
{
    private const int FallbackBuf = 4 * 1024 * 1024;
    private const int DirectBuf = 4 * 1024 * 1024;

    private readonly SafeFileHandle _handle;
    public NtfsDrive Drive { get; }

    private MftReader(SafeFileHandle handle, NtfsDrive drive)
    {
        _handle = handle;
        Drive = drive;
    }

    public static MftReader Open(NtfsDrive drive)
    {
        var h = NativeMethods.CreateFile(drive.DevicePath, NativeMethods.GENERIC_READ, NativeMethods.FILE_SHARE_READ | NativeMethods.FILE_SHARE_WRITE | NativeMethods.FILE_SHARE_DELETE, IntPtr.Zero, NativeMethods.OPEN_EXISTING, NativeMethods.FILE_FLAG_BACKUP_SEMANTICS, IntPtr.Zero);
        if (h.IsInvalid) throw new IOException($"Open drive failed: {drive.DevicePath}");
        return new MftReader(h, drive);
    }

    public ScanResult? ScanDirect(TimeSpan? maxDuration = null, Action<string>? log = null)
    {
        var started = DateTime.UtcNow;
        var recordSize = ReadMftRecordSize();
        if (recordSize is null)
        {
            log?.Invoke("direct: failed to read NTFS record size");
            return null;
        }
        log?.Invoke($"direct: record size={recordSize.Value} bytes");

        var mftPath = Path.Combine(Drive.Root, "$MFT");
        var mft = NativeMethods.CreateFile(mftPath, NativeMethods.GENERIC_READ, NativeMethods.FILE_SHARE_READ | NativeMethods.FILE_SHARE_WRITE | NativeMethods.FILE_SHARE_DELETE, IntPtr.Zero, NativeMethods.OPEN_EXISTING, NativeMethods.FILE_FLAG_BACKUP_SEMANTICS | NativeMethods.FILE_FLAG_SEQUENTIAL_SCAN, IntPtr.Zero);
        if (mft.IsInvalid)
        {
            log?.Invoke("direct: could not open $MFT");
            return null;
        }

        try
        {
            var records = new List<CompactRecord>(3_000_000);
            var names = new List<char>(40_000_000);
            var buffer = new byte[DirectBuf];
            var readChunk = new byte[DirectBuf];
            ulong mftIndex = 0;
            var leftover = 0;
            long totalRead = 0;
            var lastLog = DateTime.UtcNow;

            while (NativeMethods.ReadFile(mft, readChunk, (uint)(buffer.Length - leftover), out var read, IntPtr.Zero) && read > 0)
            {
                if (maxDuration.HasValue && (DateTime.UtcNow - started) > maxDuration.Value)
                {
                    log?.Invoke($"direct: timed out after {(DateTime.UtcNow - started).TotalSeconds:F1}s, fallback to ioctl");
                    return null;
                }
                totalRead += read;
                Buffer.BlockCopy(readChunk, 0, buffer, leftover, (int)read);
                var total = leftover + (int)read;
                var offset = 0;
                while (offset + recordSize.Value <= total)
                {
                    var slice = new Span<byte>(buffer, offset, recordSize.Value);
                    if (ApplyFixup(slice)) ParseFileRecord(slice, mftIndex, records, names);
                    mftIndex++;
                    offset += recordSize.Value;
                }
                offset = total - (total % recordSize.Value);
                leftover = total - offset;
                if (leftover > 0) Buffer.BlockCopy(buffer, offset, buffer, 0, leftover);

                var now = DateTime.UtcNow;
                if ((now - lastLog).TotalSeconds >= 2.0)
                {
                    log?.Invoke($"direct: read={totalRead / (1024 * 1024)}MB parsed={records.Count:N0}");
                    lastLog = now;
                }
            }

            log?.Invoke($"direct: completed read={totalRead / (1024 * 1024)}MB parsed={records.Count:N0}");
            return new ScanResult { Records = records, NameData = names };
        }
        finally { mft.Dispose(); }
    }

    public ScanResult Scan(Action<string>? log = null)
    {
        var records = new List<CompactRecord>(3_000_000);
        var names = new List<char>(40_000_000);
        var enumData = new MftEnumDataV0 { StartFileReferenceNumber = 0, LowUsn = 0, HighUsn = long.MaxValue };
        var outBuf = new byte[FallbackBuf];
        var loops = 0;
        var lastLog = DateTime.UtcNow;

        while (true)
        {
            var pIn = Marshal.AllocHGlobal(Marshal.SizeOf<MftEnumDataV0>());
            try
            {
                Marshal.StructureToPtr(enumData, pIn, false);
                if (!NativeMethods.DeviceIoControl(_handle, NativeMethods.FSCTL_ENUM_USN_DATA, pIn, (uint)Marshal.SizeOf<MftEnumDataV0>(), outBuf, (uint)outBuf.Length, out var returned, IntPtr.Zero))
                {
                    var err = Marshal.GetLastWin32Error();
                    log?.Invoke($"ioctl: DeviceIoControl stop err=0x{err:X8}");
                    break;
                }
                if (returned <= 8)
                {
                    log?.Invoke("ioctl: no more records");
                    break;
                }
                loops++;

                enumData.StartFileReferenceNumber = BitConverter.ToUInt64(outBuf, 0);
                var offset = 8;
                while (offset + Marshal.SizeOf<UsnRecordV2>() <= returned)
                {
                    var rec = MemoryMarshal.Read<UsnRecordV2>(new ReadOnlySpan<byte>(outBuf, offset, Marshal.SizeOf<UsnRecordV2>()));
                    if (rec.RecordLength == 0 || offset + rec.RecordLength > returned) break;
                    var nameLen = (ushort)(rec.FileNameLength / 2);
                    var nameStart = offset + rec.FileNameOffset;
                    var chars = MemoryMarshal.Cast<byte, char>(new ReadOnlySpan<byte>(outBuf, nameStart, nameLen * 2));
                    var arenaOff = (uint)names.Count;
                    AppendChars(names, chars);
                    records.Add(new CompactRecord { FileRef = rec.FileReferenceNumber, ParentRef = rec.ParentFileReferenceNumber, NameOff = arenaOff, NameLen = nameLen, IsDir = (rec.FileAttributes & 0x10) != 0 });
                    offset += (int)rec.RecordLength;
                }

                var now = DateTime.UtcNow;
                if ((now - lastLog).TotalSeconds >= 2.0)
                {
                    log?.Invoke($"ioctl: loops={loops:N0} parsed={records.Count:N0} nextRef={enumData.StartFileReferenceNumber}");
                    lastLog = now;
                }
            }
            finally { Marshal.FreeHGlobal(pIn); }
        }

        log?.Invoke($"ioctl: completed loops={loops:N0} parsed={records.Count:N0}");
        return new ScanResult { Records = records, NameData = names };
    }

    private int? ReadMftRecordSize()
    {
        if (!NativeMethods.SetFilePointerEx(_handle, 0, out _, NativeMethods.FILE_BEGIN)) return null;
        var boot = new byte[512];
        if (!NativeMethods.ReadFile(_handle, boot, 512, out var read, IntPtr.Zero) || read < 512) return null;
        if (boot[3] != (byte)'N' || boot[4] != (byte)'T' || boot[5] != (byte)'F' || boot[6] != (byte)'S') return null;
        var bps = BitConverter.ToUInt16(boot, 0x0B);
        var spc = boot[0x0D];
        var cluster = bps * spc;
        var raw = unchecked((sbyte)boot[0x40]);
        return raw > 0 ? raw * cluster : 1 << -raw;
    }

    private static bool ApplyFixup(Span<byte> rec)
    {
        if (rec.Length < 48 || rec[0] != (byte)'F' || rec[1] != (byte)'I' || rec[2] != (byte)'L' || rec[3] != (byte)'E') return false;
        var off = BitConverter.ToUInt16(rec.Slice(4, 2));
        var cnt = BitConverter.ToUInt16(rec.Slice(6, 2));
        if (cnt < 2 || off + cnt * 2 > rec.Length) return false;
        var c0 = rec[off]; var c1 = rec[off + 1];
        for (var i = 1; i < cnt; i++)
        {
            var end = i * 512 - 2;
            if (end + 1 >= rec.Length) break;
            if (rec[end] != c0 || rec[end + 1] != c1) return false;
            rec[end] = rec[off + i * 2];
            rec[end + 1] = rec[off + i * 2 + 1];
        }
        return true;
    }

    private static void ParseFileRecord(ReadOnlySpan<byte> rec, ulong mftIndex, List<CompactRecord> outRecords, List<char> outNames)
    {
        var flags = BitConverter.ToUInt16(rec.Slice(0x16, 2));
        if ((flags & 0x01) == 0) return;
        var isDir = (flags & 0x02) != 0;
        var seq = BitConverter.ToUInt16(rec.Slice(0x10, 2));
        var fileRef = mftIndex | ((ulong)seq << 48);
        var firstAttr = BitConverter.ToUInt16(rec.Slice(0x14, 2));
        var aoff = (int)firstAttr;
        byte bestNs = 255;
        (int pos, int len, ulong parent)? best = null;

        while (aoff + 8 <= rec.Length)
        {
            var atype = BitConverter.ToUInt32(rec.Slice(aoff, 4));
            if (atype == 0xFFFFFFFF) break;
            var alen = (int)BitConverter.ToUInt32(rec.Slice(aoff + 4, 4));
            if (alen == 0 || aoff + alen > rec.Length) break;
            if (atype == 0x30 && rec[aoff + 8] == 0)
            {
                var voff = (int)BitConverter.ToUInt16(rec.Slice(aoff + 20, 2));
                var vs = aoff + voff;
                if (vs + 66 <= rec.Length)
                {
                    var parent = BitConverter.ToUInt64(rec.Slice(vs, 8));
                    var nlen = rec[vs + 64];
                    var ns = rec[vs + 65];
                    if (vs + 66 + nlen * 2 <= rec.Length && ns != 2)
                    {
                        var pr = ns == 1 ? 0 : ns == 3 ? 1 : ns == 0 ? 2 : 3;
                        if (pr < bestNs)
                        {
                            bestNs = (byte)pr;
                            best = (vs + 66, nlen, parent);
                            if (pr == 0) break;
                        }
                    }
                }
            }
            aoff += alen;
        }

        if (best is { } b)
        {
            var arenaOff = (uint)outNames.Count;
            var chars = MemoryMarshal.Cast<byte, char>(rec.Slice(b.pos, b.len * 2));
            AppendChars(outNames, chars);
            outRecords.Add(new CompactRecord { FileRef = fileRef, ParentRef = b.parent, NameOff = arenaOff, NameLen = (ushort)b.len, IsDir = isDir });
        }
    }

    private static void AppendChars(List<char> target, ReadOnlySpan<char> source)
    {
        target.Capacity = Math.Max(target.Capacity, target.Count + source.Length);
        for (var i = 0; i < source.Length; i++) target.Add(source[i]);
    }

    public void Dispose() => _handle.Dispose();
}