From 49ea03a6b975c5984b9c587f5fa75aabb447f3a9 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:16:34 +0000
Subject: [PATCH] [CONSRV] Modify driver acquiring/releasing privilege
 according to r58233. The corresponding modification in services.c will be
 done at the next synchronization.

svn path=/branches/ros-csrss/; revision=58234
---
 win32ss/user/consrv/tuiconsole.c | 105 ++++++++++---------------------
 1 file changed, 32 insertions(+), 73 deletions(-)

diff --git a/win32ss/user/consrv/tuiconsole.c b/win32ss/user/consrv/tuiconsole.c
index 9669747e17e..0d22e20295b 100644
--- a/win32ss/user/consrv/tuiconsole.c
+++ b/win32ss/user/consrv/tuiconsole.c
@@ -29,62 +29,19 @@ static BOOL ConsInitialized = FALSE;
 |** BlueScreen Driver management                                             **|
 \**/
 /* Code taken and adapted from base/system/services/driver.c */
-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;
-}
-
 static DWORD
 ScmLoadDriver(LPCWSTR lpServiceName)
 {
+    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\\") */
     pszDriverPath = RtlAllocateHeap(ConSrvHeap,
-                              HEAP_ZERO_MEMORY,
-                              (52 + wcslen(lpServiceName) + 1) * sizeof(WCHAR));
+                                    HEAP_ZERO_MEMORY,
+                                    (52 + wcslen(lpServiceName) + 1) * sizeof(WCHAR));
     if (pszDriverPath == NULL)
         return ERROR_NOT_ENOUGH_MEMORY;
 
@@ -99,44 +56,44 @@ ScmLoadDriver(LPCWSTR lpServiceName)
     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("CONSRV: Cannot acquire driver-loading privilege, error = %lu\n", dwError);
+        DPRINT1("CONSRV: 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:
     RtlFreeHeap(ConSrvHeap, 0, pszDriverPath);
-
-    return dwError;
+    return RtlNtStatusToDosError(Status);
 }
 
 #ifdef BLUESCREEN_DRIVER_UNLOADING
 static DWORD
 ScmUnloadDriver(LPCWSTR lpServiceName)
 {
+    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\\") */
     pszDriverPath = RtlAllocateHeap(ConSrvHeap,
-                              HEAP_ZERO_MEMORY,
-                              (52 + wcslen(lpServiceName) + 1) * sizeof(WCHAR));
+                                    HEAP_ZERO_MEMORY,
+                                    (52 + wcslen(lpServiceName) + 1) * sizeof(WCHAR));
     if (pszDriverPath == NULL)
         return ERROR_NOT_ENOUGH_MEMORY;
 
@@ -148,29 +105,31 @@ ScmUnloadDriver(LPCWSTR lpServiceName)
     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("CONSRV: Cannot acquire driver-unloading privilege, error = %lu\n", dwError);
+        DPRINT1("CONSRV: 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:
     RtlFreeHeap(ConSrvHeap, 0, pszDriverPath);
-
-    return dwError;
+    return RtlNtStatusToDosError(Status);
 }
 #endif
 /**\