[ros-diffs] 05/08: [USETUP] Implement work-in-progress code that allows verifying whether registry hives validity.

[Hermès Bélusca-Maïto]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=455e69a0113defa7e28a6…

commit 455e69a0113defa7e28a61ba0c16fecef5f77e77
Author:     Hermès Bélusca-Maïto <hermes.belusca-maito(a)reactos.org>
AuthorDate: Mon Jun 12 00:09:16 2017 +0000
Commit:     Hermès Bélusca-Maïto <hermes.belusca-maito(a)reactos.org>
CommitDate: Thu Oct 25 00:58:23 2018 +0200

    [USETUP] Implement work-in-progress code that allows verifying whether registry hives
validity.
    
    Implement work-in-progress code that, when upgrading an existing
    ReactOS installation, verifies whether the (existing) system registry
    hives are valid (by loading & unloading them, this allows using the
    built-in repair functionality if needed), or not.
    
    If a given hive cannot be repaired successfully, it is backed up
    (with a '.brk' extension, "brk" as "broken"), then is
marked up for
    recreation.
    
    When all hives have been checked, if there are any hive that needs
    to be recreated, we recreate its hive file, mount it, but we do *NOT*
    mount the other valid existing hives for update. We create instead
    dummy registry hives so that we can actually use, as the update code,
    the same one as the one used when fully creating the registry hives
    for a clean installation (of course, this choice can be improved later on).
    
    The update code (i.e. the same as the registry clean-install one) then
    adds the registry keys & values, either putting them in the dummy
    registry hives (the ones that we don't want to recreate) or in the
    registry hive that is recreated.
    
    At the end, the (re)created registry hives are flushed back to disk,
    and a copy of them (under a '.sav' extension) are created, so that they
    can be used for restoration purposes if 2nd-stage (and up) goes berserk.
    
    Extra fixes:
    
    - Use the correct structure member field when initializing the
      'InstallDir' variable, when performing an upgrade.
    
    - CreateNestedKey() should be better analysed to see whether it correctly
      creates the full registry path compatible with volatile/non-volatile keys
      (under inspection).
    
    svn path=/branches/setup_improvements/; revision=75010
---
 base/setup/usetup/registry.c | 556 ++++++++++++++++++++++++++++++-------------
 base/setup/usetup/registry.h |   8 +-
 base/setup/usetup/usetup.c   |  95 +++++---
 3 files changed, 454 insertions(+), 205 deletions(-)

diff --git a/base/setup/usetup/registry.c b/base/setup/usetup/registry.c
index 7fda1c928c..bb79689789 100644
--- a/base/setup/usetup/registry.c
+++ b/base/setup/usetup/registry.c
@@ -28,6 +28,9 @@
 
 #include "usetup.h"
 
+// HACK!
+#include <strsafe.h>
+
 #define NDEBUG
 #include <debug.h>
 
@@ -59,6 +62,9 @@
 
 /* FUNCTIONS ****************************************************************/
 
+#define REGISTRY_SETUP_MACHINE 
L"\\Registry\\Machine\\SYSTEM\\USetup_Machine\\"
+#define REGISTRY_SETUP_USER     L"\\Registry\\Machine\\SYSTEM\\USetup_User\\"
+
 typedef struct _ROOT_KEY
 {
     PCWSTR Name;
@@ -68,11 +74,10 @@ typedef struct _ROOT_KEY
 
 ROOT_KEY RootKeys[] =
 {
-    // L"\\Registry\\Machine\\SYSTEM\\USetup_Machine\\SOFTWARE\\Classes\\"
-    { L"HKCR", L"\\Registry\\Machine\\USetup_SOFTWARE\\Classes\\",
NULL },    /* "\\Registry\\Machine\\SOFTWARE\\Classes\\" */  //
HKEY_CLASSES_ROOT
-    { L"HKCU", L"\\Registry\\User\\USetup_DEFAULT\\"             ,
NULL },    /* "\\Registry\\User\\.DEFAULT\\" */              //
HKEY_CURRENT_USER
-    { L"HKLM", L"\\Registry\\Machine\\SYSTEM\\USetup_Machine\\"  ,
NULL },    /* "\\Registry\\Machine\\"        */              //
HKEY_LOCAL_MACHINE
-    { L"HKU" , L"\\Registry\\Machine\\SYSTEM\\USetup_User\\"     ,
NULL },    /* "\\Registry\\User\\"           */              // HKEY_USERS
+    { L"HKCR", REGISTRY_SETUP_MACHINE L"SOFTWARE\\Classes\\", NULL },
  /* "\\Registry\\Machine\\SOFTWARE\\Classes\\" */ // HKEY_CLASSES_ROOT
+    { L"HKCU", REGISTRY_SETUP_USER    L".DEFAULT\\"         , NULL },
  /* "\\Registry\\User\\.DEFAULT\\" */             // HKEY_CURRENT_USER
+    { L"HKLM", REGISTRY_SETUP_MACHINE                       , NULL },   /*
"\\Registry\\Machine\\"        */             // HKEY_LOCAL_MACHINE
+    { L"HKU" , REGISTRY_SETUP_USER                          , NULL },   /*
"\\Registry\\User\\"           */             // HKEY_USERS
 #if 0
     { L"HKR", NULL, NULL },
 #endif
@@ -401,10 +406,17 @@ do_reg_operation(HANDLE KeyHandle,
   return TRUE;
 }
 
+/*
+ * This function is similar to the one in dlls/win32/advapi32/reg/reg.c
+ * TODO: I should review both of them very carefully, because they may need
+ * some adjustments in their NtCreateKey calls, especially for CreateOptions
+ * stuff etc...
+ */
 NTSTATUS
 CreateNestedKey(PHANDLE KeyHandle,
                 ACCESS_MASK DesiredAccess,
-                POBJECT_ATTRIBUTES ObjectAttributes)
+                POBJECT_ATTRIBUTES ObjectAttributes,
+                ULONG CreateOptions)
 {
     OBJECT_ATTRIBUTES LocalObjectAttributes;
     UNICODE_STRING LocalKeyName;
@@ -419,11 +431,16 @@ CreateNestedKey(PHANDLE KeyHandle,
                          ObjectAttributes,
                          0,
                          NULL,
-                         0,
+                         CreateOptions,
                          &Disposition);
     DPRINT("NtCreateKey(%wZ) called (Status %lx)\n",
ObjectAttributes->ObjectName, Status);
     if (Status != STATUS_OBJECT_NAME_NOT_FOUND)
+    {
+        if (!NT_SUCCESS(Status))
+            DPRINT1("CreateNestedKey: NtCreateKey(%wZ) failed (Status %lx)\n",
ObjectAttributes->ObjectName, Status);
+
         return Status;
+    }
 
     /* Copy object attributes */
     RtlCopyMemory(&LocalObjectAttributes,
@@ -437,28 +454,30 @@ CreateNestedKey(PHANDLE KeyHandle,
     /* Remove the last part of the key name and try to create the key again. */
     while (Status == STATUS_OBJECT_NAME_NOT_FOUND)
     {
-        Ptr = wcsrchr (LocalKeyName.Buffer, '\\');
+        Ptr = wcsrchr(LocalKeyName.Buffer, '\\');
         if (Ptr == NULL || Ptr == LocalKeyName.Buffer)
         {
             Status = STATUS_UNSUCCESSFUL;
             break;
         }
         *Ptr = (WCHAR)0;
-        LocalKeyName.Length = wcslen (LocalKeyName.Buffer) * sizeof(WCHAR);
+        LocalKeyName.Length = wcslen(LocalKeyName.Buffer) * sizeof(WCHAR);
 
         Status = NtCreateKey(&LocalKeyHandle,
-                             KEY_ALL_ACCESS,
+                             KEY_CREATE_SUB_KEY,
                              &LocalObjectAttributes,
                              0,
                              NULL,
-                             0,
+                             REG_OPTION_NON_VOLATILE,
                              &Disposition);
         DPRINT("NtCreateKey(%wZ) called (Status %lx)\n", &LocalKeyName,
Status);
+        if (!NT_SUCCESS(Status) && Status != STATUS_OBJECT_NAME_NOT_FOUND)
+            DPRINT1("CreateNestedKey: NtCreateKey(%wZ) failed (Status %lx)\n",
LocalObjectAttributes.ObjectName, Status);
     }
 
     if (!NT_SUCCESS(Status))
     {
-        RtlFreeUnicodeString (&LocalKeyName);
+        RtlFreeUnicodeString(&LocalKeyName);
         return Status;
     }
 
@@ -474,18 +493,21 @@ CreateNestedKey(PHANDLE KeyHandle,
         NtClose(LocalKeyHandle);
 
         LocalKeyName.Buffer[LocalKeyName.Length / sizeof(WCHAR)] = L'\\';
-        LocalKeyName.Length = wcslen (LocalKeyName.Buffer) * sizeof(WCHAR);
+        LocalKeyName.Length = wcslen(LocalKeyName.Buffer) * sizeof(WCHAR);
 
         Status = NtCreateKey(&LocalKeyHandle,
                              KEY_ALL_ACCESS,
                              &LocalObjectAttributes,
                              0,
                              NULL,
-                             0,
+                             CreateOptions,
                              &Disposition);
         DPRINT("NtCreateKey(%wZ) called (Status %lx)\n", &LocalKeyName,
Status);
         if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("CreateNestedKey: NtCreateKey(%wZ) failed (Status %lx)\n",
LocalObjectAttributes.ObjectName, Status);
             break;
+        }
     }
 
     RtlFreeUnicodeString(&LocalKeyName);
@@ -560,7 +582,8 @@ registry_callback(HINF hInf, PCWSTR Section, BOOLEAN Delete)
         {
             Status = CreateNestedKey(&KeyHandle,
                                      KEY_ALL_ACCESS,
-                                     &ObjectAttributes);
+                                     &ObjectAttributes,
+                                     REG_OPTION_NON_VOLATILE);
             if (!NT_SUCCESS(Status))
             {
                 DPRINT1("CreateNestedKey(%wZ) failed (Status %lx)\n",
&Name, Status);
@@ -649,29 +672,86 @@ ImportRegistryFile(
 /*
  * Should be called under privileges
  */
-// static
-NTSTATUS
+static NTSTATUS
 CreateRegistryFile(
     IN PUNICODE_STRING InstallPath,
     IN PCWSTR RegistryKey,
-    IN HANDLE ProtoKeyHandle)
+    IN BOOLEAN IsHiveNew,
+    IN HANDLE ProtoKeyHandle
+/*
+    IN PUCHAR Descriptor,
+    IN ULONG DescriptorLength
+*/
+    )
 {
+    /* '.old' is for old valid hives, while '.brk' is for old broken
hives */
+    static PCWSTR Extensions[] = {L"old", L"brk"};
+
     NTSTATUS Status;
     HANDLE FileHandle;
     UNICODE_STRING FileName;
     OBJECT_ATTRIBUTES ObjectAttributes;
     IO_STATUS_BLOCK IoStatusBlock;
+    PCWSTR Extension;
     WCHAR PathBuffer[MAX_PATH];
+    WCHAR PathBuffer2[MAX_PATH];
 
-    /* Create the file */
     CombinePaths(PathBuffer, ARRAYSIZE(PathBuffer), 3,
                  InstallPath->Buffer, L"System32\\config", RegistryKey);
+
+    Extension = Extensions[IsHiveNew ? 0 : 1];
+
+    //
+    // FIXME: The best, actually, would be to rename (move) the existing
+    // System32\config\RegistryKey file to System32\config\RegistryKey.old,
+    // and if it already existed some System32\config\RegistryKey.old, we should
+    // first rename this one into System32\config\RegistryKey_N.old before
+    // performing the original rename.
+    //
+
+    /* Check whether the registry hive file already existed, and if so, rename it */
+    if (DoesFileExist(NULL, PathBuffer))
+    {
+        // UINT i;
+
+        DPRINT1("Registry hive '%S' already exists, rename it\n",
PathBuffer);
+
+        // i = 1;
+        /* Try first by just appending the '.old' extension */
+        StringCchPrintfW(PathBuffer2, ARRAYSIZE(PathBuffer2), L"%s.%s",
PathBuffer, Extension);
+#if 0
+        while (DoesFileExist(NULL, PathBuffer2))
+        {
+            /* An old file already exists, increments its index, but not too much */
+            if (i <= 0xFFFF)
+            {
+                /* Append '_N.old' extension */
+                StringCchPrintfW(PathBuffer2, ARRAYSIZE(PathBuffer2),
L"%s_%lu.%s", PathBuffer, i, Extension);
+                ++i;
+            }
+            else
+            {
+                /*
+                 * Too many old files exist, we will rename the file
+                 * using the name of the oldest one.
+                 */
+                StringCchPrintfW(PathBuffer2, ARRAYSIZE(PathBuffer2), L"%s.%s",
PathBuffer, Extension);
+                break;
+            }
+        }
+#endif
+
+        /* Now rename the file (force the move) */
+        Status = SetupMoveFile(PathBuffer, PathBuffer2, MOVEFILE_REPLACE_EXISTING);
+    }
+
+    /* Create the file */
     RtlInitUnicodeString(&FileName, PathBuffer);
     InitializeObjectAttributes(&ObjectAttributes,
                                &FileName,
                                OBJ_CASE_INSENSITIVE,
-                               NULL, // Could have been installpath, etc...
-                               NULL);
+                               NULL,  // Could have been installpath, etc...
+                               NULL); // Descriptor
 
     Status = NtCreateFile(&FileHandle,
                           FILE_GENERIC_WRITE,
@@ -773,18 +853,19 @@ CmpLinkKeyToHive(
 /*
  * Should be called under privileges
  */
-// static
-NTSTATUS
+static NTSTATUS
 ConnectRegistry(
     IN HKEY RootKey OPTIONAL,
+    IN PCWSTR RegMountPoint,
     // IN HANDLE RootDirectory OPTIONAL,
     IN PUNICODE_STRING InstallPath,
-    IN PCWSTR RegistryKey,
-    // IN PUCHAR Descriptor,
-    // IN ULONG DescriptorLength,
-    IN PCWSTR RegMountPoint)
+    IN PCWSTR RegistryKey
+/*
+    IN PUCHAR Descriptor,
+    IN ULONG DescriptorLength
+*/
+    )
 {
-    NTSTATUS Status;
     UNICODE_STRING KeyName, FileName;
     OBJECT_ATTRIBUTES KeyObjectAttributes;
     OBJECT_ATTRIBUTES FileObjectAttributes;
@@ -795,7 +876,7 @@ ConnectRegistry(
                                &KeyName,
                                OBJ_CASE_INSENSITIVE,
                                RootKey,
-                               NULL);
+                               NULL);   // Descriptor
 
     CombinePaths(PathBuffer, ARRAYSIZE(PathBuffer), 3,
                  InstallPath->Buffer, L"System32\\config", RegistryKey);
@@ -806,56 +887,202 @@ ConnectRegistry(
                                NULL, // RootDirectory,
                                NULL);
 
-#if 0
-    IN PCMHIVE HiveToConnect;
-    /*
-     * Add security to the root key.
-     * NOTE: One can implement this using the lpSecurityAttributes
-     * parameter of RegCreateKeyExW.
-     */
-    Status = CmiCreateSecurityKey(&HiveToConnect->Hive,
-                                  HiveToConnect->Hive.BaseBlock->RootCell,
-                                  Descriptor, DescriptorLength);
+    /* Mount the registry hive in the registry namespace */
+    return NtLoadKey(&KeyObjectAttributes, &FileObjectAttributes);
+}
+
+
+static NTSTATUS
+VerifyRegistryHive(
+    // IN HKEY RootKey OPTIONAL,
+    // // IN HANDLE RootDirectory OPTIONAL,
+    IN PUNICODE_STRING InstallPath,
+    IN PCWSTR RegistryKey /* ,
+    IN PCWSTR RegMountPoint */)
+{
+    NTSTATUS Status;
+    UNICODE_STRING KeyName;
+    OBJECT_ATTRIBUTES KeyObjectAttributes;
+
+    /* Try to mount the specified registry hive */
+    Status = ConnectRegistry(NULL,
+                             L"\\Registry\\Machine\\USetup_VerifyHive",
+                             InstallPath,
+                             RegistryKey
+                             /* NULL, 0 */);
     if (!NT_SUCCESS(Status))
-        DPRINT1("Failed to add security for root key '%S'\n", Path);
-#endif
+    {
+        DPRINT1("ConnectRegistry(%S) failed, Status 0x%08lx\n", RegistryKey,
Status);
+    }
 
-    /* Mount the registry hive in the registry namespace */
-    Status = NtLoadKey(&KeyObjectAttributes, &FileObjectAttributes);
+    DPRINT1("VerifyRegistryHive: ConnectRegistry(%S) returns Status 0x%08lx\n",
RegistryKey, Status);
+
+    //
+    // TODO: Check the Status error codes: STATUS_SUCCESS, STATUS_REGISTRY_RECOVERED,
+    // STATUS_REGISTRY_HIVE_RECOVERED, STATUS_REGISTRY_CORRUPT,
STATUS_REGISTRY_IO_FAILED,
+    // STATUS_NOT_REGISTRY_FILE, STATUS_CANNOT_LOAD_REGISTRY_FILE ;
+    //(STATUS_HIVE_UNLOADED) ; STATUS_SYSTEM_HIVE_TOO_LARGE
+    //
+
+    if (Status == STATUS_REGISTRY_HIVE_RECOVERED) // NT_SUCCESS is still FALSE in this
case!
+        DPRINT1("VerifyRegistryHive: Registry hive %S was recovered but some data
may be lost (Status 0x%08lx)\n", RegistryKey, Status);
+
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("VerifyRegistryHive: Registry hive %S is corrupted (Status
0x%08lx)\n", RegistryKey, Status);
+        return Status;
+    }
+
+    if (Status == STATUS_REGISTRY_RECOVERED)
+        DPRINT1("VerifyRegistryHive: Registry hive %S succeeded recovered (Status
0x%08lx)\n", RegistryKey, Status);
+
+    /* Unmount the hive */
+    InitializeObjectAttributes(&KeyObjectAttributes,
+                               &KeyName,
+                               OBJ_CASE_INSENSITIVE,
+                               NULL,
+                               NULL);
+
+    RtlInitUnicodeString(&KeyName,
L"\\Registry\\Machine\\USetup_VerifyHive");
+    Status = NtUnloadKey(&KeyObjectAttributes);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("NtUnloadKey(%S, %wZ) failed, Status 0x%08lx\n", RegistryKey,
&KeyName, Status);
+    }
 
     return Status;
 }
 
+
+typedef enum _HIVE_UPDATE_STATE
+{
+    Create, // Create a new hive file and save possibly existing old one with a .old
extension.
+    Repair, // Re-create a new hive file and save possibly existing old one with a .brk
extension.
+    Update  // Hive update, do not need to be recreated.
+} HIVE_UPDATE_STATE;
+
+typedef struct _HIVE_LIST_ENTRY
+{
+    PCWSTR HiveName;            // HiveFileName;
+    PCWSTR HiveRegistryPath;    // HiveRegMountPoint;
+    HANDLE PredefKeyHandle;
+    PCWSTR RegSymLink;
+    HIVE_UPDATE_STATE State;
+    // PUCHAR SecurityDescriptor;
+    // ULONG  SecurityDescriptorLength;
+} HIVE_LIST_ENTRY, *PHIVE_LIST_ENTRY;
+
+#define NUMBER_OF_STANDARD_REGISTRY_HIVES   3
+
+HIVE_LIST_ENTRY RegistryHives[/*NUMBER_OF_STANDARD_REGISTRY_HIVES*/] =
+{
+    { L"SYSTEM"  , L"\\Registry\\Machine\\USetup_SYSTEM"  ,
HKEY_LOCAL_MACHINE, L"SYSTEM"  , Create /* , SystemSecurity  ,
sizeof(SystemSecurity)   */ },
+    { L"SOFTWARE", L"\\Registry\\Machine\\USetup_SOFTWARE",
HKEY_LOCAL_MACHINE, L"SOFTWARE", Create /* , SoftwareSecurity,
sizeof(SoftwareSecurity) */ },
+    { L"DEFAULT" , L"\\Registry\\User\\USetup_DEFAULT"    ,
HKEY_USERS        , L".DEFAULT", Create /* , SystemSecurity  ,
sizeof(SystemSecurity)   */ },
+
+//  { L"BCD"     , L"\\Registry\\Machine\\USetup_BCD",
HKEY_LOCAL_MACHINE, L"BCD00000000", Create /* , BcdSecurity     ,
sizeof(BcdSecurity)      */ },
+};
+C_ASSERT(_countof(RegistryHives) == NUMBER_OF_STANDARD_REGISTRY_HIVES);
+
+#define NUMBER_OF_SECURITY_REGISTRY_HIVES   2
+
+/** These hives are created by LSASS during 2nd stage setup */
+HIVE_LIST_ENTRY SecurityRegistryHives[/*NUMBER_OF_SECURITY_REGISTRY_HIVES*/] =
+{
+    { L"SAM"     , L"\\Registry\\Machine\\USetup_SAM"     ,
HKEY_LOCAL_MACHINE, L"SAM"     , Create /* , SystemSecurity  ,
sizeof(SystemSecurity)   */ },
+    { L"SECURITY", L"\\Registry\\Machine\\USetup_SECURITY",
HKEY_LOCAL_MACHINE, L"SECURITY", Create /* , NULL            , 0                
       */ },
+};
+C_ASSERT(_countof(SecurityRegistryHives) == NUMBER_OF_SECURITY_REGISTRY_HIVES);
+
+
 NTSTATUS
-RegInitializeRegistry(
-    IN PUNICODE_STRING InstallPath)
+VerifyRegistryHives(
+    IN PUNICODE_STRING InstallPath,
+    OUT PBOOLEAN ShouldUpdateRegistry)
 {
     NTSTATUS Status;
-    HANDLE KeyHandle;
-    UNICODE_STRING KeyName;
-    OBJECT_ATTRIBUTES ObjectAttributes;
     BOOLEAN PrivilegeSet[2] = {FALSE, FALSE};
-    ULONG Disposition;
     UINT i;
-    PCWSTR RegistryKeys[] =
+
+    /* Suppose first the registry hives do not have to be updated/recreated */
+    *ShouldUpdateRegistry = FALSE;
+
+    /* Acquire restore privilege */
+    Status = RtlAdjustPrivilege(SE_RESTORE_PRIVILEGE, TRUE, FALSE,
&PrivilegeSet[0]);
+    if (!NT_SUCCESS(Status))
     {
-        L"SYSTEM",
-        L"SOFTWARE",
-        L"DEFAULT", // L".DEFAULT",
-        // L"SAM",
-        // L"SECURITY",
-        // L"BCD00000000",
-    };
+        DPRINT1("RtlAdjustPrivilege(SE_RESTORE_PRIVILEGE) failed (Status
0x%08lx)\n", Status);
+        /* Exit prematurely here.... */
+        return Status;
+    }
 
-#if 0
-    /* Initialize the current session registry */
-    Status = NtInitializeRegistry(CM_BOOT_FLAG_SETUP);
+    /* Acquire backup privilege */
+    Status = RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, TRUE, FALSE, &PrivilegeSet[1]);
     if (!NT_SUCCESS(Status))
     {
-        DPRINT1("NtInitializeRegistry() failed (Status %lx)\n", Status);
+        DPRINT1("RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE) failed (Status
0x%08lx)\n", Status);
+        RtlAdjustPrivilege(SE_RESTORE_PRIVILEGE, PrivilegeSet[0], FALSE,
&PrivilegeSet[0]);
+        /* Exit prematurely here.... */
         return Status;
     }
-#endif
+
+    for (i = 0; i < ARRAYSIZE(RegistryHives); ++i)
+    {
+        Status = VerifyRegistryHive(InstallPath, RegistryHives[i].HiveName);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Registry hive '%S' needs repair!\n",
RegistryHives[i].HiveName);
+            RegistryHives[i].State = Repair;
+            *ShouldUpdateRegistry = TRUE;
+        }
+        else
+        {
+            RegistryHives[i].State = Update;
+        }
+    }
+
+    /** These hives are created by LSASS during 2nd stage setup */
+    for (i = 0; i < ARRAYSIZE(SecurityRegistryHives); ++i)
+    {
+        Status = VerifyRegistryHive(InstallPath, SecurityRegistryHives[i].HiveName);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Registry hive '%S' needs repair!\n",
SecurityRegistryHives[i].HiveName);
+            SecurityRegistryHives[i].State = Repair;
+            /*
+             * Note that it's not the role of the 1st-stage installer to fix
+             * the security hives. This should be done at 2nd-stage installation
+             * by LSASS.
+             */
+        }
+        else
+        {
+            SecurityRegistryHives[i].State = Update;
+        }
+    }
+
+    /* Reset the status (we succeeded in checking all the hives) */
+    Status = STATUS_SUCCESS;
+
+    /* Remove restore and backup privileges */
+    RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, PrivilegeSet[1], FALSE,
&PrivilegeSet[1]);
+    RtlAdjustPrivilege(SE_RESTORE_PRIVILEGE, PrivilegeSet[0], FALSE,
&PrivilegeSet[0]);
+
+    return Status;
+}
+
+
+NTSTATUS
+RegInitializeRegistry(
+    IN PUNICODE_STRING InstallPath)
+{
+    NTSTATUS Status;
+    HANDLE KeyHandle;
+    UNICODE_STRING KeyName;
+    OBJECT_ATTRIBUTES ObjectAttributes;
+    BOOLEAN PrivilegeSet[2] = {FALSE, FALSE};
+    ULONG Disposition;
+    UINT i;
 
     /* Acquire restore privilege */
     Status = RtlAdjustPrivilege(SE_RESTORE_PRIVILEGE, TRUE, FALSE,
&PrivilegeSet[0]);
@@ -885,8 +1112,7 @@ RegInitializeRegistry(
      * See
https://github.com/libyal/winreg-kb/blob/master/documentation/Registry%20fi…
      * for more information.
      */
-    RtlInitUnicodeString(&KeyName,
-                         L"\\Registry\\Machine\\SYSTEM\\$$$PROTO.HIV");
+    RtlInitUnicodeString(&KeyName,
L"\\Registry\\Machine\\SYSTEM\\$$$PROTO.HIV");
     InitializeObjectAttributes(&ObjectAttributes,
                                &KeyName,
                                OBJ_CASE_INSENSITIVE,
@@ -906,14 +1132,18 @@ RegInitializeRegistry(
     }
     NtFlushKey(KeyHandle);
 
-    for (i = 0; i < ARRAYSIZE(RegistryKeys); ++i)
+    for (i = 0; i < ARRAYSIZE(RegistryHives); ++i)
     {
+        if (RegistryHives[i].State != Create && RegistryHives[i].State !=
Repair)
+            continue;
+
         Status = CreateRegistryFile(InstallPath,
-                                    RegistryKeys[i],
+                                    RegistryHives[i].HiveName,
+                                    RegistryHives[i].State != Repair, //
RegistryHives[i].State == Create,
                                     KeyHandle);
         if (!NT_SUCCESS(Status))
         {
-            DPRINT1("CreateRegistryFile(%S) failed, Status 0x%08lx\n",
RegistryKeys[i], Status);
+            DPRINT1("CreateRegistryFile(%S) failed, Status 0x%08lx\n",
RegistryHives[i].HiveName, Status);
             /* Exit prematurely here.... */
             /* That is now done, clean everything up! */
             NtDeleteKey(KeyHandle);
@@ -928,7 +1158,7 @@ RegInitializeRegistry(
 
 
     /*
-     * Prepare the installation roots. Since we cannot create real registry keys
+     * Prepare the registry root keys. Since we cannot create real registry keys
      * inside the master keys (\Registry, \Registry\Machine or \Registry\User),
      * we need to perform some SymLink tricks instead.
      */
@@ -946,7 +1176,7 @@ RegInitializeRegistry(
                          &ObjectAttributes,
                          0,
                          NULL,
-                         REG_OPTION_VOLATILE,
+                         REG_OPTION_NON_VOLATILE, // REG_OPTION_VOLATILE, // FIXME!
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
@@ -968,7 +1198,7 @@ RegInitializeRegistry(
                          &ObjectAttributes,
                          0,
                          NULL,
-                         REG_OPTION_VOLATILE,
+                         REG_OPTION_NON_VOLATILE, // REG_OPTION_VOLATILE, // FIXME!
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
@@ -978,72 +1208,61 @@ RegInitializeRegistry(
     RootKeys[GetPredefKeyIndex(HKEY_USERS)].Handle = KeyHandle;
 
 
-
     /*
      * Now properly mount the offline hive files
      */
-
-    /* Create SYSTEM key */
-    Status =
-    ConnectRegistry(NULL,
-                    InstallPath,
-                    RegistryKeys[0],
-                    // SystemSecurity, sizeof(SystemSecurity),
-                    L"\\Registry\\Machine\\USetup_SYSTEM");
-    if (!NT_SUCCESS(Status))
-    {
-        DPRINT1("ConnectRegistry(SYSTEM) failed, Status 0x%08lx\n", Status);
-    }
-
-    /* Create the 'HKLM\SYSTEM' symlink to this key */
-    if (!CmpLinkKeyToHive(RootKeys[GetPredefKeyIndex(HKEY_LOCAL_MACHINE)].Handle,
-                          L"SYSTEM",
-                          L"\\Registry\\Machine\\USetup_SYSTEM"))
-    {
-        DPRINT1("CmpLinkKeyToHive(SYSTEM) failed!\n");
-    }
-
-
-    /* Create SOFTWARE key */
-    Status =
-    ConnectRegistry(NULL,
-                    InstallPath,
-                    RegistryKeys[1],
-                    // SoftwareSecurity, sizeof(SoftwareSecurity),
-                    L"\\Registry\\Machine\\USetup_SOFTWARE");
-    if (!NT_SUCCESS(Status))
-    {
-        DPRINT1("ConnectRegistry(SOFTWARE) failed, Status 0x%08lx\n", Status);
-    }
-
-    /* Create the 'HKLM\Software' symlink to this key */
-    if (!CmpLinkKeyToHive(RootKeys[GetPredefKeyIndex(HKEY_LOCAL_MACHINE)].Handle,
-                          L"Software",
-                          L"\\Registry\\Machine\\USetup_SOFTWARE"))
+    for (i = 0; i < ARRAYSIZE(RegistryHives); ++i)
     {
-        DPRINT1("CmpLinkKeyToHive(SOFTWARE) failed!\n");
-    }
+        // if (RegistryHives[i].State != Create && RegistryHives[i].State !=
Repair)
+            // continue;
 
+        if (RegistryHives[i].State == Create || RegistryHives[i].State == Repair)
+        {
+            Status = ConnectRegistry(NULL,
+                                     RegistryHives[i].HiveRegistryPath,
+                                     InstallPath,
+                                     RegistryHives[i].HiveName
+                                     /* SystemSecurity, sizeof(SystemSecurity) */);
+            if (!NT_SUCCESS(Status))
+            {
+                DPRINT1("ConnectRegistry(%S) failed, Status 0x%08lx\n",
RegistryHives[i].HiveName, Status);
+            }
 
-    /* Create DEFAULT key */
-    Status =
-    ConnectRegistry(NULL,
-                    InstallPath,
-                    RegistryKeys[2],
-                    // SystemSecurity, sizeof(SystemSecurity),
-                    L"\\Registry\\User\\USetup_DEFAULT");
-    if (!NT_SUCCESS(Status))
-    {
-        DPRINT1("ConnectRegistry(DEFAULT) failed, Status 0x%08lx\n", Status);
+            /* Create the registry symlink to this key */
+            if
(!CmpLinkKeyToHive(RootKeys[GetPredefKeyIndex(RegistryHives[i].PredefKeyHandle)].Handle,
+                                  RegistryHives[i].RegSymLink,
+                                  RegistryHives[i].HiveRegistryPath))
+            {
+                DPRINT1("CmpLinkKeyToHive(%S) failed!\n",
RegistryHives[i].HiveName);
+            }
+        }
+        else
+        {
+            /* Create *DUMMY* volatile hives just to make the update procedure work */
+
+            RtlInitUnicodeString(&KeyName, RegistryHives[i].RegSymLink);
+            InitializeObjectAttributes(&ObjectAttributes,
+                                       &KeyName,
+                                       OBJ_CASE_INSENSITIVE,
+                                      
RootKeys[GetPredefKeyIndex(RegistryHives[i].PredefKeyHandle)].Handle,
+                                       NULL);
+            KeyHandle = NULL;
+            Status = NtCreateKey(&KeyHandle,
+                                 KEY_ALL_ACCESS,
+                                 &ObjectAttributes,
+                                 0,
+                                 NULL,
+                                 REG_OPTION_NON_VOLATILE, // REG_OPTION_VOLATILE, //
FIXME!
+                                 &Disposition);
+            if (!NT_SUCCESS(Status))
+            {
+                DPRINT1("NtCreateKey(%wZ) failed (Status 0x%08lx)\n",
&KeyName, Status);
+                // return Status;
+            }
+            NtClose(KeyHandle);
+        }
     }
 
-    /* Create the 'HKU\.DEFAULT' symlink to this key */
-    if (!CmpLinkKeyToHive(RootKeys[GetPredefKeyIndex(HKEY_USERS)].Handle,
-                          L".DEFAULT",
-                          L"\\Registry\\User\\USetup_DEFAULT"))
-    {
-        DPRINT1("CmpLinkKeyToHive(DEFAULT) failed!\n");
-    }
 
     /* HKCU is a handle to 'HKU\.DEFAULT' */
 #if 0
@@ -1089,12 +1308,13 @@ RegInitializeRegistry(
                                NULL);
 #endif
     KeyHandle = NULL;
+    /* We use NtCreateKey instead of NtOpenKey because Software\Classes doesn't exist
originally */
     Status = NtCreateKey(&KeyHandle,
                          KEY_ALL_ACCESS,
                          &ObjectAttributes,
                          0,
                          NULL,
-                         0,
+                         REG_OPTION_NON_VOLATILE,
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
@@ -1109,31 +1329,11 @@ RegInitializeRegistry(
     RootKeys[GetPredefKeyIndex(HKEY_CLASSES_ROOT)].Handle = KeyHandle;
 
 
-#if 0
-    /* Create SAM key */
-    ConnectRegistry(NULL,
-                    &SamHive,
-                    // SystemSecurity, sizeof(SystemSecurity),
-                    L"\\Registry\\Machine\\USetup_SAM");
-
-    /* Create SECURITY key */
-    ConnectRegistry(NULL,
-                    &SecurityHive,
-                    // NULL, 0,
-                    L"\\Registry\\Machine\\USetup_SECURITY");
-
-    /* Create BCD key */
-    ConnectRegistry(NULL,
-                    &BcdHive,
-                    // BcdSecurity, sizeof(BcdSecurity),
-                    L"\\Registry\\Machine\\USetup_BCD00000000");
-#endif
-
     Status = STATUS_SUCCESS;
 
 
     /* Create the 'HKLM\SYSTEM\ControlSet001' key */
-    // L"\\Registry\\Machine\\SYSTEM\\USetup_Machine\\SYSTEM\\ControlSet001"
+    // REGISTRY_SETUP_MACHINE L"SYSTEM\\ControlSet001"
     RtlInitUnicodeString(&KeyName, L"SYSTEM\\ControlSet001");
     InitializeObjectAttributes(&ObjectAttributes,
                                &KeyName,
@@ -1163,12 +1363,15 @@ RegInitializeRegistry(
     /* Create the 'HKLM\SYSTEM\CurrentControlSet' symlink */
     if (!CmpLinkKeyToHive(RootKeys[GetPredefKeyIndex(HKEY_LOCAL_MACHINE)].Handle,
                           L"SYSTEM\\CurrentControlSet",
-                         
L"\\Registry\\Machine\\SYSTEM\\USetup_Machine\\SYSTEM\\ControlSet001"))
+                          REGISTRY_SETUP_MACHINE L"SYSTEM\\ControlSet001"))
     {
         DPRINT1("CmpLinkKeyToHive(CurrentControlSet) failed!\n");
     }
 
 
+    Status = STATUS_SUCCESS;
+
+
 Quit:
     /* Remove restore and backup privileges */
     RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, PrivilegeSet[1], FALSE,
&PrivilegeSet[1]);
@@ -1178,18 +1381,22 @@ Quit:
 }
 
 VOID
-RegCleanupRegistry(VOID)
+RegCleanupRegistry(
+    IN PUNICODE_STRING InstallPath)
 {
     NTSTATUS Status;
     UNICODE_STRING KeyName;
     OBJECT_ATTRIBUTES KeyObjectAttributes;
     BOOLEAN PrivilegeSet[2] = {FALSE, FALSE};
-    UCHAR i;
+    UINT i;
+    WCHAR SrcPath[MAX_PATH];
+    WCHAR DstPath[MAX_PATH];
 
     for (i = 0; i < ARRAYSIZE(RootKeys); ++i)
     {
         if (RootKeys[i].Handle)
         {
+            NtFlushKey(RootKeys[i].Handle);
             NtClose(RootKeys[i].Handle);
             RootKeys[i].Handle = NULL;
         }
@@ -1220,25 +1427,38 @@ RegCleanupRegistry(VOID)
                                NULL,
                                NULL);
 
-    RtlInitUnicodeString(&KeyName, L"\\Registry\\Machine\\USetup_SYSTEM");
-    Status = NtUnloadKey(&KeyObjectAttributes);
-
-    RtlInitUnicodeString(&KeyName,
L"\\Registry\\Machine\\USetup_SOFTWARE");
-    Status = NtUnloadKey(&KeyObjectAttributes);
+    for (i = 0; i < ARRAYSIZE(RegistryHives); ++i)
+    {
+        if (RegistryHives[i].State != Create && RegistryHives[i].State !=
Repair)
+            continue;
 
-    RtlInitUnicodeString(&KeyName, L"\\Registry\\User\\USetup_DEFAULT");
-    Status = NtUnloadKey(&KeyObjectAttributes);
+        RtlInitUnicodeString(&KeyName, RegistryHives[i].HiveRegistryPath);
+        Status = NtUnloadKey(&KeyObjectAttributes);
+        DPRINT1("Unmounting '%S' %s\n",
RegistryHives[i].HiveRegistryPath, NT_SUCCESS(Status) ? "succeeded" :
"failed");
+    }
 
-#if 0
-    RtlInitUnicodeString(&KeyName, L"\\Registry\\Machine\\USetup_SAM");
-    Status = NtUnloadKey(&KeyObjectAttributes);
+    //
+    // RegBackupRegistry()
+    //
+    /* Now backup the hives into .sav files */
+    for (i = 0; i < ARRAYSIZE(RegistryHives); ++i)
+    {
+        if (RegistryHives[i].State != Create && RegistryHives[i].State !=
Repair)
+            continue;
 
-    RtlInitUnicodeString(&KeyName,
L"\\Registry\\Machine\\USetup_SECURITY");
-    Status = NtUnloadKey(&KeyObjectAttributes);
+        CombinePaths(SrcPath, ARRAYSIZE(SrcPath), 3,
+                     InstallPath->Buffer, L"System32\\config",
RegistryHives[i].HiveName);
+        StringCchCopyW(DstPath, ARRAYSIZE(DstPath), SrcPath);
+        StringCchCatW(DstPath, ARRAYSIZE(DstPath), L".sav");
 
-    RtlInitUnicodeString(&KeyName,
L"\\Registry\\Machine\\USetup_BCD00000000");
-    Status = NtUnloadKey(&KeyObjectAttributes);
-#endif
+        DPRINT1("Copy hive: %S ==> %S\n", SrcPath, DstPath);
+        Status = SetupCopyFile(SrcPath, DstPath, FALSE);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("SetupCopyFile() failed (Status %lx)\n", Status);
+            // return Status;
+        }
+    }
 
     /* Remove restore and backup privileges */
     RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, PrivilegeSet[1], FALSE,
&PrivilegeSet[1]);
diff --git a/base/setup/usetup/registry.h b/base/setup/usetup/registry.h
index c97d897a78..cae50e8271 100644
--- a/base/setup/usetup/registry.h
+++ b/base/setup/usetup/registry.h
@@ -43,12 +43,18 @@ ImportRegistryFile(
     LCID LocaleId,
     BOOLEAN Delete);
 
+NTSTATUS
+VerifyRegistryHives(
+    IN PUNICODE_STRING InstallPath,
+    OUT PBOOLEAN ShouldUpdateRegistry);
+
 NTSTATUS
 RegInitializeRegistry(
     IN PUNICODE_STRING InstallPath);
 
 VOID
-RegCleanupRegistry(VOID);
+RegCleanupRegistry(
+    IN PUNICODE_STRING InstallPath);
 
 VOID
 SetDefaultPagefile(
diff --git a/base/setup/usetup/usetup.c b/base/setup/usetup/usetup.c
index 4eb2558fdf..9318f7e250 100644
--- a/base/setup/usetup/usetup.c
+++ b/base/setup/usetup/usetup.c
@@ -551,6 +551,7 @@ CheckUnattendedSetup(VOID)
     wcscpy(UnattendInstallationDirectory, Value);
 
     IsUnattendedSetup = TRUE;
+    DPRINT("Running unattended setup\n");
 
     /* Search for 'MBRInstallType' in the 'Unattend' section */
     if (SetupFindFirstLineW(UnattendInf, L"Unattend",
L"MBRInstallType", &Context))
@@ -589,8 +590,6 @@ CheckUnattendedSetup(VOID)
     }
 
     SetupCloseInfFile(UnattendInf);
-
-    DPRINT("Running unattended setup\n");
 }
 
 
@@ -1503,7 +1502,7 @@ IsDiskSizeValid(PPARTENTRY PartEntry)
 
     if (size < RequiredPartitionDiskSpace)
     {
-        /* partition is too small so ask for another partition */
+        /* Partition is too small so ask for another one */
         DPRINT1("Partition is too small (size: %I64u MB), required disk space is %lu
MB\n", size, RequiredPartitionDiskSpace);
         return FALSE;
     }
@@ -3291,6 +3290,22 @@ InstallDirectoryPage(PINPUT_RECORD Ir)
     DiskEntry = PartitionList->CurrentDisk;
     PartEntry = PartitionList->CurrentPartition;
 
+#if 0
+    if (RepairUpdateFlag)
+    {
+        if (!IsValidPath(CurrentInstallation->PathComponent)) // SystemNtPath
+        {
+            /* FIXME: Log the error? */
+            return QUIT_PAGE;
+        }
+
+        BuildInstallPaths(CurrentInstallation->PathComponent, // SystemNtPath
+                          DiskEntry,
+                          PartEntry);
+
+        return PREPARE_COPY_PAGE;
+    }
+#endif
     if (IsUnattendedSetup)
     {
         if (!IsValidPath(UnattendInstallationDirectory))
@@ -4057,20 +4072,34 @@ FileCopyPage(PINPUT_RECORD Ir)
 static PAGE_NUMBER
 RegistryPage(PINPUT_RECORD Ir)
 {
+    NTSTATUS Status;
     INFCONTEXT InfContext;
     PWSTR Action;
     PWSTR File;
     PWSTR Section;
     BOOLEAN Delete;
-    NTSTATUS Status;
 
     MUIDisplayPage(REGISTRY_PAGE);
 
     if (RepairUpdateFlag)
     {
-        // FIXME!
-        DPRINT1("FIXME: Updating / repairing the registry is NOT implemented
yet!\n");
-        return SUCCESS_PAGE;
+        BOOLEAN ShouldUpdateRegistry = FALSE;
+
+        DPRINT1("TODO: Updating / repairing the registry is not completely
implemented yet!\n");
+
+        /* Verify the registry hives and check whether we need to update or repair any of
them */
+        Status = VerifyRegistryHives(&DestinationPath, &ShouldUpdateRegistry);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("VerifyRegistryHives failed, Status 0x%08lx\n", Status);
+            ShouldUpdateRegistry = FALSE;
+        }
+        if (!ShouldUpdateRegistry)
+        {
+            DPRINT1("No need to update the registry\n");
+            // return SUCCESS_PAGE;
+            goto Quit;
+        }
     }
 
     /* Initialize the registry and setup the default installation hives */
@@ -4087,7 +4116,7 @@ RegistryPage(PINPUT_RECORD Ir)
         else
         /*************************************/
         {
-            /* Something else (correct) failed */
+            /* Something else failed */
             MUIDisplayError(ERROR_CREATE_HIVE, Ir, POPUP_WAIT_ENTER);
         }
         return QUIT_PAGE;
@@ -4099,9 +4128,8 @@ RegistryPage(PINPUT_RECORD Ir)
     if (!SetupFindFirstLineW(SetupInf, L"HiveInfs.Install", NULL,
&InfContext))
     {
         DPRINT1("SetupFindFirstLine() failed\n");
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_FIND_REGISTRY, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     do
@@ -4116,27 +4144,19 @@ RegistryPage(PINPUT_RECORD Ir)
             break; // Hackfix
 
         if (!_wcsicmp(Action, L"AddReg"))
-        {
             Delete = FALSE;
-        }
         else if (!_wcsicmp(Action, L"DelReg"))
-        {
             Delete = TRUE;
-        }
         else
-        {
             continue;
-        }
 
         CONSOLE_SetStatusText(MUIGetString(STRING_IMPORTFILE), File);
 
         if (!ImportRegistryFile(File, Section, LanguageId, Delete))
         {
             DPRINT1("Importing %S failed\n", File);
-
-            RegCleanupRegistry();
             MUIDisplayError(ERROR_IMPORT_HIVE, Ir, POPUP_WAIT_ENTER);
-            return QUIT_PAGE;
+            goto Cleanup;
         }
     } while (SetupFindNextLine(&InfContext, &InfContext));
 
@@ -4144,35 +4164,31 @@ RegistryPage(PINPUT_RECORD Ir)
     CONSOLE_SetStatusText(MUIGetString(STRING_DISPLAYETTINGSUPDATE));
     if (!ProcessDisplayRegistry(SetupInf, DisplayList))
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_UPDATE_DISPLAY_SETTINGS, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     /* Set the locale */
     CONSOLE_SetStatusText(MUIGetString(STRING_LOCALESETTINGSUPDATE));
     if (!ProcessLocaleRegistry(LanguageList))
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_UPDATE_LOCALESETTINGS, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     /* Add keyboard layouts */
     CONSOLE_SetStatusText(MUIGetString(STRING_ADDKBLAYOUTS));
     if (!AddKeyboardLayouts())
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_ADDING_KBLAYOUTS, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     /* Set GeoID */
     if (!SetGeoID(MUIGetGeoID()))
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_UPDATE_GEOID, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     if (!IsUnattendedSetup)
@@ -4181,9 +4197,8 @@ RegistryPage(PINPUT_RECORD Ir)
         CONSOLE_SetStatusText(MUIGetString(STRING_KEYBOARDSETTINGSUPDATE));
         if (!ProcessKeyboardLayoutRegistry(LayoutList))
         {
-            RegCleanupRegistry();
             MUIDisplayError(ERROR_UPDATE_KBSETTINGS, Ir, POPUP_WAIT_ENTER);
-            return QUIT_PAGE;
+            goto Cleanup;
         }
     }
 
@@ -4191,26 +4206,34 @@ RegistryPage(PINPUT_RECORD Ir)
     CONSOLE_SetStatusText(MUIGetString(STRING_CODEPAGEINFOUPDATE));
     if (!AddCodePage())
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_ADDING_CODEPAGE, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     /* Set the default pagefile entry */
     SetDefaultPagefile(DestinationDriveLetter);
 
     /* Update the mounted devices list */
+    // FIXME: This should technically be done by mountmgr (if AutoMount is enabled)!
     SetMountedDeviceValues(PartitionList);
 
+Cleanup:
     //
     // TODO: Unload all the registry stuff, perform cleanup,
-    // and copy the created hive files into .sav ones.
+    // and copy the created hive files into .sav files.
     //
-    RegCleanupRegistry();
+    RegCleanupRegistry(&DestinationPath);
 
-    CONSOLE_SetStatusText(MUIGetString(STRING_DONE));
-
-    return BOOT_LOADER_PAGE;
+Quit:
+    if (NT_SUCCESS(Status))
+    {
+        CONSOLE_SetStatusText(MUIGetString(STRING_DONE));
+        return BOOT_LOADER_PAGE;
+    }
+    else
+    {
+        return QUIT_PAGE;
+    }
 }