FastSeekWpf/Core/MftReader.cs

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 _volumeHandle;
    private readonly NtfsDrive _drive;
    private bool _disposed;

    private const int FALLBACK_BUF = 4 * 1024 * 1024;
    private const int DIRECT_BUF = 4 * 1024 * 1024;

    public NtfsDrive Drive => _drive;

    public MftReader(NtfsDrive drive)
    {
        _drive = drive;
        _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_BACKUP_SEMANTICS,
            IntPtr.Zero);

        if (_volumeHandle == new IntPtr(-1))
        {
            int err = Marshal.GetLastWin32Error();
            throw new IOException($"Failed to open volume {drive.DevicePath}: {Win32Error(err)} (code={err})");
        }
    }

    // ------------------------------------------------------------------
    //  Try direct first, fall back to FSCTL (matches Rust main.rs exactly)
    // ------------------------------------------------------------------

    public (ScanResult scan, string method) ScanAny()
    {
        var direct = ScanDirect();
        if (direct != null)
            return (direct, "direct");
        return (Scan(), "ioctl");
    }

    // ------------------------------------------------------------------
    //  Primary: direct $MFT file read (matches Rust scan_direct())
    // ------------------------------------------------------------------

    public ScanResult? ScanDirect()
    {
        int? recordSize = ReadMftRecordSize();
        if (!recordSize.HasValue) return null;

        string mftPath = $"{_drive.Root}$MFT";
        Logger.Log($"[{Drive.Letter}] Opening $MFT: {mftPath}");

        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))
        {
            int err = Marshal.GetLastWin32Error();
            Logger.Log($"[{Drive.Letter}] CreateFileW({mftPath}) failed: {Win32Error(err)} (code={err})");
            return null;
        }

        try
        {
            var records = new List<CompactRecord>();
            var nameData = new List<char>();
            byte[] buffer = new byte[DIRECT_BUF];
            ulong mftIndex = 0;
            int leftover = 0;

            while (true)
            {
                uint bytesRead = 0;
                bool ok = Win32Api.ReadFile(
                    mftHandle, buffer, leftover, (uint)(buffer.Length - leftover), out bytesRead, IntPtr.Zero);

                if (!ok || bytesRead == 0)
                    break;

                int total = leftover + (int)bytesRead;
                int offset = 0;

                while (offset + recordSize.Value <= total)
                {
                    bool applied = ApplyFixup(buffer, offset, recordSize.Value);
                    if (applied)
                    {
                        ParseFileRecord(
                            buffer.AsSpan(offset, recordSize.Value),
                            mftIndex,
                            records,
                            nameData);
                    }

                    mftIndex += 1;
                    offset += recordSize.Value;
                }

                // Align down to record boundary (matches Rust: offset = total - (total % record_size))
                offset = total - (total % recordSize.Value);

                leftover = total - offset;
                if (leftover > 0)
                {
                    // Copy tail to front of buffer (matches Rust std::ptr::copy)
                    Buffer.BlockCopy(buffer, offset, buffer, 0, leftover);
                }
            }

            Logger.Log($"[{Drive.Letter}] Direct scan: {records.Count} records from {mftIndex} positions");
            return new ScanResult(records, nameData);
        }
        finally
        {
            Win32Api.CloseHandle(mftHandle);
        }
    }

    // ------------------------------------------------------------------
    //  Fallback: FSCTL_ENUM_USN_DATA (4 MB buffer)
    //  Matches Rust scan() exactly.
    // ------------------------------------------------------------------

    public ScanResult Scan()
    {
        var records = new List<CompactRecord>();
        var nameData = new List<char>();
        byte[] buffer = new byte[FALLBACK_BUF];

        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
        {
            int iterations = 0;
            while (true)
            {
                iterations++;
                IntPtr bufferPtr = Marshal.AllocHGlobal(buffer.Length);
                try
                {
                    uint bytesReturned = 0;
                    bool ok = Win32Api.DeviceIoControl(
                        _volumeHandle,
                        Win32Api.FSCTL_ENUM_USN_DATA,
                        enumDataPtr,
                        (uint)enumDataSize,
                        bufferPtr,
                        (uint)buffer.Length,
                        out bytesReturned,
                        IntPtr.Zero);

                    if (!ok)
                    {
                        int error = Marshal.GetLastWin32Error();
                        uint code = (uint)error;
                        if (code == 0x80070026) // ERROR_HANDLE_EOF
                        {
                            Logger.Log($"[{Drive.Letter}] FSCTL_ENUM_USN_DATA: EOF after {iterations} iterations");
                            break;
                        }
                        if (error == 122) // ERROR_INSUFFICIENT_BUFFER
                        {
                            Logger.Log($"[{Drive.Letter}] FSCTL buffer too small ({buffer.Length}), retrying");
                            buffer = new byte[buffer.Length * 2];
                            continue;
                        }
                        Logger.Log($"[{Drive.Letter}] FSCTL_ENUM_USN_DATA failed: {Win32Error(error)} (code={error})");
                        break;
                    }

                    if (bytesReturned <= 8)
                        break;

                    ulong nextRef = (ulong)Marshal.ReadInt64(bufferPtr);
                    enumData.StartFileReferenceNumber = nextRef;
                    Marshal.StructureToPtr(enumData, enumDataPtr, false);

                    int offset = 8;
                    int usnRecordSize = Marshal.SizeOf<USN_RECORD_V2>();
                    while (offset + usnRecordSize <= (int)bytesReturned)
                    {
                        IntPtr recordPtr = IntPtr.Add(bufferPtr, offset);
                        var record = Marshal.PtrToStructure<USN_RECORD_V2>(recordPtr);

                        int recLen = (int)record.RecordLength;
                        if (recLen == 0 || offset + recLen > (int)bytesReturned)
                            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 += recLen;
                    }
                }
                finally
                {
                    Marshal.FreeHGlobal(bufferPtr);
                }
            }
        }
        finally
        {
            Marshal.FreeHGlobal(enumDataPtr);
        }

        Logger.Log($"[{Drive.Letter}] FSCTL scan: {records.Count} records");
        return new ScanResult(records, nameData);
    }

    // ------------------------------------------------------------------
    //  NTFS helpers (match Rust exactly)
    // ------------------------------------------------------------------

    private int? ReadMftRecordSize()
    {
        // Seek to beginning of volume
        Win32Api.SetFilePointerEx(_volumeHandle, 0, out _, 0);

        byte[] boot = new byte[512];
        uint br = 0;
        bool ok = Win32Api.ReadFile(_volumeHandle, boot, 512, out 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(boot.AsSpan(0x0B, 2));
        byte sectorsPerCluster = boot[0x0D];
        int clusterSize = bytesPerSector * sectorsPerCluster;

        sbyte raw = (sbyte)boot[0x40];
        int recordSize = raw > 0 ? raw * clusterSize : 1 << (-raw);

        Logger.Log($"[{Drive.Letter}] Boot: bps={bytesPerSector}, spc={sectorsPerCluster}, clusterSize={clusterSize}, recordSize={recordSize}");
        return recordSize;
    }

    /// <summary>
    /// Apply NTFS multi-sector fixup. Returns false if record is invalid.
    /// CRITICAL: restores bytes from fixup array back to sector ends (matches Rust exactly).
    /// </summary>
    private static bool ApplyFixup(byte[] record, int offset, int recordSize)
    {
        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.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[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[offset + end] != check0 || record[offset + end + 1] != check1)
                return false;

            // RESTORE real bytes from fixup array (this is what Rust does)
            record[offset + end] = record[offset + fixupOff + i * 2];
            record[offset + end + 1] = record[offset + fixupOff + i * 2 + 1];
        }

        return true;
    }

    /// <summary>
    /// Parse one MFT FILE record (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

        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) // $FILE_NAME, resident
            {
                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) // DOS 8.3 name — skip
                        {
                            aoff += (int)alen;
                            continue;
                        }

                        byte priority = ns switch
                        {
                            1 => 0, // Win32
                            3 => 1, // Win32 + DOS
                            0 => 2, // POSIX
                            _ => 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
            });
        }
    }

    private static string Win32Error(int code)
    {
        try
        {
            return new System.ComponentModel.Win32Exception(code).Message;
        }
        catch
        {
            return $"error {code} (0x{code:X})";
        }
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            Win32Api.CloseHandle(_volumeHandle);
            _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;
    }
}