[ekohl] 59020: [SYSSETUP] - Move set privilege code into a separate function. - Set the restore privilege to make NtUnloadKey work properly.

Author: ekohl
Date: Wed May 15 21:26:07 2013
New Revision: 59020

URL: http://svn.reactos.org/svn/reactos?rev=59020&view=rev
Log:
[SYSSETUP]
- Move set privilege code into a separate function.
- Set the restore privilege to make NtUnloadKey work properly.

Modified:
    trunk/reactos/dll/win32/syssetup/install.c

Modified: trunk/reactos/dll/win32/syssetup/install.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/dll/win32/syssetup/install…
==============================================================================
--- trunk/reactos/dll/win32/syssetup/install.c	[iso-8859-1] (original)
+++ trunk/reactos/dll/win32/syssetup/install.c	[iso-8859-1] Wed May 15 21:26:07 2013
@@ -806,12 +806,58 @@
     return TRUE;
 }
 
+
+static
+BOOL
+SetPrivilege(IN LPTSTR lpPrivilegeName,
+             IN DWORD dwAttribute)
+{
+    TOKEN_PRIVILEGES privs;
+    HANDLE hToken = NULL;
+    BOOL bResult = TRUE;
+
+    if (!OpenProcessToken(GetCurrentProcess(),
+                          TOKEN_ADJUST_PRIVILEGES,
+                          &hToken))
+    {
+        FatalError("OpenProcessToken() failed!");
+        return FALSE;
+    }
+
+    if (!LookupPrivilegeValue(NULL,
+                              lpPrivilegeName,
+                              &privs.Privileges[0].Luid))
+    {
+        FatalError("LookupPrivilegeValue() failed!");
+        bResult = FALSE;
+        goto done;
+    }
+
+    privs.PrivilegeCount = 1;
+    privs.Privileges[0].Attributes = dwAttribute;
+    if (AdjustTokenPrivileges(hToken,
+                              FALSE,
+                              &privs,
+                              0,
+                              (PTOKEN_PRIVILEGES)NULL,
+                              NULL) == 0)
+    {
+        FatalError("AdjustTokenPrivileges() failed!");
+        bResult = FALSE;
+    }
+
+done:
+    if (hToken != NULL)
+        CloseHandle(hToken);
+
+    return bResult;
+}
+
+
 DWORD WINAPI
 InstallReactOS(HINSTANCE hInstance)
 {
     TCHAR szBuffer[MAX_PATH];
-    HANDLE token;
-    TOKEN_PRIVILEGES privs;
     HKEY hKey;
 
     InitializeSetupActionLog(FALSE);
@@ -868,6 +914,8 @@
         FatalError("CreateShortcuts() failed");
         return 0;
     }
+
+    SetPrivilege(SE_RESTORE_NAME, SE_PRIVILEGE_ENABLED);
 
     /* ROS HACK, as long as NtUnloadKey is not implemented */
     {
@@ -899,39 +947,15 @@
     }
     /* END OF ROS HACK */
 
+    SetPrivilege(SE_RESTORE_NAME, 0);
+
     SetupCloseInfFile(hSysSetupInf);
     SetSetupType(0);
 
     LogItem(SYSSETUP_SEVERITY_INFORMATION, L"Installing ReactOS done");
     TerminateSetupActionLog();
 
-    /* Get shutdown privilege */
-    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token))
-    {
-        FatalError("OpenProcessToken() failed!");
-        return 0;
-    }
-    if (!LookupPrivilegeValue(
-        NULL,
-        SE_SHUTDOWN_NAME,
-        &privs.Privileges[0].Luid))
-    {
-        FatalError("LookupPrivilegeValue() failed!");
-        return 0;
-    }
-    privs.PrivilegeCount = 1;
-    privs.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
-    if (AdjustTokenPrivileges(
-        token,
-        FALSE,
-        &privs,
-        0,
-        (PTOKEN_PRIVILEGES)NULL,
-        NULL) == 0)
-    {
-        FatalError("AdjustTokenPrivileges() failed!");
-        return 0;
-    }
+    SetPrivilege(SE_SHUTDOWN_NAME, SE_PRIVILEGE_ENABLED);
 
     ExitWindowsEx(EWX_REBOOT, 0);
     return 0;