diff --git a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.ProvisioningStatusTracking.cs b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.ProvisioningStatusTracking.cs index 29137342..c9b2ca4c 100644 --- a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.ProvisioningStatusTracking.cs +++ b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.ProvisioningStatusTracking.cs @@ -2,7 +2,7 @@ using System; using System.Collections.Generic; using System.Linq; using System.Text.Json; -using System.Threading; +using AutopilotMonitor.Agent.Core.Monitoring.Core; using AutopilotMonitor.Agent.Core.Monitoring.Interop; using AutopilotMonitor.Shared; using AutopilotMonitor.Shared.Models; @@ -30,7 +30,7 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors /// - Emit once when categorySucceeded resolves to true or false (final outcome) /// - Fire EspFailureDetected on subcategory failure even if categorySucceeded is still null /// - /// Uses RegNotifyChangeKeyValue for instant registry change detection (no polling timer). + /// Uses RegistryWatcher (RegNotifyChangeKeyValue) for instant registry change detection. /// public partial class EspAndHelloTracker { @@ -65,159 +65,80 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors _provisioningCategoriesResolved = new HashSet(StringComparer.OrdinalIgnoreCase); _lastSubcategoryStates = new Dictionary>(StringComparer.OrdinalIgnoreCase); - _registryWatcherStopEvent = RegistryWatcherNativeMethods.CreateEvent(IntPtr.Zero, true, false, null); - if (_registryWatcherStopEvent == IntPtr.Zero) + // Try to start immediately; if key doesn't exist yet, retry every 2s + if (!TryStartProvisioningWatcher()) { - _logger.Warning("Failed to create registry watcher stop event — provisioning status tracking disabled"); - return; - } - - _registryWatcherThread = new Thread(RegistryWatcherLoop) - { - IsBackground = true, - Name = "ProvisioningStatusWatcher" - }; - _registryWatcherThread.Start(); - - _logger.Info("Provisioning status registry watcher started"); - } - - private void StopProvisioningStatusWatcher(string reason) - { - if (_registryWatcherStopEvent == IntPtr.Zero) - return; - - try - { - // Signal stop event to wake up the watcher thread - RegistryWatcherNativeMethods.SetEvent(_registryWatcherStopEvent); - - // Wait for thread to exit - if (_registryWatcherThread != null && _registryWatcherThread.IsAlive) - { - _registryWatcherThread.Join(TimeSpan.FromSeconds(3)); - } - - RegistryWatcherNativeMethods.CloseHandle(_registryWatcherStopEvent); - _registryWatcherStopEvent = IntPtr.Zero; - _registryWatcherThread = null; - - _logger.Info($"Provisioning status watcher stopped: {reason}"); - } - catch (Exception ex) - { - _logger.Error("Error stopping provisioning status watcher", ex); + _logger.Info("Provisioning status registry key not yet present — retrying every 2s"); + _provisioningWatcherRetryTimer = new System.Threading.Timer( + _ => + { + if (TryStartProvisioningWatcher()) + { + _provisioningWatcherRetryTimer?.Dispose(); + _provisioningWatcherRetryTimer = null; + } + }, + null, + TimeSpan.FromSeconds(2), + TimeSpan.FromSeconds(2)); } } /// - /// Background thread: opens the provisioning registry key and waits for value changes - /// via RegNotifyChangeKeyValue. Falls back to retry loop if key doesn't exist yet. + /// Attempts to create and start the RegistryWatcher for the provisioning key. + /// Returns false if the key doesn't exist yet. /// - private void RegistryWatcherLoop() + private bool TryStartProvisioningWatcher() { - IntPtr hKey = IntPtr.Zero; - IntPtr hNotifyEvent = IntPtr.Zero; - try { - // Wait for registry key to appear (enrollment may not have started yet) - while (true) + using (var key = Registry.LocalMachine.OpenSubKey(ProvisioningStatusRegistryPath)) { - int result = RegistryWatcherNativeMethods.RegOpenKeyEx( - RegistryWatcherNativeMethods.HKEY_LOCAL_MACHINE, - ProvisioningStatusRegistryPath, - 0, - RegistryWatcherNativeMethods.KEY_NOTIFY | RegistryWatcherNativeMethods.KEY_READ, - out hKey); - - if (result == 0) - break; - - // Key doesn't exist yet — wait 2 seconds or until stop signaled - uint waitResult = RegistryWatcherNativeMethods.WaitForSingleObject(_registryWatcherStopEvent, 2000); - if (waitResult == RegistryWatcherNativeMethods.WAIT_OBJECT_0) - return; // Stop requested + if (key == null) + return false; } - // Capture initial state before notification loop. - // Any changes between this read and the first RegNotifyChangeKeyValue registration - // are covered by the 15s heartbeat fallback. + // Capture initial state before watcher starts CheckProvisioningStatus(); - // Create notification event - hNotifyEvent = RegistryWatcherNativeMethods.CreateEvent(IntPtr.Zero, false, false, null); - if (hNotifyEvent == IntPtr.Zero) + // Create and start watcher + _provisioningWatcher = new RegistryWatcher( + RegistryHive.LocalMachine, + ProvisioningStatusRegistryPath, + watchSubtree: true, + filter: RegistryNativeMethods.RegChangeNotifyFilter.LastSet); + + _provisioningWatcher.Changed += (s, e) => CheckProvisioningStatus(); + _provisioningWatcher.Error += (s, ex) => _logger.Warning($"Provisioning watcher handler error: {ex.Message}"); + + _provisioningWatcher.Start(); + _logger.Info("Provisioning status registry watcher started"); + return true; + } + catch (Exception ex) + { + _logger.Warning($"Failed to start provisioning watcher: {ex.Message}"); + return false; + } + } + + private void StopProvisioningStatusWatcher(string reason) + { + try + { + _provisioningWatcherRetryTimer?.Dispose(); + _provisioningWatcherRetryTimer = null; + + if (_provisioningWatcher != null) { - _logger.Warning("Failed to create registry notify event — watcher thread exiting"); - return; - } - - var waitHandles = new[] { hNotifyEvent, _registryWatcherStopEvent }; - - // Main watch loop: register → wait → read - // IMPORTANT: CheckProvisioningStatus() must run AFTER WaitForMultipleObjects, not between - // RegNotifyChangeKeyValue and Wait. The managed Registry.OpenSubKey() call interferes with - // pending RegNotifyChangeKeyValue notifications when called before Wait (proven in Build 416). - // The 15s heartbeat ensures changes are detected even if a notification is missed. - const uint heartbeatMs = 15_000; - - while (true) - { - // 1. Register for value change notifications - int regResult = RegistryWatcherNativeMethods.RegNotifyChangeKeyValue( - hKey, - true, // Watch subtree - RegistryWatcherNativeMethods.REG_NOTIFY_CHANGE_LAST_SET, - hNotifyEvent, - true); // Asynchronous - - if (regResult != 0) - { - _logger.Warning($"RegNotifyChangeKeyValue failed (error {regResult}) — watcher thread exiting"); - return; - } - - // 2. Wait for registry change notification or heartbeat timeout - uint waitResult = RegistryWatcherNativeMethods.WaitForMultipleObjects( - (uint)waitHandles.Length, - waitHandles, - false, - heartbeatMs); - - if (waitResult == RegistryWatcherNativeMethods.WAIT_OBJECT_0) - { - // Registry notification fired — instant detection - _logger.Debug("Provisioning status watcher: registry change notification received"); - } - else if (waitResult == RegistryWatcherNativeMethods.WAIT_OBJECT_0 + 1) - { - return; // Stop requested - } - else if (waitResult == RegistryWatcherNativeMethods.WAIT_TIMEOUT) - { - // Heartbeat — re-read to catch any missed notifications - } - else - { - _logger.Warning($"Registry watcher unexpected wait result: {waitResult} — exiting"); - return; - } - - // 3. Read current state AFTER wait (not between register and wait) - CheckProvisioningStatus(); + _provisioningWatcher.Dispose(); + _provisioningWatcher = null; + _logger.Info($"Provisioning status watcher stopped: {reason}"); } } catch (Exception ex) { - _logger.Warning($"Registry watcher thread failed: {ex.Message}"); - } - finally - { - if (hNotifyEvent != IntPtr.Zero) - RegistryWatcherNativeMethods.CloseHandle(hNotifyEvent); - if (hKey != IntPtr.Zero) - RegistryWatcherNativeMethods.RegCloseKey(hKey); + _logger.Error("Error stopping provisioning status watcher", ex); } } @@ -255,16 +176,15 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors } } - // Self-termination: only auto-stop when all categories resolved with success (no failures) - // Signal the stop event directly (we're on the watcher thread — can't call StopProvisioningStatusWatcher - // which would Join() on ourselves) + // Self-termination: only auto-stop when all categories resolved with success (no failures). + // Uses RequestStop() which is safe to call from within the Changed handler (non-blocking). if (_provisioningCategoriesResolved.Count > 0 && _lastProvisioningJson.Count > 0 && _provisioningCategoriesResolved.Count >= _lastProvisioningJson.Count && !_provisioningFailureFired.Any()) { - _logger.Info("Provisioning status watcher: all categories resolved with success — signaling stop"); - RegistryWatcherNativeMethods.SetEvent(_registryWatcherStopEvent); + _logger.Info("Provisioning status watcher: all categories resolved with success — requesting stop"); + _provisioningWatcher?.RequestStop(); } // Fire EspFailureDetected AFTER event emission diff --git a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.cs b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.cs index 5249929f..78d8b7e9 100644 --- a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.cs +++ b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.cs @@ -3,6 +3,7 @@ using System.Collections.Generic; using System.Diagnostics.Eventing.Reader; using System.Threading; using AutopilotMonitor.Agent.Core.Logging; +using AutopilotMonitor.Agent.Core.Monitoring.Core; using AutopilotMonitor.Shared.Models; using Microsoft.Win32; @@ -36,8 +37,8 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors private System.Threading.Timer _policyCheckTimer; private System.Threading.Timer _helloWaitTimer; private System.Threading.Timer _helloCompletionTimer; - private Thread _registryWatcherThread; - private IntPtr _registryWatcherStopEvent = IntPtr.Zero; + private RegistryWatcher _provisioningWatcher; + private System.Threading.Timer _provisioningWatcherRetryTimer; private bool _isPolicyConfigured = false; private bool _isHelloPolicyEnabled = false; // true when Hello policy is explicitly detected as enabled diff --git a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Core/RegistryWatcher.cs b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Core/RegistryWatcher.cs new file mode 100644 index 00000000..149010c8 --- /dev/null +++ b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Core/RegistryWatcher.cs @@ -0,0 +1,202 @@ +using System; +using System.ComponentModel; +using System.Threading; +using System.Threading.Tasks; +using AutopilotMonitor.Agent.Core.Monitoring.Interop; +using Microsoft.Win32; +using Microsoft.Win32.SafeHandles; + +namespace AutopilotMonitor.Agent.Core.Monitoring.Core +{ + /// + /// Watches a registry key for changes using RegNotifyChangeKeyValue. + /// Fires the Changed event on each detected change. + /// + /// Usage: + /// var watcher = new RegistryWatcher(RegistryHive.LocalMachine, @"SOFTWARE\...", watchSubtree: true); + /// watcher.Changed += (s, e) => { /* read values */ }; + /// watcher.Start(); + /// // later: + /// await watcher.StopAsync(); + /// watcher.Dispose(); + /// + /// Important: one change per registration — the watcher re-registers automatically after each notification. + /// The watched scope is a key (including values); to track a specific value, compare in the Changed handler. + /// + internal sealed class RegistryWatcher : IDisposable + { + private readonly UIntPtr _hive; + private readonly string _subKey; + private readonly bool _watchSubtree; + private readonly RegistryNativeMethods.RegChangeNotifyFilter _filter; + + private CancellationTokenSource _cts; + private Task _watchTask; + private readonly object _sync = new object(); + + /// + /// Fired when the watched key (or subtree) changes. Handler runs on a thread pool thread. + /// + public event EventHandler Changed; + + /// + /// Fired when the Changed handler throws. Allows the consumer to log without crashing the watcher. + /// + public event EventHandler Error; + + public RegistryWatcher( + RegistryHive hive, + string subKey, + bool watchSubtree = false, + RegistryNativeMethods.RegChangeNotifyFilter filter = + RegistryNativeMethods.RegChangeNotifyFilter.LastSet) + { + _hive = HiveToPointer(hive); + _subKey = subKey ?? throw new ArgumentNullException(nameof(subKey)); + _watchSubtree = watchSubtree; + _filter = filter; + } + + /// + /// Starts watching on a background thread pool thread. + /// Throws Win32Exception if the registry key cannot be opened. + /// + public void Start() + { + lock (_sync) + { + if (_watchTask != null) + throw new InvalidOperationException("Watcher is already running."); + + _cts = new CancellationTokenSource(); + _watchTask = Task.Run(() => WatchLoopAsync(_cts.Token)); + } + } + + /// + /// Non-blocking stop request. Safe to call from within the Changed handler + /// (does not wait for the loop to exit, so no deadlock). + /// + public void RequestStop() + { + lock (_sync) + { + _cts?.Cancel(); + } + } + + /// + /// Cancels the watcher and waits for the background loop to exit. + /// Do NOT call from within the Changed handler — use RequestStop() instead. + /// + public async Task StopAsync() + { + Task taskToWait; + + lock (_sync) + { + if (_cts == null) + return; + + _cts.Cancel(); + taskToWait = _watchTask; + } + + if (taskToWait != null) + { + try + { + await taskToWait.ConfigureAwait(false); + } + catch (OperationCanceledException) { } + catch (Win32Exception) { } + } + + lock (_sync) + { + _cts?.Dispose(); + _cts = null; + _watchTask = null; + } + } + + private async Task WatchLoopAsync(CancellationToken cancellationToken) + { + int openResult = RegistryNativeMethods.RegOpenKeyEx( + _hive, _subKey, 0, RegistryNativeMethods.KEY_NOTIFY, + out SafeRegistryHandle keyHandle); + + if (openResult != 0) + throw new Win32Exception(openResult, $"RegOpenKeyEx failed for '{_subKey}'"); + + using (keyHandle) + using (var changedEvent = new AutoResetEvent(false)) + using (var cancelEvent = new ManualResetEvent(false)) + using (cancellationToken.Register(() => cancelEvent.Set())) + { + var waitHandles = new WaitHandle[] { changedEvent, cancelEvent }; + + while (!cancellationToken.IsCancellationRequested) + { + // Register for exactly ONE change notification + int notifyResult = RegistryNativeMethods.RegNotifyChangeKeyValue( + keyHandle, + _watchSubtree, + _filter, + changedEvent.SafeWaitHandle, + true); + + if (notifyResult != 0) + throw new Win32Exception(notifyResult, "RegNotifyChangeKeyValue failed"); + + // Wait for change or cancellation + int signaled = WaitHandle.WaitAny(waitHandles); + + if (signaled == 1) + break; // cancel requested + + // Notify consumer + try + { + Changed?.Invoke(this, EventArgs.Empty); + } + catch (Exception ex) + { + Error?.Invoke(this, ex); + } + + await Task.Yield(); + } + + cancellationToken.ThrowIfCancellationRequested(); + } + } + + public void Dispose() + { + try + { + StopAsync().GetAwaiter().GetResult(); + } + catch + { + // Dispose must not throw + } + } + + // ===== Helpers ===== + + private static UIntPtr HiveToPointer(RegistryHive hive) + { + switch (hive) + { + case RegistryHive.ClassesRoot: return (UIntPtr)0x80000000u; + case RegistryHive.CurrentUser: return (UIntPtr)0x80000001u; + case RegistryHive.LocalMachine: return (UIntPtr)0x80000002u; + case RegistryHive.Users: return (UIntPtr)0x80000003u; + case RegistryHive.CurrentConfig: return (UIntPtr)0x80000005u; + default: throw new ArgumentOutOfRangeException(nameof(hive)); + } + } + } +} diff --git a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryNativeMethods.cs b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryNativeMethods.cs new file mode 100644 index 00000000..3d176935 --- /dev/null +++ b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryNativeMethods.cs @@ -0,0 +1,36 @@ +using System; +using System.Runtime.InteropServices; +using Microsoft.Win32.SafeHandles; + +namespace AutopilotMonitor.Agent.Core.Monitoring.Interop +{ + /// + /// P/Invoke declarations for registry change notification via RegNotifyChangeKeyValue. + /// Uses SafeRegistryHandle and SafeWaitHandle for proper handle lifecycle management. + /// Used by RegistryWatcher to detect registry changes instantly. + /// + internal static class RegistryNativeMethods + { + public const int KEY_NOTIFY = 0x0010; + + [Flags] + public enum RegChangeNotifyFilter : uint + { + Name = 1, + Attributes = 2, + LastSet = 4, + Security = 8 + } + + [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)] + public static extern int RegOpenKeyEx( + UIntPtr hKey, string lpSubKey, uint ulOptions, + int samDesired, out SafeRegistryHandle phkResult); + + [DllImport("advapi32.dll", SetLastError = true)] + public static extern int RegNotifyChangeKeyValue( + SafeRegistryHandle hKey, bool bWatchSubtree, + RegChangeNotifyFilter dwNotifyFilter, + SafeWaitHandle hEvent, bool fAsynchronous); + } +} diff --git a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryWatcherNativeMethods.cs b/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryWatcherNativeMethods.cs deleted file mode 100644 index 3be50b9f..00000000 --- a/src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryWatcherNativeMethods.cs +++ /dev/null @@ -1,51 +0,0 @@ -using System; -using System.Runtime.InteropServices; - -namespace AutopilotMonitor.Agent.Core.Monitoring.Interop -{ - /// - /// P/Invoke declarations for registry change notification via RegNotifyChangeKeyValue. - /// Used by EspAndHelloTracker.ProvisioningStatusTracking to detect provisioning status changes instantly. - /// - internal static class RegistryWatcherNativeMethods - { - public static readonly IntPtr HKEY_LOCAL_MACHINE = new IntPtr(unchecked((int)0x80000002)); - - public const uint KEY_NOTIFY = 0x0010; - public const uint KEY_READ = 0x20019; - public const uint REG_NOTIFY_CHANGE_LAST_SET = 0x00000004; - public const uint WAIT_OBJECT_0 = 0; - public const uint WAIT_TIMEOUT = 258; - public const uint INFINITE = 0xFFFFFFFF; - - [DllImport("advapi32.dll", SetLastError = true)] - public static extern int RegOpenKeyEx( - IntPtr hKey, string subKey, uint options, uint samDesired, out IntPtr result); - - [DllImport("advapi32.dll", SetLastError = true)] - public static extern int RegNotifyChangeKeyValue( - IntPtr hKey, bool watchSubtree, uint notifyFilter, IntPtr hEvent, bool asynchronous); - - [DllImport("advapi32.dll", SetLastError = true)] - public static extern int RegCloseKey(IntPtr hKey); - - [DllImport("kernel32.dll", SetLastError = true)] - public static extern IntPtr CreateEvent( - IntPtr lpEventAttributes, bool bManualReset, bool bInitialState, string lpName); - - [DllImport("kernel32.dll", SetLastError = true)] - [return: MarshalAs(UnmanagedType.Bool)] - public static extern bool SetEvent(IntPtr hEvent); - - [DllImport("kernel32.dll", SetLastError = true)] - public static extern uint WaitForSingleObject(IntPtr hHandle, uint dwMilliseconds); - - [DllImport("kernel32.dll", SetLastError = true)] - public static extern uint WaitForMultipleObjects( - uint nCount, IntPtr[] lpHandles, bool bWaitAll, uint dwMilliseconds); - - [DllImport("kernel32.dll", SetLastError = true)] - [return: MarshalAs(UnmanagedType.Bool)] - public static extern bool CloseHandle(IntPtr hObject); - } -} diff --git a/src/Agent/AutopilotMonitor.Agent/buildcounter.txt b/src/Agent/AutopilotMonitor.Agent/buildcounter.txt index 1bb7ac53..9524ef49 100644 --- a/src/Agent/AutopilotMonitor.Agent/buildcounter.txt +++ b/src/Agent/AutopilotMonitor.Agent/buildcounter.txt @@ -1 +1 @@ -421 +424