added native methods necessary for requesting computer restart permis…

…sion fixes Topshelf#193
hongweichang · Mar 26, 2015 · 64f4bb5 · 64f4bb5
1 parent 0f9e001
commit 64f4bb5
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 1 deletion.
diff --git a/src/Topshelf/Runtime/Windows/NativeMethods.cs b/src/Topshelf/Runtime/Windows/NativeMethods.cs
@@ -131,6 +131,14 @@ public enum ACCESS_MASK : uint
             WINSTA_ALL_ACCESS = 0x0000037f
         }
 
+        [Flags]
+        public enum SYSTEM_ACCESS : uint
+        {
+            SE_PRIVILEGE_ENABLED = 0x00000002,
+            TOKEN_QUERY = 0x00000008,
+            TOKEN_ADJUST_PRIVILEGES = 0x00000020
+        }
+
         [DllImport("advapi32.dll", EntryPoint = "OpenSCManagerW", ExactSpelling = true, CharSet = CharSet.Unicode,
             SetLastError = true)]
         public static extern IntPtr OpenSCManager(string machineName, string databaseName, uint dwAccess);
@@ -146,6 +154,39 @@ public enum ACCESS_MASK : uint
         public static extern bool ChangeServiceConfig2(IntPtr serviceHandle, uint infoLevel,
             IntPtr lpInfo);
 
+        [DllImport("advapi32.dll", SetLastError = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        public static extern bool OpenProcessToken(IntPtr ProcessHandle, int DesiredAccess, out IntPtr TokenHandle);
+
+        [DllImport("advapi32.dll", SetLastError = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        public static extern bool AdjustTokenPrivileges(IntPtr TokenHandle,
+            [MarshalAs(UnmanagedType.Bool)]bool DisableAllPrivileges,
+            ref TOKEN_PRIVILEGES NewState,
+            UInt32 BufferLength,
+            IntPtr PreviousState,
+            IntPtr ReturnLength);
+
+        [DllImport("advapi32.dll", SetLastError = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        public static extern bool LookupPrivilegeValue(string lpSystemName, string lpName, out LUID lpLuid);
+
+        public struct LUID
+        {
+            public int LowPart;
+            public int HighPart;
+        }
+        public struct LUID_AND_ATTRIBUTES
+        {
+            public LUID pLuid;
+            public int Attributes;
+        }
+        public struct TOKEN_PRIVILEGES
+        {
+            public int PrivilegeCount;
+            public LUID_AND_ATTRIBUTES Privileges;
+        }
+
         [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
         public struct SERVICE_FAILURE_ACTIONS
         {

diff --git a/src/Topshelf/Runtime/Windows/WindowsServiceRecoveryController.cs b/src/Topshelf/Runtime/Windows/WindowsServiceRecoveryController.cs
@@ -14,6 +14,7 @@ namespace Topshelf.Runtime.Windows
 {
     using System;
     using System.Collections.Generic;
+    using System.ComponentModel;
     using System.Linq;
     using System.Runtime.InteropServices;
     using System.Security.Permissions;
@@ -85,10 +86,14 @@ public void SetServiceRecoveryOptions(HostSettings settings, ServiceRecoveryOpti
 
                 Marshal.StructureToPtr(failureActions, lpInfo, false);
 
+                // If user specified a Restart option, get shutdown privileges
+                if(options.Actions.Any(x => x.GetType() == typeof(RestartSystemRecoveryAction)))
+                    RequestShutdownPrivileges();
+
                 if (!NativeMethods.ChangeServiceConfig2(serviceHandle,
                     NativeMethods.SERVICE_CONFIG_FAILURE_ACTIONS, lpInfo))
                 {
-                    throw new TopshelfException("Failed to change service recovery options");
+                    throw new TopshelfException(string.Format("Failed to change service recovery options. Windows Error: {0}", new Win32Exception().Message));
                 }
 
                 if (false == options.RecoverOnCrashOnly)
@@ -127,5 +132,32 @@ public void SetServiceRecoveryOptions(HostSettings settings, ServiceRecoveryOpti
                     NativeMethods.CloseServiceHandle(scmHandle);
             }
         }
+
+        private void RequestShutdownPrivileges()
+        {
+            IntPtr hToken;
+            ThrowOnFail(
+                NativeMethods.OpenProcessToken(System.Diagnostics.Process.GetCurrentProcess().Handle,
+                (int)NativeMethods.SYSTEM_ACCESS.TOKEN_ADJUST_PRIVILEGES | 
+                (int)NativeMethods.SYSTEM_ACCESS.TOKEN_QUERY, out hToken));
+
+            NativeMethods.TOKEN_PRIVILEGES tkp;
+            tkp.PrivilegeCount = 1;
+            tkp.Privileges.Attributes = (int)NativeMethods.SYSTEM_ACCESS.SE_PRIVILEGE_ENABLED;
+            const string SE_SHUTDOWN_NAME = "SeShutdownPrivilege";
+            ThrowOnFail(
+                NativeMethods.LookupPrivilegeValue("", SE_SHUTDOWN_NAME, out tkp.Privileges.pLuid));
+
+            ThrowOnFail(
+                NativeMethods.AdjustTokenPrivileges(hToken, false, ref tkp, 0U, IntPtr.Zero, IntPtr.Zero));
+        }
+
+        private static void ThrowOnFail(bool success)
+        {
+            if(!success)
+                throw new TopshelfException(string.Format(
+                    "Computer shutdown was specified as a recovery option, but privileges could not be acquired. Windows Error: {0}",
+                    new Win32Exception().Message));
+        }
     }
 }