- Sync with trunk r58248 to bring the latest changes from Amine (headers) and others...

[reactos.git] / base / system / services / driver.c

diff --git a/base/system/services/driver.c b/base/system/services/driver.c
index b8023dc..8526b66 100644 (file)
--- a/base/system/services/driver.c
+++ b/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);
 }