FastSeekWpf/Core/UsnWatcher.cs

using System;
using System.Buffers.Binary;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using FastSeekWpf.NativeInterop;

namespace FastSeekWpf.Core;

public class UsnWatcher : IDisposable
{
    private readonly IntPtr _handle;
    private readonly NtfsDrive _drive;
    private readonly Action<IndexEvent> _onEvent;
    private bool _disposed;
    private bool _running;

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

    public UsnWatcher(NtfsDrive drive, Action<IndexEvent> onEvent, JournalCheckpoint? checkpoint)
    {
        _drive = drive;
        _onEvent = onEvent;

        _handle = Win32Api.CreateFileW(
            drive.DevicePath,
            0,
            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}:");

        var journalData = new USN_JOURNAL_DATA_V0();
        IntPtr outBuf = Marshal.AllocHGlobal(Marshal.SizeOf<USN_JOURNAL_DATA_V0>());
        try
        {
            bool ok = Win32Api.DeviceIoControl(
                _handle, Win32Api.FSCTL_QUERY_USN_JOURNAL,
                IntPtr.Zero, 0,
                outBuf, (uint)Marshal.SizeOf<USN_JOURNAL_DATA_V0>(),
                out uint bytesReturned,
                IntPtr.Zero);

            if (!ok)
                throw new IOException("Failed to query USN journal");

            journalData = Marshal.PtrToStructure<USN_JOURNAL_DATA_V0>(outBuf);
        }
        finally
        {
            Marshal.FreeHGlobal(outBuf);
        }

        if (checkpoint != null)
        {
            if (checkpoint.JournalId != journalData.UsnJournalID)
                throw new IOException("Journal ID mismatch — rescan needed");
            if (checkpoint.NextUsn < journalData.FirstUsn || checkpoint.NextUsn > journalData.NextUsn)
                throw new IOException("Saved USN outside journal range — rescan needed");
            NextUsn = checkpoint.NextUsn;
        }
        else
        {
            NextUsn = journalData.NextUsn;
        }

        JournalId = journalData.UsnJournalID;
    }

    public static UsnWatcher CreateForCheckpoint(NtfsDrive drive, Action<IndexEvent> onEvent, JournalCheckpoint checkpoint)
    {
        return new UsnWatcher(drive, onEvent, checkpoint);
    }

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

    public int Drain()
    {
        int count = 0;
        while (true)
        {
            long before = NextUsn;
            Poll();
            if (NextUsn == before) break;
            count++;
        }
        return count;
    }

    /// <summary>Synchronous run loop — matches Rust UsnWatcher::run()</summary>
    public void Run()
    {
        byte[] buffer = new byte[65536];
        while (!_disposed)
        {
            Thread.Sleep(500);
            Poll();
        }
    }

    /// <summary>Synchronous run with shared checkpoint updates — matches Rust UsnWatcher::run_shared()</summary>
    public void RunShared(List<JournalCheckpoint> sharedCheckpoints)
    {
        byte[] buffer = new byte[65536];
        while (!_disposed)
        {
            Thread.Sleep(500);
            Poll();
            lock (sharedCheckpoints)
            {
                sharedCheckpoints.RemoveAll(c => c.DriveLetter == _drive.Letter);
                sharedCheckpoints.Add(Checkpoint());
            }
        }
    }

    public async Task RunAsync(CancellationToken ct)
    {
        _running = true;
        byte[] buffer = new byte[65536];

        while (!_disposed && !ct.IsCancellationRequested && _running)
        {
            Poll();
            await Task.Delay(500, ct);
        }
    }

    private void Poll()
    {
        var readData = new READ_USN_JOURNAL_DATA_V0
        {
            StartUsn = NextUsn,
            ReasonMask = Win32Api.USN_REASON_FILE_CREATE
                | Win32Api.USN_REASON_FILE_DELETE
                | Win32Api.USN_REASON_RENAME_NEW_NAME
                | Win32Api.USN_REASON_RENAME_OLD_NAME,
            ReturnOnlyOnClose = 0,
            Timeout = 0,
            BytesToWaitFor = 0,
            UsnJournalID = JournalId
        };

        int readDataSize = Marshal.SizeOf<READ_USN_JOURNAL_DATA_V0>();
        IntPtr readDataPtr = Marshal.AllocHGlobal(readDataSize);
        IntPtr bufferPtr = Marshal.AllocHGlobal(65536);

        try
        {
            Marshal.StructureToPtr(readData, readDataPtr, false);

            bool ok = Win32Api.DeviceIoControl(
                _handle, Win32Api.FSCTL_READ_USN_JOURNAL,
                readDataPtr, (uint)readDataSize,
                bufferPtr, 65536,
                out uint bytesReturned,
                IntPtr.Zero);

            if (!ok || bytesReturned <= 8) return;

            NextUsn = Marshal.ReadInt64(bufferPtr);
            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;

                ProcessRecord(record, bufferPtr, offset);
                offset += (int)record.RecordLength;
            }
        }
        finally
        {
            Marshal.FreeHGlobal(readDataPtr);
            Marshal.FreeHGlobal(bufferPtr);
        }
    }

    private void ProcessRecord(USN_RECORD_V2 record, IntPtr bufferBase, int offset)
    {
        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(bufferBase, nameOffset + i * 2);
        string name = new string(nameChars);

        bool isDir = (record.FileAttributes & 0x10) != 0;
        ulong fileRef = record.FileReferenceNumber;
        ulong parentRef = record.ParentFileReferenceNumber;
        uint reason = record.Reason;
        var kind = isDir ? FileKind.Directory : FileKind.File;

        if ((reason & Win32Api.USN_REASON_FILE_DELETE) != 0)
        {
            _onEvent(new IndexEvent.Deleted(fileRef));
            return;
        }

        if ((reason & Win32Api.USN_REASON_RENAME_NEW_NAME) != 0)
        {
            _onEvent(new IndexEvent.Moved(fileRef, parentRef, name, kind));
            return;
        }

        if ((reason & Win32Api.USN_REASON_FILE_CREATE) != 0)
        {
            _onEvent(new IndexEvent.Created(new FileRecord
            {
                FileRef = fileRef,
                ParentRef = parentRef,
                Name = name,
                Kind = kind
            }));
        }
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            _running = false;
            Win32Api.CloseHandle(_handle);
            _disposed = true;
        }
        GC.SuppressFinalize(this);
    }

    ~UsnWatcher() => Dispose();
}