| | 1 | | using System; |
| | 2 | | using System.Collections.Generic; |
| | 3 | | using System.ComponentModel; |
| | 4 | | using System.Diagnostics; |
| | 5 | | using System.Diagnostics.CodeAnalysis; |
| | 6 | | using System.IO; |
| | 7 | | using System.Linq; |
| | 8 | | using System.Text; |
| | 9 | |
|
| | 10 | | namespace LockCheck.Windows; |
| | 11 | |
|
| | 12 | | internal static class RestartManager |
| | 13 | | { |
| | 14 | | public static HashSet<ProcessInfo> GetLockingProcessInfos(string[] paths, |
| | 15 | | #if NET |
| | 16 | | [NotNullIfNotNull(nameof(directories))] |
| | 17 | | #endif |
| | 18 | | ref List<string>? directories) |
| | 19 | | { |
| 2 | 20 | | if (paths == null) |
| 2 | 21 | | throw new ArgumentNullException(nameof(paths)); |
| | 22 | |
|
| | 23 | | const int maxRetries = 6; |
| | 24 | |
|
| | 25 | | // See http://blogs.msdn.com/b/oldnewthing/archive/2012/02/17/10268840.aspx. |
| 2 | 26 | | var key = new StringBuilder(new string('\0', NativeMethods.CCH_RM_SESSION_KEY + 1)); |
| | 27 | |
|
| | 28 | | uint handle; |
| 2 | 29 | | int res = NativeMethods.RmStartSession(out handle, 0, key); |
| 2 | 30 | | if (res != 0) |
| 0 | 31 | | throw GetException(res, "Failed to begin restart manager session"); |
| | 32 | |
|
| | 33 | | try |
| | 34 | | { |
| 2 | 35 | | var files = new HashSet<string>(paths.Length, StringComparer.OrdinalIgnoreCase); |
| 2 | 36 | | foreach (string path in paths) |
| | 37 | | { |
| 2 | 38 | | if (Directory.Exists(path)) |
| | 39 | | { |
| 2 | 40 | | directories?.Add(path); |
| | 41 | | } |
| | 42 | | else |
| | 43 | | { |
| 2 | 44 | | files.Add(path); |
| | 45 | | } |
| | 46 | | } |
| | 47 | |
|
| 2 | 48 | | res = NativeMethods.RmRegisterResources(handle, (uint)files.Count, files.ToArray(), 0, null, 0, null); |
| 2 | 49 | | if (res != 0) |
| 0 | 50 | | throw GetException(res, "Could not register resources"); |
| | 51 | |
|
| | 52 | | // |
| | 53 | | // Obtain the list of affected applications/services. |
| | 54 | | // |
| | 55 | | // NOTE: Restart Manager returns the results into the buffer allocated by the caller. The first call to |
| | 56 | | // RmGetList() will return the size of the buffer (i.e. nProcInfoNeeded) the caller needs to allocate. |
| | 57 | | // The caller then needs to allocate the buffer (i.e. rgAffectedApps) and make another RmGetList() |
| | 58 | | // call to ask Restart Manager to write the results into the buffer. However, since Restart Manager |
| | 59 | | // refreshes the list every time RmGetList()is called, it is possible that the size returned by the first |
| | 60 | | // RmGetList()call is not sufficient to hold the results discovered by the second RmGetList() call. Therefor |
| | 61 | | // it is recommended that the caller follows the following practice to handle this race condition: |
| | 62 | | // |
| | 63 | | // Use a loop to call RmGetList() in case the buffer allocated according to the size returned in previous |
| | 64 | | // call is not enough. |
| | 65 | | // |
| 2 | 66 | | uint pnProcInfo = 0; |
| 2 | 67 | | NativeMethods.RM_PROCESS_INFO[]? rgAffectedApps = null; |
| 2 | 68 | | int retry = 0; |
| | 69 | | do |
| | 70 | | { |
| 2 | 71 | | uint lpdwRebootReasons = (uint)NativeMethods.RM_REBOOT_REASON.RmRebootReasonNone; |
| | 72 | | uint pnProcInfoNeeded; |
| 2 | 73 | | res = NativeMethods.RmGetList(handle, out pnProcInfoNeeded, ref pnProcInfo, rgAffectedApps, ref lpdwRebo |
| 2 | 74 | | if (res == 0) |
| | 75 | | { |
| | 76 | | // If pnProcInfo == 0, then there is simply no locking process (found), in this case rgAffectedApps |
| 2 | 77 | | if (pnProcInfo == 0) |
| 2 | 78 | | return []; |
| | 79 | |
|
| | 80 | | Debug.Assert(rgAffectedApps != null); |
| 2 | 81 | | var lockInfos = new HashSet<ProcessInfo>((int)pnProcInfo); |
| 2 | 82 | | for (int i = 0; i < pnProcInfo; i++) |
| | 83 | | { |
| 2 | 84 | | var info = ProcessInfoWindows.Create(rgAffectedApps![i]); |
| 2 | 85 | | if (info != null) |
| | 86 | | { |
| 2 | 87 | | lockInfos.Add(info); |
| | 88 | | } |
| | 89 | | } |
| 2 | 90 | | return lockInfos; |
| | 91 | | } |
| | 92 | |
|
| 2 | 93 | | if (res != NativeMethods.ERROR_MORE_DATA) |
| 0 | 94 | | throw GetException(res, $"Failed to get entries (retry {retry})"); |
| | 95 | |
|
| 2 | 96 | | pnProcInfo = pnProcInfoNeeded; |
| 2 | 97 | | rgAffectedApps = new NativeMethods.RM_PROCESS_INFO[pnProcInfo]; |
| 2 | 98 | | } while ((res == NativeMethods.ERROR_MORE_DATA) && (retry++ < maxRetries)); |
| 0 | 99 | | } |
| | 100 | | finally |
| | 101 | | { |
| 2 | 102 | | res = NativeMethods.RmEndSession(handle); |
| 2 | 103 | | if (res != 0) |
| 0 | 104 | | throw GetException(res, "Failed to end the restart manager session"); |
| 2 | 105 | | } |
| | 106 | |
|
| 0 | 107 | | return []; |
| 2 | 108 | | } |
| | 109 | |
|
| | 110 | | internal static Win32Exception GetException(int res, string message) |
| | 111 | | { |
| 0 | 112 | | return new Win32Exception(res, $"{message}: {NativeMethods.GetMessage(res)}"); |
| | 113 | | } |
| | 114 | | } |