FastSeek.Core/UsnWatcher.cs

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

namespace FastSeek.Core.Mft;

public sealed class UsnWatcher : IDisposable
{
    private const int BufferSize = 64 * 1024;
    private readonly SafeFileHandle _handle;
    private readonly NtfsDrive _drive;
    private readonly BlockingCollection<IndexEvent> _sender;

    public long NextUsn { get; private set; }
    public ulong JournalId { get; private set; }

    private UsnWatcher(SafeFileHandle handle, NtfsDrive drive, BlockingCollection<IndexEvent> sender, long nextUsn, ulong journalId)
    {
        _handle = handle;
        _drive = drive;
        _sender = sender;
        NextUsn = nextUsn;
        JournalId = journalId;
    }

    public static UsnWatcher New(NtfsDrive drive, BlockingCollection<IndexEvent> sender) => NewFrom(drive, sender, null);

    public static UsnWatcher NewFrom(NtfsDrive drive, BlockingCollection<IndexEvent> sender, JournalCheckpoint? checkpoint)
    {
        var h = NativeMethods.CreateFile(drive.DevicePath, 0, 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("USN open failed");

        var outPtr = Marshal.AllocHGlobal(Marshal.SizeOf<UsnJournalDataV0>());
        try
        {
            if (!NativeMethods.DeviceIoControl(h, NativeMethods.FSCTL_QUERY_USN_JOURNAL, IntPtr.Zero, 0, outPtr, (uint)Marshal.SizeOf<UsnJournalDataV0>(), out _, IntPtr.Zero))
                throw new IOException("USN query failed");

            var data = Marshal.PtrToStructure<UsnJournalDataV0>(outPtr);
            long nextUsn;
            if (checkpoint is not null)
            {
                if (checkpoint.JournalId != data.UsnJournalID || checkpoint.NextUsn < data.FirstUsn || checkpoint.NextUsn > data.NextUsn)
                    throw new IOException("Journal mismatch - rescan needed");
                nextUsn = checkpoint.NextUsn;
            }
            else nextUsn = data.NextUsn;

            return new UsnWatcher(h, drive, sender, nextUsn, data.UsnJournalID);
        }
        finally { Marshal.FreeHGlobal(outPtr); }
    }

    public JournalCheckpoint Checkpoint() => new() { NextUsn = NextUsn, JournalId = JournalId, DriveLetter = _drive.Letter };

    public void RunShared(List<JournalCheckpoint> checkpoints, object checkpointLock, CancellationToken token)
    {
        var buffer = new byte[BufferSize];
        while (!token.IsCancellationRequested)
        {
            Thread.Sleep(500);
            Poll(buffer);
            lock (checkpointLock)
            {
                checkpoints.RemoveAll(c => c.DriveLetter == _drive.Letter);
                checkpoints.Add(Checkpoint());
            }
        }
    }

    public int Drain()
    {
        var buffer = new byte[BufferSize];
        var count = 0;
        while (true)
        {
            var before = NextUsn;
            Poll(buffer);
            if (NextUsn == before) break;
            count++;
        }
        return count;
    }

    private void Poll(byte[] buffer)
    {
        var read = new ReadUsnJournalDataV0
        {
            StartUsn = NextUsn,
            ReasonMask = NativeMethods.USN_REASON_FILE_CREATE | NativeMethods.USN_REASON_FILE_DELETE | NativeMethods.USN_REASON_RENAME_NEW_NAME | NativeMethods.USN_REASON_RENAME_OLD_NAME,
            ReturnOnlyOnClose = 0,
            Timeout = 0,
            BytesToWaitFor = 0,
            UsnJournalID = JournalId,
        };

        var inPtr = Marshal.AllocHGlobal(Marshal.SizeOf<ReadUsnJournalDataV0>());
        try
        {
            Marshal.StructureToPtr(read, inPtr, false);
            if (!NativeMethods.DeviceIoControl(_handle, NativeMethods.FSCTL_READ_USN_JOURNAL, inPtr, (uint)Marshal.SizeOf<ReadUsnJournalDataV0>(), buffer, (uint)buffer.Length, out var returned, IntPtr.Zero) || returned <= 8) return;
            NextUsn = BitConverter.ToInt64(buffer, 0);

            var offset = 8;
            while (offset + Marshal.SizeOf<UsnRecordV2>() <= returned)
            {
                var rec = MemoryMarshal.Read<UsnRecordV2>(new ReadOnlySpan<byte>(buffer, offset, Marshal.SizeOf<UsnRecordV2>()));
                if (rec.RecordLength == 0) break;
                ProcessRecord(rec, buffer, offset);
                offset += (int)rec.RecordLength;
            }
        }
        finally { Marshal.FreeHGlobal(inPtr); }
    }

    private void ProcessRecord(UsnRecordV2 record, byte[] buffer, int offset)
    {
        var nameLen = record.FileNameLength / 2;
        var nameStart = offset + record.FileNameOffset;
        var name = new string(MemoryMarshal.Cast<byte, char>(new ReadOnlySpan<byte>(buffer, nameStart, nameLen * 2)));
        var isDir = (record.FileAttributes & 0x10) != 0;
        var kind = isDir ? FileKind.Directory : FileKind.File;

        if ((record.Reason & NativeMethods.USN_REASON_FILE_DELETE) != 0) { _sender.Add(new IndexEvent.Deleted(record.FileReferenceNumber)); return; }
        if ((record.Reason & NativeMethods.USN_REASON_RENAME_NEW_NAME) != 0) { _sender.Add(new IndexEvent.Moved(record.FileReferenceNumber, record.ParentFileReferenceNumber, name, kind)); return; }
        if ((record.Reason & NativeMethods.USN_REASON_FILE_CREATE) != 0) _sender.Add(new IndexEvent.Created(new FileRecord { FileRef = record.FileReferenceNumber, ParentRef = record.ParentFileReferenceNumber, Name = name, Kind = kind }));
    }

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