From 4ad0f2245b09f9c2cc192ecaa45773d570a17b57 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Herm=C3=A8s=20B=C3=A9lusca-Ma=C3=AFto?=
 <hermes.belusca-maito@reactos.org>
Date: Sat, 26 Jan 2013 23:14:05 +0000
Subject: [PATCH] [SERVICES] Fix acquiring/releasing driver-loading and
 unloading privilege (see r57754) by using the RtlAdjustPrivilege API instead
 of using the standard user-mode Win32 API (thus get rid of calling LSA when
 calling LookupPrivilegeValueW).

svn path=/trunk/; revision=58233
---
 reactos/base/system/services/driver.c   | 98 +++++++------------------
 reactos/base/system/services/services.h |  1 +
 2 files changed, 29 insertions(+), 70 deletions(-)

diff --git a/reactos/base/system/services/driver.c b/reactos/base/system/services/driver.c
index b8023dcedb9..8526b663211 100644
--- a/reactos/base/system/services/driver.c
+++ b/reactos/base/system/services/driver.c
@@ -16,57 +16,13 @@
 
 /* FUNCTIONS ****************************************************************/
 
-static DWORD EnablePrivilege(LPCWSTR lpszPrivilegeName, BOOL bEnablePrivilege)
-{
-    DWORD  dwRet  = ERROR_SUCCESS;
-    HANDLE hToken = NULL;
-
-    if (OpenProcessToken(GetCurrentProcess(),
-                         TOKEN_ADJUST_PRIVILEGES,
-                         &hToken))
-    {
-        TOKEN_PRIVILEGES tp;
-
-        tp.PrivilegeCount = 1;
-        tp.Privileges[0].Attributes = (bEnablePrivilege ? SE_PRIVILEGE_ENABLED : 0);
-
-        if (LookupPrivilegeValueW(NULL,
-                                  lpszPrivilegeName,
-                                  &tp.Privileges[0].Luid))
-        {
-            if (AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL))
-            {
-                if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
-                    dwRet = ERROR_NOT_ALL_ASSIGNED;
-            }
-            else
-            {
-                dwRet = GetLastError();
-            }
-        }
-        else
-        {
-            dwRet = GetLastError();
-        }
-
-        CloseHandle(hToken);
-    }
-    else
-    {
-        dwRet = GetLastError();
-    }
-
-    return dwRet;
-}
-
-
 DWORD
 ScmLoadDriver(PSERVICE lpService)
 {
+    NTSTATUS Status = STATUS_SUCCESS;
+    BOOLEAN WasPrivilegeEnabled = FALSE;
     PWSTR pszDriverPath;
     UNICODE_STRING DriverPath;
-    NTSTATUS Status;
-    DWORD dwError = ERROR_SUCCESS;
 
     /* Build the driver path */
     /* 52 = wcslen(L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\") */
@@ -87,38 +43,38 @@ ScmLoadDriver(PSERVICE lpService)
     DPRINT("  Path: %wZ\n", &DriverPath);
 
     /* Acquire driver-loading privilege */
-    dwError = EnablePrivilege(SE_LOAD_DRIVER_NAME, TRUE);
-    if (dwError != ERROR_SUCCESS)
+    Status = RtlAdjustPrivilege(SE_LOAD_DRIVER_PRIVILEGE,
+                                TRUE,
+                                FALSE,
+                                &WasPrivilegeEnabled);
+    if (!NT_SUCCESS(Status))
     {
         /* We encountered a failure, exit properly */
-        DPRINT1("SERVICES: Cannot acquire driver-loading privilege, error = %lu\n", dwError);
+        DPRINT1("SERVICES: Cannot acquire driver-loading privilege, Status = 0x%08lx\n", Status);
         goto done;
     }
 
     Status = NtLoadDriver(&DriverPath);
 
     /* Release driver-loading privilege */
-    EnablePrivilege(SE_LOAD_DRIVER_NAME, FALSE);
-
-    if (!NT_SUCCESS(Status))
-    {
-        dwError = RtlNtStatusToDosError(Status);
-    }
+    RtlAdjustPrivilege(SE_LOAD_DRIVER_PRIVILEGE,
+                       WasPrivilegeEnabled,
+                       FALSE,
+                       &WasPrivilegeEnabled);
 
 done:
     HeapFree(GetProcessHeap(), 0, pszDriverPath);
-
-    return dwError;
+    return RtlNtStatusToDosError(Status);
 }
 
 
 DWORD
 ScmUnloadDriver(PSERVICE lpService)
 {
+    NTSTATUS Status = STATUS_SUCCESS;
+    BOOLEAN WasPrivilegeEnabled = FALSE;
     PWSTR pszDriverPath;
     UNICODE_STRING DriverPath;
-    NTSTATUS Status;
-    DWORD dwError = ERROR_SUCCESS;
 
     /* Build the driver path */
     /* 52 = wcslen(L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\") */
@@ -136,29 +92,31 @@ ScmUnloadDriver(PSERVICE lpService)
     RtlInitUnicodeString(&DriverPath,
                          pszDriverPath);
 
+    DPRINT("  Path: %wZ\n", &DriverPath);
+
     /* Acquire driver-unloading privilege */
-    dwError = EnablePrivilege(SE_LOAD_DRIVER_NAME, TRUE);
-    if (dwError != ERROR_SUCCESS)
+    Status = RtlAdjustPrivilege(SE_LOAD_DRIVER_PRIVILEGE,
+                                TRUE,
+                                FALSE,
+                                &WasPrivilegeEnabled);
+    if (!NT_SUCCESS(Status))
     {
         /* We encountered a failure, exit properly */
-        DPRINT1("SERVICES: Cannot acquire driver-unloading privilege, error = %lu\n", dwError);
+        DPRINT1("SERVICES: Cannot acquire driver-unloading privilege, Status = 0x%08lx\n", Status);
         goto done;
     }
 
     Status = NtUnloadDriver(&DriverPath);
 
     /* Release driver-unloading privilege */
-    EnablePrivilege(SE_LOAD_DRIVER_NAME, FALSE);
-
-    if (!NT_SUCCESS(Status))
-    {
-        dwError = RtlNtStatusToDosError(Status);
-    }
+    RtlAdjustPrivilege(SE_LOAD_DRIVER_PRIVILEGE,
+                       WasPrivilegeEnabled,
+                       FALSE,
+                       &WasPrivilegeEnabled);
 
 done:
     HeapFree(GetProcessHeap(), 0, pszDriverPath);
-
-    return dwError;
+    return RtlNtStatusToDosError(Status);
 }
 
 
diff --git a/reactos/base/system/services/services.h b/reactos/base/system/services/services.h
index 076ea38cbd8..8e13a3e342a 100644
--- a/reactos/base/system/services/services.h
+++ b/reactos/base/system/services/services.h
@@ -14,6 +14,7 @@
 #include <ndk/iofuncs.h>
 #include <ndk/obfuncs.h>
 #include <ndk/rtlfuncs.h>
+#include <ndk/setypes.h>
 #include <services/services.h>
 #include <svcctl_s.h>