[MKHIVE] Fix some compilation warnings.
[reactos.git] / sdk / tools / mkhive / registry.c
index e7a8cc9..8ee8d36 100644 (file)
  * PROJECT:         ReactOS hive maker
  * FILE:            tools/mkhive/registry.c
  * PURPOSE:         Registry code
- * PROGRAMMER:      Hervé Poussineau
+ * PROGRAMMERS:     Hervé Poussineau
+ *                  Hermès Bélusca-Maïto
  */
 
-/*
- * TODO:
- *   - Implement RegDeleteKeyW() and RegDeleteValueW()
- */
-
-#include <stdlib.h>
-#include <string.h>
-#include <stdio.h>
+/* INCLUDES *****************************************************************/
 
 #define NDEBUG
 #include "mkhive.h"
 
+/* DATA *********************************************************************/
+
+typedef struct _REPARSE_POINT
+{
+    LIST_ENTRY ListEntry;
+    PCMHIVE SourceHive;
+    HCELL_INDEX SourceKeyCellOffset;
+    PCMHIVE DestinationHive;
+    HCELL_INDEX DestinationKeyCellOffset;
+} REPARSE_POINT, *PREPARSE_POINT;
+
+typedef struct _MEMKEY
+{
+    /* Information on hard disk structure */
+    HCELL_INDEX KeyCellOffset;
+    PCMHIVE RegistryHive;
+} MEMKEY, *PMEMKEY;
+
+#define HKEY_TO_MEMKEY(hKey) ((PMEMKEY)(hKey))
+#define MEMKEY_TO_HKEY(memKey) ((HKEY)(memKey))
+
 static CMHIVE RootHive;
 static PMEMKEY RootKey;
-CMHIVE DefaultHive;  /* \Registry\User\.DEFAULT */
-CMHIVE SamHive;      /* \Registry\Machine\SAM */
-CMHIVE SecurityHive; /* \Registry\Machine\SECURITY */
-CMHIVE SoftwareHive; /* \Registry\Machine\SOFTWARE */
-CMHIVE SystemHive;   /* \Registry\Machine\SYSTEM */
-CMHIVE BcdHive;      /* \Registry\Machine\BCD00000000 */
+
+static CMHIVE SystemHive;   /* \Registry\Machine\SYSTEM */
+static CMHIVE SoftwareHive; /* \Registry\Machine\SOFTWARE */
+static CMHIVE DefaultHive;  /* \Registry\User\.DEFAULT */
+static CMHIVE SamHive;      /* \Registry\Machine\SAM */
+static CMHIVE SecurityHive; /* \Registry\Machine\SECURITY */
+static CMHIVE BcdHive;      /* \Registry\Machine\BCD00000000 */
 
 //
 // TODO: Write these values in a more human-readable form.
@@ -56,7 +72,7 @@ CMHIVE BcdHive;      /* \Registry\Machine\BCD00000000 */
 // A cross-check was subsequently done with the system hives to verify that
 // the security descriptors were the same.
 //
-UCHAR BcdSecurity[] =
+static UCHAR BcdSecurity[] =
 {
     // SECURITY_DESCRIPTOR_RELATIVE
     0x01,                   // Revision
@@ -118,7 +134,7 @@ UCHAR BcdSecurity[] =
     0x01, 0x02, 0x00, 0x00
 };
 
-UCHAR SoftwareSecurity[] =
+static UCHAR SoftwareSecurity[] =
 {
     // SECURITY_DESCRIPTOR_RELATIVE
     0x01,                   // Revision
@@ -234,7 +250,7 @@ UCHAR SoftwareSecurity[] =
 };
 
 // Same security for SYSTEM, SAM and .DEFAULT
-UCHAR SystemSecurity[] =
+static UCHAR SystemSecurity[] =
 {
     // SECURITY_DESCRIPTOR_RELATIVE
     0x01,                   // Revision
@@ -330,6 +346,26 @@ UCHAR SystemSecurity[] =
     0x01, 0x02, 0x00, 0x00
 };
 
+/* GLOBALS ******************************************************************/
+
+HIVE_LIST_ENTRY RegistryHives[/*MAX_NUMBER_OF_REGISTRY_HIVES*/] =
+{
+    /* Special Setup system registry hive */
+    // WARNING: Please *keep* it in first position!
+    { "SETUPREG", L"Registry\\Machine\\SYSTEM"     , &SystemHive  , SystemSecurity  , sizeof(SystemSecurity)   },
+
+    /* Regular registry hives */
+    { "SYSTEM"  , L"Registry\\Machine\\SYSTEM"     , &SystemHive  , SystemSecurity  , sizeof(SystemSecurity)   },
+    { "SOFTWARE", L"Registry\\Machine\\SOFTWARE"   , &SoftwareHive, SoftwareSecurity, sizeof(SoftwareSecurity) },
+    { "DEFAULT" , L"Registry\\User\\.DEFAULT"      , &DefaultHive , SystemSecurity  , sizeof(SystemSecurity)   },
+    { "SAM"     , L"Registry\\Machine\\SAM"        , &SamHive     , SystemSecurity  , sizeof(SystemSecurity)   },
+    { "SECURITY", L"Registry\\Machine\\SECURITY"   , &SecurityHive, NULL            , 0                        },
+    { "BCD"     , L"Registry\\Machine\\BCD00000000", &BcdHive     , BcdSecurity     , sizeof(BcdSecurity)      },
+};
+C_ASSERT(_countof(RegistryHives) == MAX_NUMBER_OF_REGISTRY_HIVES);
+
+/* FUNCTIONS ****************************************************************/
+
 static PMEMKEY
 CreateInMemoryStructure(
     IN PCMHIVE RegistryHive,
@@ -346,30 +382,30 @@ CreateInMemoryStructure(
     return Key;
 }
 
+LIST_ENTRY CmiHiveListHead;
 LIST_ENTRY CmiReparsePointsHead;
 
 static LONG
-RegpOpenOrCreateKey(
+RegpCreateOrOpenKey(
     IN HKEY hParentKey,
     IN PCWSTR KeyName,
     IN BOOL AllowCreation,
     IN BOOL Volatile,
     OUT PHKEY Key)
 {
+    NTSTATUS Status;
     PWSTR LocalKeyName;
     PWSTR End;
     UNICODE_STRING KeyString;
-    NTSTATUS Status;
     PREPARSE_POINT CurrentReparsePoint;
     PMEMKEY CurrentKey;
     PCMHIVE ParentRegistryHive;
     HCELL_INDEX ParentCellOffset;
     PCM_KEY_NODE ParentKeyCell;
     PLIST_ENTRY Ptr;
-    PCM_KEY_NODE SubKeyCell;
     HCELL_INDEX BlockOffset;
 
-    DPRINT("RegpCreateOpenKey('%S')\n", KeyName);
+    DPRINT("RegpCreateOrOpenKey('%S')\n", KeyName);
 
     if (*KeyName == OBJ_NAME_PATH_SEPARATOR)
     {
@@ -410,7 +446,7 @@ RegpOpenOrCreateKey(
 
         ParentKeyCell = (PCM_KEY_NODE)HvGetCell(&ParentRegistryHive->Hive, ParentCellOffset);
         if (!ParentKeyCell)
-            return STATUS_UNSUCCESSFUL;
+            return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
 
         VERIFY_KEY_CELL(ParentKeyCell);
 
@@ -442,11 +478,19 @@ RegpOpenOrCreateKey(
                                   Volatile,
                                   &BlockOffset);
         }
+        else // if (BlockOffset == HCELL_NIL)
+        {
+            Status = STATUS_OBJECT_NAME_NOT_FOUND;
+        }
 
         HvReleaseCell(&ParentRegistryHive->Hive, ParentCellOffset);
 
         if (!NT_SUCCESS(Status))
-            return ERROR_UNSUCCESSFUL;
+        {
+            DPRINT("RegpCreateOrOpenKey('%S'): Could not create or open subkey '%.*S', Status 0x%08x\n",
+                   KeyName, (int)(KeyString.Length / sizeof(WCHAR)), KeyString.Buffer, Status);
+            return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
+        }
 
         ParentCellOffset = BlockOffset;
         if (End)
@@ -457,7 +501,7 @@ RegpOpenOrCreateKey(
 
     CurrentKey = CreateInMemoryStructure(ParentRegistryHive, ParentCellOffset);
     if (!CurrentKey)
-        return ERROR_OUTOFMEMORY;
+        return ERROR_NOT_ENOUGH_MEMORY; // STATUS_NO_MEMORY;
 
     *Key = MEMKEY_TO_HKEY(CurrentKey);
 
@@ -465,31 +509,24 @@ RegpOpenOrCreateKey(
 }
 
 LONG WINAPI
-RegCreateKeyW(
-    IN HKEY hKey,
-    IN LPCWSTR lpSubKey,
-    OUT PHKEY phkResult)
+RegCloseKey(
+    IN HKEY hKey)
 {
-    return RegpOpenOrCreateKey(hKey, lpSubKey, TRUE, FALSE, phkResult);
-}
+    PMEMKEY Key = HKEY_TO_MEMKEY(hKey); // ParentKey
+
+    /* Free the object */
+    free(Key);
 
-LONG WINAPI
-RegDeleteKeyW(
-    IN HKEY hKey,
-    IN LPCWSTR lpSubKey)
-{
-    DPRINT1("RegDeleteKeyW(0x%p, '%S') is UNIMPLEMENTED!\n",
-            hKey, (lpSubKey ? lpSubKey : L""));
     return ERROR_SUCCESS;
 }
 
 LONG WINAPI
-RegOpenKeyW(
+RegCreateKeyW(
     IN HKEY hKey,
     IN LPCWSTR lpSubKey,
     OUT PHKEY phkResult)
 {
-    return RegpOpenOrCreateKey(hKey, lpSubKey, FALSE, FALSE, phkResult);
+    return RegpCreateOrOpenKey(hKey, lpSubKey, TRUE, FALSE, phkResult);
 }
 
 LONG WINAPI
@@ -504,13 +541,117 @@ RegCreateKeyExW(
     OUT PHKEY phkResult,
     OUT LPDWORD lpdwDisposition OPTIONAL)
 {
-    return RegpOpenOrCreateKey(hKey,
+    return RegpCreateOrOpenKey(hKey,
                                lpSubKey,
                                TRUE,
                                (dwOptions & REG_OPTION_VOLATILE) != 0,
                                phkResult);
 }
 
+LONG WINAPI
+RegDeleteKeyW(
+    IN HKEY hKey,
+    IN LPCWSTR lpSubKey)
+{
+    LONG rc;
+    NTSTATUS Status;
+    HKEY hTargetKey;
+    PMEMKEY Key; // ParentKey
+    PHHIVE Hive;
+    PCM_KEY_NODE KeyNode; // ParentNode
+    PCM_KEY_NODE Parent;
+    HCELL_INDEX ParentCell;
+
+    if (lpSubKey)
+    {
+        rc = RegOpenKeyW(hKey, lpSubKey, &hTargetKey);
+        if (rc != ERROR_SUCCESS)
+            return rc;
+    }
+    else
+    {
+        hTargetKey = hKey;
+        rc = ERROR_SUCCESS;
+    }
+
+    /* Don't allow deleting the root */
+    if (hTargetKey == RootKey)
+    {
+        /* Fail */
+        rc = ERROR_ACCESS_DENIED; // STATUS_CANNOT_DELETE;
+        goto Quit;
+    }
+
+    /* Get the hive and node */
+    Key = HKEY_TO_MEMKEY(hTargetKey);
+    Hive = &Key->RegistryHive->Hive;
+
+    /* Get the key node */
+    KeyNode = (PCM_KEY_NODE)HvGetCell(Hive, Key->KeyCellOffset);
+    if (!KeyNode)
+    {
+        rc = ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
+        goto Quit;
+    }
+
+    ASSERT(KeyNode->Signature == CM_KEY_NODE_SIGNATURE);
+
+    /* Check if we don't have any children */
+    if (!(KeyNode->SubKeyCounts[Stable] + KeyNode->SubKeyCounts[Volatile]) &&
+        !(KeyNode->Flags & KEY_NO_DELETE))
+    {
+        /* Get the parent and free the cell */
+        ParentCell = KeyNode->Parent;
+        Status = CmpFreeKeyByCell(Hive, Key->KeyCellOffset, TRUE);
+        if (NT_SUCCESS(Status))
+        {
+            /* Get the parent node */
+            Parent = (PCM_KEY_NODE)HvGetCell(Hive, ParentCell);
+            if (Parent)
+            {
+                /* Make sure we're dirty */
+                ASSERT(HvIsCellDirty(Hive, ParentCell));
+
+                /* Update the write time */
+                KeQuerySystemTime(&Parent->LastWriteTime);
+
+                /* Release the cell */
+                HvReleaseCell(Hive, ParentCell);
+            }
+
+            rc = ERROR_SUCCESS;
+        }
+        else
+        {
+            /* Fail */
+            rc = ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
+        }
+    }
+    else
+    {
+        /* Fail */
+        rc = ERROR_ACCESS_DENIED; // STATUS_CANNOT_DELETE;
+    }
+
+    /* Release the cell */
+    HvReleaseCell(Hive, Key->KeyCellOffset);
+
+Quit:
+    if (lpSubKey)
+        RegCloseKey(hTargetKey);
+
+    return rc;
+}
+
+LONG WINAPI
+RegOpenKeyW(
+    IN HKEY hKey,
+    IN LPCWSTR lpSubKey,
+    OUT PHKEY phkResult)
+{
+    return RegpCreateOrOpenKey(hKey, lpSubKey, FALSE, FALSE, phkResult);
+}
+
 LONG WINAPI
 RegSetValueExW(
     IN HKEY hKey,
@@ -524,6 +665,7 @@ RegSetValueExW(
     PHHIVE Hive;
     PCM_KEY_NODE KeyNode; // ParentNode
     PCM_KEY_VALUE ValueCell;
+    ULONG ChildIndex;
     HCELL_INDEX CellIndex;
     UNICODE_STRING ValueNameString;
 
@@ -537,7 +679,7 @@ RegSetValueExW(
 
         /* Special handling of registry links */
         if (cbData != sizeof(PVOID))
-            return STATUS_INVALID_PARAMETER;
+            return ERROR_INVALID_PARAMETER; // STATUS_INVALID_PARAMETER;
 
         DestKey = HKEY_TO_MEMKEY(*(PHKEY)lpData);
 
@@ -545,20 +687,20 @@ RegSetValueExW(
 
         /* Create the link in registry hive (if applicable) */
         if (Key->RegistryHive != DestKey->RegistryHive)
-            return STATUS_SUCCESS;
+            return ERROR_SUCCESS;
 
         DPRINT1("Save link to registry\n");
-        return STATUS_NOT_IMPLEMENTED;
+        return ERROR_INVALID_FUNCTION; // STATUS_NOT_IMPLEMENTED;
     }
 
     if ((cbData & ~CM_KEY_VALUE_SPECIAL_SIZE) != cbData)
-        return STATUS_UNSUCCESSFUL;
+        return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
 
     Hive = &Key->RegistryHive->Hive;
 
     KeyNode = (PCM_KEY_NODE)HvGetCell(Hive, Key->KeyCellOffset);
     if (!KeyNode)
-        return ERROR_UNSUCCESSFUL;
+        return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
 
     ASSERT(KeyNode->Signature == CM_KEY_NODE_SIGNATURE);
 
@@ -567,12 +709,23 @@ RegSetValueExW(
 
     /* Initialize value name string */
     RtlInitUnicodeString(&ValueNameString, lpValueName);
-    CellIndex = CmpFindValueByName(Hive, KeyNode, &ValueNameString);
+    if (!CmpFindNameInList(Hive,
+                           &KeyNode->ValueList,
+                           &ValueNameString,
+                           &ChildIndex,
+                           &CellIndex))
+    {
+        /* Sanity check */
+        ASSERT(CellIndex == HCELL_NIL);
+        /* Fail */
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+    }
     if (CellIndex == HCELL_NIL)
     {
         /* The value doesn't exist, create a new one */
         Status = CmiAddValueKey(Key->RegistryHive,
                                 KeyNode,
+                                ChildIndex,
                                 &ValueNameString,
                                 &ValueCell,
                                 &CellIndex);
@@ -588,7 +741,7 @@ RegSetValueExW(
     // /**/HvReleaseCell(Hive, CellIndex);/**/
 
     if (!NT_SUCCESS(Status))
-        return ERROR_UNSUCCESSFUL;
+        return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
 
     /* Get size of the allocated cell (if any) */
     if (!(ValueCell->DataLength & CM_KEY_VALUE_SPECIAL_SIZE) &&
@@ -596,7 +749,7 @@ RegSetValueExW(
     {
         DataCell = HvGetCell(Hive, ValueCell->Data);
         if (!DataCell)
-            return ERROR_UNSUCCESSFUL;
+            return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
 
         DataCellSize = (ULONG)(-HvGetCellSize(Hive, DataCell));
     }
@@ -631,7 +784,7 @@ RegSetValueExW(
             if (NewOffset == HCELL_NIL)
             {
                 DPRINT("HvAllocateCell() has failed!\n");
-                return ERROR_UNSUCCESSFUL;
+                return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
             }
 
             if (DataCell)
@@ -665,15 +818,15 @@ RegSetValueExW(
 }
 
 
-// Synced with freeldr/windows/registry.c
+// Synced with freeldr/ntldr/registry.c
 static
 VOID
 RepGetValueData(
     IN PHHIVE Hive,
     IN PCM_KEY_VALUE ValueCell,
-    OUT ULONG* Type OPTIONAL,
+    OUT PULONG Type OPTIONAL,
     OUT PUCHAR Data OPTIONAL,
-    IN OUT ULONG* DataSize OPTIONAL)
+    IN OUT PULONG DataSize OPTIONAL)
 {
     ULONG DataLength;
     PVOID DataCell;
@@ -704,7 +857,7 @@ RepGetValueData(
     }
 }
 
-// Similar to RegQueryValue in freeldr/windows/registry.c
+// Similar to RegQueryValue in freeldr/ntldr/registry.c
 LONG WINAPI
 RegQueryValueExW(
     IN HKEY hKey,
@@ -723,7 +876,7 @@ RegQueryValueExW(
 
     KeyNode = (PCM_KEY_NODE)HvGetCell(Hive, ParentKey->KeyCellOffset);
     if (!KeyNode)
-        return ERROR_UNSUCCESSFUL;
+        return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
 
     ASSERT(KeyNode->Signature == CM_KEY_NODE_SIGNATURE);
 
@@ -731,7 +884,7 @@ RegQueryValueExW(
     RtlInitUnicodeString(&ValueNameString, lpValueName);
     CellIndex = CmpFindValueByName(Hive, KeyNode, &ValueNameString);
     if (CellIndex == HCELL_NIL)
-        return ERROR_FILE_NOT_FOUND;
+        return ERROR_FILE_NOT_FOUND; // STATUS_OBJECT_NAME_NOT_FOUND;
 
     /* Get the value cell */
     ValueCell = HvGetCell(Hive, CellIndex);
@@ -749,9 +902,108 @@ RegDeleteValueW(
     IN HKEY hKey,
     IN LPCWSTR lpValueName OPTIONAL)
 {
-    DPRINT1("RegDeleteValueW(0x%p, '%S') is UNIMPLEMENTED!\n",
-            hKey, (lpValueName ? lpValueName : L""));
-    return ERROR_UNSUCCESSFUL;
+    LONG rc;
+    NTSTATUS Status;
+    PMEMKEY Key = HKEY_TO_MEMKEY(hKey); // ParentKey
+    PHHIVE Hive = &Key->RegistryHive->Hive;
+    PCM_KEY_NODE KeyNode; // ParentNode
+    PCM_KEY_VALUE ValueCell;
+    HCELL_INDEX CellIndex;
+    ULONG ChildIndex;
+    UNICODE_STRING ValueNameString;
+
+    KeyNode = (PCM_KEY_NODE)HvGetCell(Hive, Key->KeyCellOffset);
+    if (!KeyNode)
+        return ERROR_GEN_FAILURE; // STATUS_UNSUCCESSFUL;
+
+    ASSERT(KeyNode->Signature == CM_KEY_NODE_SIGNATURE);
+
+    /* Initialize value name string */
+    RtlInitUnicodeString(&ValueNameString, lpValueName);
+    if (!CmpFindNameInList(Hive,
+                           &KeyNode->ValueList,
+                           &ValueNameString,
+                           &ChildIndex,
+                           &CellIndex))
+    {
+        /* Sanity check */
+        ASSERT(CellIndex == HCELL_NIL);
+    }
+    if (CellIndex == HCELL_NIL)
+    {
+        rc = ERROR_FILE_NOT_FOUND; // STATUS_OBJECT_NAME_NOT_FOUND;
+        goto Quit;
+    }
+
+    /* We found the value, mark all relevant cells dirty */
+    HvMarkCellDirty(Hive, Key->KeyCellOffset, FALSE);
+    HvMarkCellDirty(Hive, KeyNode->ValueList.List, FALSE);
+    HvMarkCellDirty(Hive, CellIndex, FALSE);
+
+    /* Get the key value */
+    ValueCell = (PCM_KEY_VALUE)HvGetCell(Hive, CellIndex);
+    ASSERT(ValueCell);
+
+    /* Mark it and all related data as dirty */
+    if (!CmpMarkValueDataDirty(Hive, ValueCell))
+    {
+        /* Not enough log space, fail */
+        rc = ERROR_NO_LOG_SPACE; // STATUS_NO_LOG_SPACE;
+        goto Quit;
+    }
+
+    /* Sanity checks */
+    ASSERT(HvIsCellDirty(Hive, KeyNode->ValueList.List));
+    ASSERT(HvIsCellDirty(Hive, CellIndex));
+
+    /* Remove the value from the child list */
+    Status = CmpRemoveValueFromList(Hive, ChildIndex, &KeyNode->ValueList);
+    if (!NT_SUCCESS(Status))
+    {
+        /* Set known error */
+        rc = ERROR_NO_SYSTEM_RESOURCES; // STATUS_INSUFFICIENT_RESOURCES;
+        goto Quit;
+    }
+
+    /* Remove the value and its data itself */
+    if (!CmpFreeValue(Hive, CellIndex))
+    {
+        /* Failed to free the value, fail */
+        rc = ERROR_NO_SYSTEM_RESOURCES; // STATUS_INSUFFICIENT_RESOURCES;
+        goto Quit;
+    }
+
+    /* Set the last write time */
+    KeQuerySystemTime(&KeyNode->LastWriteTime);
+
+    /* Sanity check */
+    ASSERT(HvIsCellDirty(Hive, Key->KeyCellOffset));
+
+    /* Check if the value list is empty now */
+    if (!KeyNode->ValueList.Count)
+    {
+        /* Then clear key node data */
+        KeyNode->MaxValueNameLen = 0;
+        KeyNode->MaxValueDataLen = 0;
+    }
+
+    /* Change default Status to success */
+    rc = ERROR_SUCCESS;
+
+Quit:
+    /* Check if we had a value */
+    if (ValueCell)
+    {
+        /* Release the child cell */
+        ASSERT(CellIndex != HCELL_NIL);
+        HvReleaseCell(Hive, CellIndex);
+    }
+
+    /* Release the parent cell, if any */
+    if (KeyNode)
+        HvReleaseCell(Hive, Key->KeyCellOffset);
+
+    return rc;
 }
 
 
@@ -759,16 +1011,16 @@ static BOOL
 ConnectRegistry(
     IN HKEY RootKey,
     IN PCMHIVE HiveToConnect,
-    IN PUCHAR Descriptor,
-    IN ULONG DescriptorLength,
-    IN LPCWSTR Path)
+    IN PUCHAR SecurityDescriptor,
+    IN ULONG SecurityDescriptorLength,
+    IN PCWSTR Path)
 {
     NTSTATUS Status;
+    LONG rc;
     PREPARSE_POINT ReparsePoint;
     PMEMKEY NewKey;
-    LONG rc;
 
-    ReparsePoint = (PREPARSE_POINT)malloc(sizeof(REPARSE_POINT));
+    ReparsePoint = (PREPARSE_POINT)malloc(sizeof(*ReparsePoint));
     if (!ReparsePoint)
         return FALSE;
 
@@ -794,11 +1046,11 @@ ConnectRegistry(
      */
     Status = CmiCreateSecurityKey(&HiveToConnect->Hive,
                                   HiveToConnect->Hive.BaseBlock->RootCell,
-                                  Descriptor, DescriptorLength);
+                                  SecurityDescriptor, SecurityDescriptorLength);
     if (!NT_SUCCESS(Status))
         DPRINT1("Failed to add security for root key '%S'\n", Path);
 
-    /* Create key */
+    /* Create the key */
     rc = RegCreateKeyExW(RootKey,
                          Path,
                          0,
@@ -821,18 +1073,70 @@ ConnectRegistry(
     ReparsePoint->DestinationHive = NewKey->RegistryHive;
     ReparsePoint->DestinationKeyCellOffset = NewKey->KeyCellOffset;
     InsertTailList(&CmiReparsePointsHead, &ReparsePoint->ListEntry);
+
     return TRUE;
 }
 
-LIST_ENTRY CmiHiveListHead;
+static BOOL
+CreateSymLink(
+    IN PCWSTR LinkKeyPath OPTIONAL,
+    IN OUT PHKEY LinkKeyHandle OPTIONAL,
+    // IN PCWSTR TargetKeyPath OPTIONAL,
+    IN HKEY TargetKeyHandle)
+{
+    LONG rc;
+    PMEMKEY LinkKey, TargetKey;
+    PREPARSE_POINT ReparsePoint;
+
+    ReparsePoint = (PREPARSE_POINT)malloc(sizeof(*ReparsePoint));
+    if (!ReparsePoint)
+        return FALSE;
+
+    if (LinkKeyPath && !(LinkKeyHandle && *LinkKeyHandle))
+    {
+        /* Create the link key */
+        rc = RegCreateKeyExW(NULL,
+                             LinkKeyPath,
+                             0,
+                             NULL,
+                             REG_OPTION_VOLATILE,
+                             0,
+                             NULL,
+                             (PHKEY)&LinkKey,
+                             NULL);
+        if (rc != ERROR_SUCCESS)
+        {
+            free(ReparsePoint);
+            return FALSE;
+        }
+    }
+    else if (LinkKeyHandle)
+    {
+        /* Use the user-provided link key handle */
+        LinkKey = HKEY_TO_MEMKEY(*LinkKeyHandle);
+    }
+
+    if (LinkKeyHandle)
+        *LinkKeyHandle = MEMKEY_TO_HKEY(LinkKey);
+
+    TargetKey = HKEY_TO_MEMKEY(TargetKeyHandle);
+
+    ReparsePoint->SourceHive = LinkKey->RegistryHive;
+    ReparsePoint->SourceKeyCellOffset = LinkKey->KeyCellOffset;
+    ReparsePoint->DestinationHive = TargetKey->RegistryHive;
+    ReparsePoint->DestinationKeyCellOffset = TargetKey->KeyCellOffset;
+    InsertTailList(&CmiReparsePointsHead, &ReparsePoint->ListEntry);
+
+    return TRUE;
+}
 
 VOID
-RegInitializeRegistry(VOID)
+RegInitializeRegistry(
+    IN PCSTR HiveList)
 {
-    UNICODE_STRING RootKeyName = RTL_CONSTANT_STRING(L"\\");
     NTSTATUS Status;
-    PMEMKEY ControlSetKey, CurrentControlSetKey;
-    PREPARSE_POINT ReparsePoint;
+    UINT i;
+    HKEY ControlSetKey;
 
     InitializeListHead(&CmiHiveListHead);
     InitializeListHead(&CmiReparsePointsHead);
@@ -847,70 +1151,58 @@ RegInitializeRegistry(VOID)
     RootKey = CreateInMemoryStructure(&RootHive,
                                       RootHive.Hive.BaseBlock->RootCell);
 
-    /* Create DEFAULT key */
-    ConnectRegistry(NULL,
-                    &DefaultHive,
-                    SystemSecurity, sizeof(SystemSecurity),
-                    L"Registry\\User\\.DEFAULT");
-
-    /* Create SAM key */
-    ConnectRegistry(NULL,
-                    &SamHive,
-                    SystemSecurity, sizeof(SystemSecurity),
-                    L"Registry\\Machine\\SAM");
-
-    /* Create SECURITY key */
-    ConnectRegistry(NULL,
-                    &SecurityHive,
-                    NULL, 0,
-                    L"Registry\\Machine\\SECURITY");
-
-    /* Create SOFTWARE key */
-    ConnectRegistry(NULL,
-                    &SoftwareHive,
-                    SoftwareSecurity, sizeof(SoftwareSecurity),
-                    L"Registry\\Machine\\SOFTWARE");
-
-    /* Create BCD key */
-    ConnectRegistry(NULL,
-                    &BcdHive,
-                    BcdSecurity, sizeof(BcdSecurity),
-                    L"Registry\\Machine\\BCD00000000");
-
-    /* Create SYSTEM key */
-    ConnectRegistry(NULL,
-                    &SystemHive,
-                    SystemSecurity, sizeof(SystemSecurity),
-                    L"Registry\\Machine\\SYSTEM");
-
-    /* Create 'ControlSet001' key */
+    for (i = 0; i < _countof(RegistryHives); ++i)
+    {
+        /* Skip this registry hive if it's not in the list */
+        if (!strstr(HiveList, RegistryHives[i].HiveName))
+            continue;
+
+        /* Create the registry key */
+        ConnectRegistry(NULL,
+                        RegistryHives[i].CmHive,
+                        RegistryHives[i].SecurityDescriptor,
+                        RegistryHives[i].SecurityDescriptorLength,
+                        RegistryHives[i].HiveRegistryPath);
+
+        /* If we happen to deal with the special setup registry hive, stop there */
+        // if (strcmp(RegistryHives[i].HiveName, "SETUPREG") == 0)
+        if (i == 0)
+            break;
+    }
+
+    /* Create the 'ControlSet001' key */
     RegCreateKeyW(NULL,
                   L"Registry\\Machine\\SYSTEM\\ControlSet001",
-                  (HKEY*)&ControlSetKey);
-
-    /* Create 'CurrentControlSet' key */
-    RegCreateKeyExW(NULL,
-                    L"Registry\\Machine\\SYSTEM\\CurrentControlSet",
-                    0,
-                    NULL,
-                    REG_OPTION_VOLATILE,
-                    0,
-                    NULL,
-                    (HKEY*)&CurrentControlSetKey,
-                    NULL);
-
-    /* Connect 'CurrentControlSet' to 'ControlSet001' */
-    ReparsePoint = (PREPARSE_POINT)malloc(sizeof(REPARSE_POINT));
-    ReparsePoint->SourceHive = CurrentControlSetKey->RegistryHive;
-    ReparsePoint->SourceKeyCellOffset = CurrentControlSetKey->KeyCellOffset;
-    ReparsePoint->DestinationHive = ControlSetKey->RegistryHive;
-    ReparsePoint->DestinationKeyCellOffset = ControlSetKey->KeyCellOffset;
-    InsertTailList(&CmiReparsePointsHead, &ReparsePoint->ListEntry);
+                  &ControlSetKey);
+
+    /* Create the 'CurrentControlSet' key as a symlink to 'ControlSet001' */
+    CreateSymLink(L"Registry\\Machine\\SYSTEM\\CurrentControlSet",
+                  NULL, ControlSetKey);
+
+    RegCloseKey(ControlSetKey);
+
+#if 0
+    /* Link SECURITY to SAM */
+    CmpLinkKeyToHive(L"\\Registry\\Machine\\Security\\SAM", L"\\Registry\\Machine\\SAM\\SAM");
+    /* Link S-1-5-18 to .Default */
+    CmpLinkKeyToHive(L"\\Registry\\User\\S-1-5-18", L"\\Registry\\User\\.Default");
+#endif
 }
 
 VOID
 RegShutdownRegistry(VOID)
 {
+    PLIST_ENTRY Entry;
+    PREPARSE_POINT ReparsePoint;
+
+    /* Clean up the reparse points list */
+    while (!IsListEmpty(&CmiReparsePointsHead))
+    {
+        Entry = RemoveHeadList(&CmiReparsePointsHead);
+        ReparsePoint = CONTAINING_RECORD(Entry, REPARSE_POINT, ListEntry);
+        free(ReparsePoint);
+    }
+
     /* FIXME: clean up the complete hive */
 
     free(RootKey);