RestartManagerSession.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;

namespace RestartManager
{
    public sealed class RestartManagerSession : IDisposable
    {
        #region Interop
        [DllImport("rstrtmgr.dll", CharSet = CharSet.Auto)]
        static extern int RmStartSession(out uint pSessionHandle, int dwSessionFlags, string strSessionKey);

        [DllImport("rstrtmgr.dll")]
        static extern int RmEndSession(uint pSessionHandle);

        [StructLayout(LayoutKind.Sequential)]
        struct RM_UNIQUE_PROCESS
        {
            public int dwProcessId;
            public System.Runtime.InteropServices.ComTypes.FILETIME ProcessStartTime;
        }

        const int RmRebootReasonNone = 0;
        const int CCH_RM_MAX_APP_NAME = 255;
        const int CCH_RM_MAX_SVC_NAME = 63;

        enum RM_APP_TYPE
        {
            RmUnknownApp = 0,
            RmMainWindow = 1,
            RmOtherWindow = 2,
            RmService = 3,
            RmExplorer = 4,
            RmConsole = 5,
            RmCritical = 1000
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        struct RM_PROCESS_INFO
        {
            public RM_UNIQUE_PROCESS Process;

            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = CCH_RM_MAX_APP_NAME + 1)]
            public string strAppName;

            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = CCH_RM_MAX_SVC_NAME + 1)]
            public string strServiceShortName;

            public RM_APP_TYPE ApplicationType;
            public uint AppStatus;
            public uint TSSessionId;
            [MarshalAs(UnmanagedType.Bool)]
            public bool bRestartable;
        }

        [DllImport("rstrtmgr.dll", CharSet = CharSet.Unicode)]
        static extern int RmRegisterResources(uint pSessionHandle,
                                              UInt32 nFiles,
                                              string[] rgsFilenames,
                                              UInt32 nApplications,
                                              [In] RM_UNIQUE_PROCESS[] rgApplications,
                                              UInt32 nServices,
                                              string[] rgsServiceNames);

        [DllImport("rstrtmgr.dll")]
        static extern int RmGetList(uint dwSessionHandle,
                                    out uint pnProcInfoNeeded,
                                    ref uint pnProcInfo,
                                    [In, Out] RM_PROCESS_INFO[] rgAffectedApps,
                                    ref uint lpdwRebootReasons); 
        #endregion

        private readonly string _key = Guid.NewGuid().ToString();
        private readonly uint _handle;
        private readonly RM_PROCESS_INFO[] _processesInfo = new RM_PROCESS_INFO[64];

        public RestartManagerSession()
        {
            if (RmStartSession(out _handle, 0, _key) != 0)
                throw new Exception("Could not begin restart session.  Unable to determine file locker");
        }

        public void RegisterResources(string[] resources)
        {
            if(RmRegisterResources(_handle, (uint)resources.Length, resources, 0, null, 0, null) != 0)
                throw new Exception("Could not register resource");
        }

        public IEnumerable<Process> GetProcesses()
        {
            return GetProcessesIds().Select(Process.GetProcessById);
        }

        public int[] GetProcessesIds()
        {
            uint pnProcInfoNeeded = 0, pnProcInfo = (uint)_processesInfo.Length, lpdwRebootReasons = RmRebootReasonNone;

            // Get the list
            if (RmGetList(_handle, out pnProcInfoNeeded, ref pnProcInfo, _processesInfo, ref lpdwRebootReasons) != 0)
                throw new Exception("Could not list processes locking resource");

            return _processesInfo.Select(p => p.Process.dwProcessId).Take((int)pnProcInfoNeeded).ToArray();
        }

        public void Dispose()
        {
            if(RmEndSession(_handle) != 0)
                throw new Exception("Could not release restart api session handle");
        }
    }
}