[USETUP] Implement work-in-progress code that allows verifying whether registry hives...
authorHermès Bélusca-Maïto <hermes.belusca-maito@reactos.org>
Mon, 12 Jun 2017 00:09:16 +0000 (00:09 +0000)
committerHermès Bélusca-Maïto <hermes.belusca-maito@reactos.org>
Wed, 24 Oct 2018 22:58:23 +0000 (00:58 +0200)
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
base/setup/usetup/registry.h
base/setup/usetup/usetup.c

index 7fda1c9..bb79689 100644 (file)
@@ -28,6 +28,9 @@
 
 #include "usetup.h"
 
 
 #include "usetup.h"
 
+// HACK!
+#include <strsafe.h>
+
 #define NDEBUG
 #include <debug.h>
 
 #define NDEBUG
 #include <debug.h>
 
@@ -59,6 +62,9 @@
 
 /* FUNCTIONS ****************************************************************/
 
 
 /* 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;
 typedef struct _ROOT_KEY
 {
     PCWSTR Name;
@@ -68,11 +74,10 @@ typedef struct _ROOT_KEY
 
 ROOT_KEY RootKeys[] =
 {
 
 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
 #if 0
     { L"HKR", NULL, NULL },
 #endif
@@ -401,10 +406,17 @@ do_reg_operation(HANDLE KeyHandle,
   return TRUE;
 }
 
   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,
 NTSTATUS
 CreateNestedKey(PHANDLE KeyHandle,
                 ACCESS_MASK DesiredAccess,
-                POBJECT_ATTRIBUTES ObjectAttributes)
+                POBJECT_ATTRIBUTES ObjectAttributes,
+                ULONG CreateOptions)
 {
     OBJECT_ATTRIBUTES LocalObjectAttributes;
     UNICODE_STRING LocalKeyName;
 {
     OBJECT_ATTRIBUTES LocalObjectAttributes;
     UNICODE_STRING LocalKeyName;
@@ -419,11 +431,16 @@ CreateNestedKey(PHANDLE KeyHandle,
                          ObjectAttributes,
                          0,
                          NULL,
                          ObjectAttributes,
                          0,
                          NULL,
-                         0,
+                         CreateOptions,
                          &Disposition);
     DPRINT("NtCreateKey(%wZ) called (Status %lx)\n", ObjectAttributes->ObjectName, Status);
     if (Status != STATUS_OBJECT_NAME_NOT_FOUND)
                          &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;
         return Status;
+    }
 
     /* Copy object attributes */
     RtlCopyMemory(&LocalObjectAttributes,
 
     /* 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)
     {
     /* 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;
         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,
 
         Status = NtCreateKey(&LocalKeyHandle,
-                             KEY_ALL_ACCESS,
+                             KEY_CREATE_SUB_KEY,
                              &LocalObjectAttributes,
                              0,
                              NULL,
                              &LocalObjectAttributes,
                              0,
                              NULL,
-                             0,
+                             REG_OPTION_NON_VOLATILE,
                              &Disposition);
         DPRINT("NtCreateKey(%wZ) called (Status %lx)\n", &LocalKeyName, Status);
                              &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))
     {
     }
 
     if (!NT_SUCCESS(Status))
     {
-        RtlFreeUnicodeString (&LocalKeyName);
+        RtlFreeUnicodeString(&LocalKeyName);
         return Status;
     }
 
         return Status;
     }
 
@@ -474,18 +493,21 @@ CreateNestedKey(PHANDLE KeyHandle,
         NtClose(LocalKeyHandle);
 
         LocalKeyName.Buffer[LocalKeyName.Length / sizeof(WCHAR)] = L'\\';
         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,
 
         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))
                              &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;
             break;
+        }
     }
 
     RtlFreeUnicodeString(&LocalKeyName);
     }
 
     RtlFreeUnicodeString(&LocalKeyName);
@@ -560,7 +582,8 @@ registry_callback(HINF hInf, PCWSTR Section, BOOLEAN Delete)
         {
             Status = CreateNestedKey(&KeyHandle,
                                      KEY_ALL_ACCESS,
         {
             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);
             if (!NT_SUCCESS(Status))
             {
                 DPRINT1("CreateNestedKey(%wZ) failed (Status %lx)\n", &Name, Status);
@@ -649,29 +672,86 @@ ImportRegistryFile(
 /*
  * Should be called under privileges
  */
 /*
  * Should be called under privileges
  */
-// static
-NTSTATUS
+static NTSTATUS
 CreateRegistryFile(
     IN PUNICODE_STRING InstallPath,
     IN PCWSTR RegistryKey,
 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;
     NTSTATUS Status;
     HANDLE FileHandle;
     UNICODE_STRING FileName;
     OBJECT_ATTRIBUTES ObjectAttributes;
     IO_STATUS_BLOCK IoStatusBlock;
+    PCWSTR Extension;
     WCHAR PathBuffer[MAX_PATH];
     WCHAR PathBuffer[MAX_PATH];
+    WCHAR PathBuffer2[MAX_PATH];
 
 
-    /* Create the file */
     CombinePaths(PathBuffer, ARRAYSIZE(PathBuffer), 3,
                  InstallPath->Buffer, L"System32\\config", RegistryKey);
     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,
     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,
 
     Status = NtCreateFile(&FileHandle,
                           FILE_GENERIC_WRITE,
@@ -773,18 +853,19 @@ CmpLinkKeyToHive(
 /*
  * Should be called under privileges
  */
 /*
  * Should be called under privileges
  */
-// static
-NTSTATUS
+static NTSTATUS
 ConnectRegistry(
     IN HKEY RootKey OPTIONAL,
 ConnectRegistry(
     IN HKEY RootKey OPTIONAL,
+    IN PCWSTR RegMountPoint,
     // IN HANDLE RootDirectory OPTIONAL,
     IN PUNICODE_STRING InstallPath,
     // 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;
     UNICODE_STRING KeyName, FileName;
     OBJECT_ATTRIBUTES KeyObjectAttributes;
     OBJECT_ATTRIBUTES FileObjectAttributes;
@@ -795,7 +876,7 @@ ConnectRegistry(
                                &KeyName,
                                OBJ_CASE_INSENSITIVE,
                                RootKey,
                                &KeyName,
                                OBJ_CASE_INSENSITIVE,
                                RootKey,
-                               NULL);
+                               NULL);   // Descriptor
 
     CombinePaths(PathBuffer, ARRAYSIZE(PathBuffer), 3,
                  InstallPath->Buffer, L"System32\\config", RegistryKey);
 
     CombinePaths(PathBuffer, ARRAYSIZE(PathBuffer), 3,
                  InstallPath->Buffer, L"System32\\config", RegistryKey);
@@ -806,56 +887,202 @@ ConnectRegistry(
                                NULL, // RootDirectory,
                                NULL);
 
                                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))
     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;
 }
 
 
     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
 NTSTATUS
-RegInitializeRegistry(
-    IN PUNICODE_STRING InstallPath)
+VerifyRegistryHives(
+    IN PUNICODE_STRING InstallPath,
+    OUT PBOOLEAN ShouldUpdateRegistry)
 {
     NTSTATUS Status;
 {
     NTSTATUS Status;
-    HANDLE KeyHandle;
-    UNICODE_STRING KeyName;
-    OBJECT_ATTRIBUTES ObjectAttributes;
     BOOLEAN PrivilegeSet[2] = {FALSE, FALSE};
     BOOLEAN PrivilegeSet[2] = {FALSE, FALSE};
-    ULONG Disposition;
     UINT i;
     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))
     {
     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;
     }
         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]);
 
     /* 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%20files.asciidoc
      * for more information.
      */
      * See https://github.com/libyal/winreg-kb/blob/master/documentation/Registry%20files.asciidoc
      * 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,
     InitializeObjectAttributes(&ObjectAttributes,
                                &KeyName,
                                OBJ_CASE_INSENSITIVE,
@@ -906,14 +1132,18 @@ RegInitializeRegistry(
     }
     NtFlushKey(KeyHandle);
 
     }
     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,
         Status = CreateRegistryFile(InstallPath,
-                                    RegistryKeys[i],
+                                    RegistryHives[i].HiveName,
+                                    RegistryHives[i].State != Repair, // RegistryHives[i].State == Create,
                                     KeyHandle);
         if (!NT_SUCCESS(Status))
         {
                                     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);
             /* 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.
      */
      * 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,
                          &ObjectAttributes,
                          0,
                          NULL,
-                         REG_OPTION_VOLATILE,
+                         REG_OPTION_NON_VOLATILE, // REG_OPTION_VOLATILE, // FIXME!
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
@@ -968,7 +1198,7 @@ RegInitializeRegistry(
                          &ObjectAttributes,
                          0,
                          NULL,
                          &ObjectAttributes,
                          0,
                          NULL,
-                         REG_OPTION_VOLATILE,
+                         REG_OPTION_NON_VOLATILE, // REG_OPTION_VOLATILE, // FIXME!
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
@@ -978,72 +1208,61 @@ RegInitializeRegistry(
     RootKeys[GetPredefKeyIndex(HKEY_USERS)].Handle = KeyHandle;
 
 
     RootKeys[GetPredefKeyIndex(HKEY_USERS)].Handle = KeyHandle;
 
 
-
     /*
      * Now properly mount the offline hive files
      */
     /*
      * 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
 
     /* HKCU is a handle to 'HKU\.DEFAULT' */
 #if 0
@@ -1089,12 +1308,13 @@ RegInitializeRegistry(
                                NULL);
 #endif
     KeyHandle = NULL;
                                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,
     Status = NtCreateKey(&KeyHandle,
                          KEY_ALL_ACCESS,
                          &ObjectAttributes,
                          0,
                          NULL,
-                         0,
+                         REG_OPTION_NON_VOLATILE,
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
                          &Disposition);
     if (!NT_SUCCESS(Status))
     {
@@ -1109,31 +1329,11 @@ RegInitializeRegistry(
     RootKeys[GetPredefKeyIndex(HKEY_CLASSES_ROOT)].Handle = KeyHandle;
 
 
     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 */
     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,
     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",
     /* 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");
     }
 
 
     {
         DPRINT1("CmpLinkKeyToHive(CurrentControlSet) failed!\n");
     }
 
 
+    Status = STATUS_SUCCESS;
+
+
 Quit:
     /* Remove restore and backup privileges */
     RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, PrivilegeSet[1], FALSE, &PrivilegeSet[1]);
 Quit:
     /* Remove restore and backup privileges */
     RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, PrivilegeSet[1], FALSE, &PrivilegeSet[1]);
@@ -1178,18 +1381,22 @@ Quit:
 }
 
 VOID
 }
 
 VOID
-RegCleanupRegistry(VOID)
+RegCleanupRegistry(
+    IN PUNICODE_STRING InstallPath)
 {
     NTSTATUS Status;
     UNICODE_STRING KeyName;
     OBJECT_ATTRIBUTES KeyObjectAttributes;
     BOOLEAN PrivilegeSet[2] = {FALSE, FALSE};
 {
     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)
         {
 
     for (i = 0; i < ARRAYSIZE(RootKeys); ++i)
     {
         if (RootKeys[i].Handle)
         {
+            NtFlushKey(RootKeys[i].Handle);
             NtClose(RootKeys[i].Handle);
             RootKeys[i].Handle = NULL;
         }
             NtClose(RootKeys[i].Handle);
             RootKeys[i].Handle = NULL;
         }
@@ -1220,25 +1427,38 @@ RegCleanupRegistry(VOID)
                                NULL,
                                NULL);
 
                                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]);
 
     /* Remove restore and backup privileges */
     RtlAdjustPrivilege(SE_BACKUP_PRIVILEGE, PrivilegeSet[1], FALSE, &PrivilegeSet[1]);
index c97d897..cae50e8 100644 (file)
@@ -43,12 +43,18 @@ ImportRegistryFile(
     LCID LocaleId,
     BOOLEAN Delete);
 
     LCID LocaleId,
     BOOLEAN Delete);
 
+NTSTATUS
+VerifyRegistryHives(
+    IN PUNICODE_STRING InstallPath,
+    OUT PBOOLEAN ShouldUpdateRegistry);
+
 NTSTATUS
 RegInitializeRegistry(
     IN PUNICODE_STRING InstallPath);
 
 VOID
 NTSTATUS
 RegInitializeRegistry(
     IN PUNICODE_STRING InstallPath);
 
 VOID
-RegCleanupRegistry(VOID);
+RegCleanupRegistry(
+    IN PUNICODE_STRING InstallPath);
 
 VOID
 SetDefaultPagefile(
 
 VOID
 SetDefaultPagefile(
index 4eb2558..9318f7e 100644 (file)
@@ -551,6 +551,7 @@ CheckUnattendedSetup(VOID)
     wcscpy(UnattendInstallationDirectory, Value);
 
     IsUnattendedSetup = TRUE;
     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))
 
     /* Search for 'MBRInstallType' in the 'Unattend' section */
     if (SetupFindFirstLineW(UnattendInf, L"Unattend", L"MBRInstallType", &Context))
@@ -589,8 +590,6 @@ CheckUnattendedSetup(VOID)
     }
 
     SetupCloseInfFile(UnattendInf);
     }
 
     SetupCloseInfFile(UnattendInf);
-
-    DPRINT("Running unattended setup\n");
 }
 
 
 }
 
 
@@ -1503,7 +1502,7 @@ IsDiskSizeValid(PPARTENTRY PartEntry)
 
     if (size < RequiredPartitionDiskSpace)
     {
 
     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;
     }
         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;
 
     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))
     if (IsUnattendedSetup)
     {
         if (!IsValidPath(UnattendInstallationDirectory))
@@ -4057,20 +4072,34 @@ FileCopyPage(PINPUT_RECORD Ir)
 static PAGE_NUMBER
 RegistryPage(PINPUT_RECORD Ir)
 {
 static PAGE_NUMBER
 RegistryPage(PINPUT_RECORD Ir)
 {
+    NTSTATUS Status;
     INFCONTEXT InfContext;
     PWSTR Action;
     PWSTR File;
     PWSTR Section;
     BOOLEAN Delete;
     INFCONTEXT InfContext;
     PWSTR Action;
     PWSTR File;
     PWSTR Section;
     BOOLEAN Delete;
-    NTSTATUS Status;
 
     MUIDisplayPage(REGISTRY_PAGE);
 
     if (RepairUpdateFlag)
     {
 
     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 */
     }
 
     /* Initialize the registry and setup the default installation hives */
@@ -4087,7 +4116,7 @@ RegistryPage(PINPUT_RECORD Ir)
         else
         /*************************************/
         {
         else
         /*************************************/
         {
-            /* Something else (correct) failed */
+            /* Something else failed */
             MUIDisplayError(ERROR_CREATE_HIVE, Ir, POPUP_WAIT_ENTER);
         }
         return QUIT_PAGE;
             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");
     if (!SetupFindFirstLineW(SetupInf, L"HiveInfs.Install", NULL, &InfContext))
     {
         DPRINT1("SetupFindFirstLine() failed\n");
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_FIND_REGISTRY, Ir, POPUP_WAIT_ENTER);
         MUIDisplayError(ERROR_FIND_REGISTRY, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     do
     }
 
     do
@@ -4116,27 +4144,19 @@ RegistryPage(PINPUT_RECORD Ir)
             break; // Hackfix
 
         if (!_wcsicmp(Action, L"AddReg"))
             break; // Hackfix
 
         if (!_wcsicmp(Action, L"AddReg"))
-        {
             Delete = FALSE;
             Delete = FALSE;
-        }
         else if (!_wcsicmp(Action, L"DelReg"))
         else if (!_wcsicmp(Action, L"DelReg"))
-        {
             Delete = TRUE;
             Delete = TRUE;
-        }
         else
         else
-        {
             continue;
             continue;
-        }
 
         CONSOLE_SetStatusText(MUIGetString(STRING_IMPORTFILE), File);
 
         if (!ImportRegistryFile(File, Section, LanguageId, Delete))
         {
             DPRINT1("Importing %S failed\n", File);
 
         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);
             MUIDisplayError(ERROR_IMPORT_HIVE, Ir, POPUP_WAIT_ENTER);
-            return QUIT_PAGE;
+            goto Cleanup;
         }
     } while (SetupFindNextLine(&InfContext, &InfContext));
 
         }
     } while (SetupFindNextLine(&InfContext, &InfContext));
 
@@ -4144,35 +4164,31 @@ RegistryPage(PINPUT_RECORD Ir)
     CONSOLE_SetStatusText(MUIGetString(STRING_DISPLAYETTINGSUPDATE));
     if (!ProcessDisplayRegistry(SetupInf, DisplayList))
     {
     CONSOLE_SetStatusText(MUIGetString(STRING_DISPLAYETTINGSUPDATE));
     if (!ProcessDisplayRegistry(SetupInf, DisplayList))
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_UPDATE_DISPLAY_SETTINGS, Ir, POPUP_WAIT_ENTER);
         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))
     {
     }
 
     /* Set the locale */
     CONSOLE_SetStatusText(MUIGetString(STRING_LOCALESETTINGSUPDATE));
     if (!ProcessLocaleRegistry(LanguageList))
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_UPDATE_LOCALESETTINGS, Ir, POPUP_WAIT_ENTER);
         MUIDisplayError(ERROR_UPDATE_LOCALESETTINGS, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     /* Add keyboard layouts */
     CONSOLE_SetStatusText(MUIGetString(STRING_ADDKBLAYOUTS));
     if (!AddKeyboardLayouts())
     {
     }
 
     /* Add keyboard layouts */
     CONSOLE_SetStatusText(MUIGetString(STRING_ADDKBLAYOUTS));
     if (!AddKeyboardLayouts())
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_ADDING_KBLAYOUTS, Ir, POPUP_WAIT_ENTER);
         MUIDisplayError(ERROR_ADDING_KBLAYOUTS, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     /* Set GeoID */
     if (!SetGeoID(MUIGetGeoID()))
     {
     }
 
     /* Set GeoID */
     if (!SetGeoID(MUIGetGeoID()))
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_UPDATE_GEOID, Ir, POPUP_WAIT_ENTER);
         MUIDisplayError(ERROR_UPDATE_GEOID, Ir, POPUP_WAIT_ENTER);
-        return QUIT_PAGE;
+        goto Cleanup;
     }
 
     if (!IsUnattendedSetup)
     }
 
     if (!IsUnattendedSetup)
@@ -4181,9 +4197,8 @@ RegistryPage(PINPUT_RECORD Ir)
         CONSOLE_SetStatusText(MUIGetString(STRING_KEYBOARDSETTINGSUPDATE));
         if (!ProcessKeyboardLayoutRegistry(LayoutList))
         {
         CONSOLE_SetStatusText(MUIGetString(STRING_KEYBOARDSETTINGSUPDATE));
         if (!ProcessKeyboardLayoutRegistry(LayoutList))
         {
-            RegCleanupRegistry();
             MUIDisplayError(ERROR_UPDATE_KBSETTINGS, Ir, POPUP_WAIT_ENTER);
             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())
     {
     CONSOLE_SetStatusText(MUIGetString(STRING_CODEPAGEINFOUPDATE));
     if (!AddCodePage())
     {
-        RegCleanupRegistry();
         MUIDisplayError(ERROR_ADDING_CODEPAGE, Ir, POPUP_WAIT_ENTER);
         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 */
     }
 
     /* 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);
 
     SetMountedDeviceValues(PartitionList);
 
+Cleanup:
     //
     // TODO: Unload all the registry stuff, perform 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;
+    }
 }
 
 
 }