File size: 5,639 Bytes
1c4658f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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();
}