xref: /reactos/boot/environ/lib/misc/bootreg.c (revision b09b5584) 
1 /*
2  * COPYRIGHT:       See COPYING.ARM in the top level directory
3  * PROJECT:         ReactOS UEFI Boot Library
4  * FILE:            boot/environ/lib/misc/bootreg.c
5  * PURPOSE:         Boot Library Boot Registry Wrapper for CMLIB
6  * PROGRAMMER:      Alex Ionescu (alex.ionescu@reactos.org)
7  */
8 
9 /* INCLUDES ******************************************************************/
10 
11 #include "bl.h"
12 #include <bcd.h>
13 
14 /* DEFINITIONS ***************************************************************/
15 
16 #define BI_FLUSH_HIVE       0x01
17 #define BI_HIVE_WRITEABLE   0x02
18 
19 /* DATA STRUCTURES ***********************************************************/
20 
21 typedef struct _BI_KEY_HIVE
22 {
23     PHBASE_BLOCK BaseBlock;
24     ULONG HiveSize;
25     PBL_FILE_PATH_DESCRIPTOR FilePath;
26     CMHIVE Hive;
27     LONG ReferenceCount;
28     ULONG Flags;
29     PCM_KEY_NODE RootNode;
30 } BI_KEY_HIVE, *PBI_KEY_HIVE;
31 
32 typedef struct _BI_KEY_OBJECT
33 {
34     PBI_KEY_HIVE KeyHive;
35     PCM_KEY_NODE KeyNode;
36     HCELL_INDEX KeyCell;
37     PWCHAR KeyName;
38 } BI_KEY_OBJECT, *PBI_KEY_OBJECT;
39 
40 /* GLOBALS *******************************************************************/
41 
42 BOOLEAN BiHiveHashLibraryInitialized;
43 ULONGLONG HvSymcryptSeed;
44 
45 /* FUNCTIONS *****************************************************************/
46 
47 BOOLEAN
48 HvIsInPlaceBaseBlockValid (
49     _In_ PHBASE_BLOCK BaseBlock
50     )
51 {
52     ULONG HiveLength, HeaderSum;
53     BOOLEAN Valid;
54 
55     /* Assume failure */
56     Valid = FALSE;
57 
58     /* Check for incorrect signature, type, version, or format */
59     if ((BaseBlock->Signature == 'fger') &&
60         (BaseBlock->Type == 0) &&
61         (BaseBlock->Major <= 1) &&
62         (BaseBlock->Minor <= 5) &&
63         (BaseBlock->Minor >= 3) &&
64         (BaseBlock->Format == 1))
65     {
66         /* Check for invalid hive size */
67         HiveLength = BaseBlock->Length;
68         if (HiveLength)
69         {
70             /* Check for misaligned or too large hive size */
71             if (!(HiveLength & 0xFFF) && HiveLength <= 0x7FFFE000)
72             {
73                 /* Check for invalid header checksum */
74                 HeaderSum = HvpHiveHeaderChecksum(BaseBlock);
75                 if (HeaderSum == BaseBlock->CheckSum)
76                 {
77                     /* All good */
78                     Valid = TRUE;
79                 }
80             }
81         }
82     }
83 
84     /* Return validity */
85     return Valid;
86 }
87 
88 PVOID
89 NTAPI
90 CmpAllocate (
91     _In_ SIZE_T Size,
92     _In_ BOOLEAN Paged,
93     _In_ ULONG Tag
94     )
95 {
96     UNREFERENCED_PARAMETER(Paged);
97     UNREFERENCED_PARAMETER(Tag);
98 
99     /* Call the heap allocator */
100     return BlMmAllocateHeap(Size);
101 }
102 
103 VOID
104 NTAPI
105 CmpFree (
106     _In_ PVOID Ptr,
107     _In_ ULONG Quota
108     )
109 {
110     UNREFERENCED_PARAMETER(Quota);
111 
112     /* Call the heap allocator */
113     BlMmFreeHeap(Ptr);
114 }
115 
116 VOID
117 BiDereferenceHive (
118     _In_ HANDLE KeyHandle
119     )
120 {
121     PBI_KEY_OBJECT KeyObject;
122 
123     /* Get the key object */
124     KeyObject = (PBI_KEY_OBJECT)KeyHandle;
125 
126     /* Drop a reference on the parent hive */
127     --KeyObject->KeyHive->ReferenceCount;
128 }
129 
130 VOID
131 BiFlushHive (
132     _In_ HANDLE KeyHandle
133     )
134 {
135     /* Not yet implemented */
136     EfiPrintf(L"NO reg flush\r\n");
137     return;
138 }
139 
140 VOID
141 BiCloseKey (
142     _In_ HANDLE KeyHandle
143     )
144 {
145     PBI_KEY_HIVE KeyHive;
146     PBI_KEY_OBJECT KeyObject;
147 
148     /* Get the key object and hive */
149     KeyObject = (PBI_KEY_OBJECT)KeyHandle;
150     KeyHive = KeyObject->KeyHive;
151 
152     /* Check if we have a hive, or name, or key node */
153     if ((KeyHive) || (KeyObject->KeyNode) || (KeyObject->KeyName))
154     {
155         /* Drop a reference, see if it's the last one */
156         BiDereferenceHive(KeyHandle);
157         if (!KeyHive->ReferenceCount)
158         {
159             /* Check if we should flush it */
160             if (KeyHive->Flags & BI_FLUSH_HIVE)
161             {
162                 BiFlushHive(KeyHandle);
163             }
164 
165             /* Unmap the hive */
166             MmPapFreePages(KeyHive->BaseBlock, BL_MM_INCLUDE_MAPPED_ALLOCATED);
167 
168             /* Free the hive and hive path */
169             BlMmFreeHeap(KeyHive->FilePath);
170             BlMmFreeHeap(KeyHive);
171         }
172 
173         /* Check if a key name is present */
174         if (KeyObject->KeyName)
175         {
176             /* Free it */
177             BlMmFreeHeap(KeyObject->KeyName);
178         }
179     }
180 
181     /* Free the object */
182     BlMmFreeHeap(KeyObject);
183 }
184 
185 NTSTATUS
186 BiOpenKey(
187     _In_ HANDLE ParentHandle,
188     _In_ PWCHAR KeyName,
189     _Out_ PHANDLE Handle
190     )
191 {
192     PBI_KEY_OBJECT ParentKey, NewKey;
193     PBI_KEY_HIVE ParentHive;
194     NTSTATUS Status;
195     SIZE_T NameLength, SubNameLength, NameBytes;
196     PWCHAR NameStart, NameBuffer;
197     UNICODE_STRING KeyString;
198     HCELL_INDEX KeyCell;
199     PHHIVE Hive;
200     PCM_KEY_NODE ParentNode;
201 
202     /* Convert from a handle to our key object */
203     ParentKey = (PBI_KEY_OBJECT)ParentHandle;
204 
205     /* Extract the hive and node information */
206     ParentHive = ParentKey->KeyHive;
207     ParentNode = ParentKey->KeyNode;
208     Hive = &ParentKey->KeyHive->Hive.Hive;
209 
210     /* Initialize variables */
211     KeyCell = HCELL_NIL;
212     Status = STATUS_SUCCESS;
213     NameBuffer = NULL;
214 
215     /* Loop as long as there's still portions of the key name in play */
216     NameLength = wcslen(KeyName);
217     while (NameLength)
218     {
219         /* Find the first path separator */
220         NameStart = wcschr(KeyName, OBJ_NAME_PATH_SEPARATOR);
221         if (NameStart)
222         {
223             /* Look only at the key before the separator */
224             SubNameLength = NameStart - KeyName;
225             ++NameStart;
226         }
227         else
228         {
229             /* No path separator, this is the final leaf key */
230             SubNameLength = NameLength;
231         }
232 
233         /* Free the name buffer from the previous pass if needed */
234         if (NameBuffer)
235         {
236             BlMmFreeHeap(NameBuffer);
237         }
238 
239         /* Allocate a buffer to hold the name of this specific subkey only */
240         NameBytes = SubNameLength * sizeof(WCHAR);
241         NameBuffer = BlMmAllocateHeap(NameBytes + sizeof(UNICODE_NULL));
242         if (!NameBuffer)
243         {
244             Status = STATUS_NO_MEMORY;
245             goto Quickie;
246         }
247 
248         /* Copy and null-terminate the name of the subkey */
249         RtlCopyMemory(NameBuffer, KeyName, NameBytes);
250         NameBuffer[SubNameLength] = UNICODE_NULL;
251 
252         /* Convert it into a UNICODE_STRING and try to find it */
253         RtlInitUnicodeString(&KeyString, NameBuffer);
254         KeyCell = CmpFindSubKeyByName(Hive, ParentNode, &KeyString);
255         if (KeyCell == HCELL_NIL)
256         {
257             Status = STATUS_OBJECT_NAME_NOT_FOUND;
258             goto Quickie;
259         }
260 
261         /* We found it -- get the key node out of it */
262         ParentNode = (PCM_KEY_NODE)HvGetCell(Hive, KeyCell);
263         if (!ParentNode)
264         {
265             Status = STATUS_REGISTRY_CORRUPT;
266             goto Quickie;
267         }
268 
269         /* Update the key name to the next remaining path element */
270         KeyName = NameStart;
271         if (NameStart)
272         {
273             /* Update the length to the remainder of the path */
274             NameLength += -1 - SubNameLength;
275         }
276         else
277         {
278             /* There's nothing left, this was the leaf key */
279             NameLength = 0;
280         }
281     }
282 
283     /* Allocate a key object */
284     NewKey = BlMmAllocateHeap(sizeof(*NewKey));
285     if (!NewKey)
286     {
287         /* Bail out if we had no memory for it */
288         Status = STATUS_NO_MEMORY;
289         goto Quickie;
290     }
291 
292     /* Fill out the key object data */
293     NewKey->KeyNode = ParentNode;
294     NewKey->KeyHive = ParentHive;
295     NewKey->KeyName = NameBuffer;
296     NewKey->KeyCell = KeyCell;
297 
298     /* Add a reference to the hive */
299     ++ParentHive->ReferenceCount;
300 
301     /* Return the object back to the caller */
302     *Handle = NewKey;
303 
304 Quickie:
305     /* If we had a name buffer, free it */
306     if (NameBuffer)
307     {
308         BlMmFreeHeap(NameBuffer);
309     }
310 
311     /* Return status of the open operation */
312     return Status;
313 }
314 
315 NTSTATUS
316 BiInitializeAndValidateHive (
317     _In_ PBI_KEY_HIVE Hive
318     )
319 {
320     ULONG HiveSize;
321     CM_CHECK_REGISTRY_STATUS CmStatusCode;
322     NTSTATUS Status;
323 
324     /* Make sure the hive is at least the size of a base block */
325     if (Hive->HiveSize < sizeof(HBASE_BLOCK))
326     {
327         return STATUS_REGISTRY_CORRUPT;
328     }
329 
330     /* Make sure that the base block accurately describes the size of the hive */
331     HiveSize = Hive->BaseBlock->Length + sizeof(HBASE_BLOCK);
332     if ((HiveSize < sizeof(HBASE_BLOCK)) || (HiveSize > Hive->HiveSize))
333     {
334         return STATUS_REGISTRY_CORRUPT;
335     }
336 
337     /* Initialize a flat memory hive */
338     RtlZeroMemory(&Hive->Hive, sizeof(Hive->Hive));
339     Status = HvInitialize(&Hive->Hive.Hive,
340                           HINIT_FLAT,
341                           0,
342                           0,
343                           Hive->BaseBlock,
344                           CmpAllocate,
345                           CmpFree,
346                           NULL,
347                           NULL,
348                           NULL,
349                           NULL,
350                           0,
351                           NULL);
352     if (NT_SUCCESS(Status))
353     {
354         /* Cleanup volatile/old data */
355         CmStatusCode = CmCheckRegistry(Hive->Hive, CM_CHECK_REGISTRY_BOOTLOADER_PURGE_VOLATILES | CM_CHECK_REGISTRY_VALIDATE_HIVE);
356         if (!CM_CHECK_REGISTRY_SUCCESS(CmStatusCode))
357         {
358             return STATUS_REGISTRY_CORRUPT;
359         }
360 
361         Status = STATUS_SUCCESS;
362     }
363 
364     /* Return the final status */
365     return Status;
366 }
367 
368 NTSTATUS
369 BiLoadHive (
370     _In_ PBL_FILE_PATH_DESCRIPTOR FilePath,
371     _Out_ PHANDLE HiveHandle
372     )
373 {
374     ULONG DeviceId;
375     PHBASE_BLOCK BaseBlock, NewBaseBlock;
376     PBI_KEY_OBJECT KeyObject;
377     PBI_KEY_HIVE BcdHive;
378     PBL_DEVICE_DESCRIPTOR BcdDevice;
379     ULONG PathLength, DeviceLength, HiveSize, HiveLength, NewHiveSize;
380     PWCHAR HiveName, LogName;
381     BOOLEAN HaveWriteAccess;
382     NTSTATUS Status;
383     PVOID LogData;
384     PHHIVE Hive;
385     UNICODE_STRING KeyString;
386     PCM_KEY_NODE RootNode;
387     HCELL_INDEX CellIndex;
388 
389     /* Initialize variables */
390     DeviceId = -1;
391     BaseBlock = NULL;
392     BcdHive = NULL;
393     KeyObject = NULL;
394     LogData = NULL;
395     LogName = NULL;
396 
397     /* Initialize the crypto seed */
398     if (!BiHiveHashLibraryInitialized)
399     {
400         HvSymcryptSeed = 0x82EF4D887A4E55C5;
401         BiHiveHashLibraryInitialized = TRUE;
402     }
403 
404     /* Extract and validate the input path */
405     BcdDevice = (PBL_DEVICE_DESCRIPTOR)&FilePath->Path;
406     PathLength = FilePath->Length;
407     DeviceLength = BcdDevice->Size;
408     HiveName = (PWCHAR)((ULONG_PTR)BcdDevice + BcdDevice->Size);
409     if (PathLength <= DeviceLength)
410     {
411         /* Doesn't make sense, bail out */
412         Status = STATUS_INVALID_PARAMETER;
413         goto Quickie;
414     }
415 
416     /* Attempt to open the underlying device for RW access */
417     HaveWriteAccess = TRUE;
418     Status = BlpDeviceOpen(BcdDevice,
419                            BL_DEVICE_READ_ACCESS | BL_DEVICE_WRITE_ACCESS,
420                            0,
421                            &DeviceId);
422     if (!NT_SUCCESS(Status))
423     {
424         /* Try for RO access instead */
425         HaveWriteAccess = FALSE;
426         Status = BlpDeviceOpen(BcdDevice, BL_DEVICE_READ_ACCESS, 0, &DeviceId);
427         if (!NT_SUCCESS(Status))
428         {
429             /* No access at all -- bail out */
430             goto Quickie;
431         }
432     }
433 
434     /* Now try to load the hive on disk */
435     Status = BlImgLoadImageWithProgress2(DeviceId,
436                                          BlLoaderRegistry,
437                                          HiveName,
438                                          (PVOID*)&BaseBlock,
439                                          &HiveSize,
440                                          0,
441                                          FALSE,
442                                          NULL,
443                                          NULL);
444     if (!NT_SUCCESS(Status))
445     {
446         EfiPrintf(L"Hive read failure: % lx\r\n", Status);
447         goto Quickie;
448     }
449 
450     /* Allocate a hive structure */
451     BcdHive = BlMmAllocateHeap(sizeof(*BcdHive));
452     if (!BcdHive)
453     {
454         Status = STATUS_NO_MEMORY;
455         goto Quickie;
456     }
457 
458     /* Initialize it */
459     RtlZeroMemory(BcdHive, sizeof(*BcdHive));
460     BcdHive->BaseBlock = BaseBlock;
461     BcdHive->HiveSize = HiveSize;
462     if (HaveWriteAccess)
463     {
464         BcdHive->Flags |= BI_HIVE_WRITEABLE;
465     }
466 
467     /* Make sure the hive was at least one bin long */
468     if (HiveSize < sizeof(*BaseBlock))
469     {
470         Status = STATUS_REGISTRY_CORRUPT;
471         goto Quickie;
472     }
473 
474     /* Make sure the hive contents are at least one bin long */
475     HiveLength = BaseBlock->Length;
476     if (BaseBlock->Length < sizeof(*BaseBlock))
477     {
478         Status = STATUS_REGISTRY_CORRUPT;
479         goto Quickie;
480     }
481 
482     /* Validate the initial bin (the base block) */
483     if (!HvIsInPlaceBaseBlockValid(BaseBlock))
484     {
485         EfiPrintf(L"Recovery not implemented\r\n");
486         Status = STATUS_REGISTRY_CORRUPT;
487         goto Quickie;
488     }
489 
490     /* Check if there's log recovery that needs to happen */
491     if (BaseBlock->Sequence1 != BaseBlock->Sequence2)
492     {
493         EfiPrintf(L"Log fix not implemented: %lx %lx\r\n");
494         Status = STATUS_REGISTRY_CORRUPT;
495         goto Quickie;
496     }
497 
498     /*
499      * Check if the whole hive doesn't fit in the buffer.
500      * Note: HiveLength does not include the size of the baseblock itself
501      */
502     if (HiveSize < (HiveLength + sizeof(*BaseBlock)))
503     {
504         EfiPrintf(L"Need bigger hive buffer path\r\n");
505 
506         /* Allocate a slightly bigger buffer */
507         NewHiveSize = HiveLength + sizeof(*BaseBlock);
508         Status = MmPapAllocatePagesInRange((PVOID*)&NewBaseBlock,
509                                            BlLoaderRegistry,
510                                            NewHiveSize >> PAGE_SHIFT,
511                                            0,
512                                            0,
513                                            NULL,
514                                            0);
515         if (!NT_SUCCESS(Status))
516         {
517             goto Quickie;
518         }
519 
520         /* Copy the current data in there */
521         RtlCopyMemory(NewBaseBlock, BaseBlock, HiveSize);
522 
523         /* Free the old data */
524         MmPapFreePages(BaseBlock, BL_MM_INCLUDE_MAPPED_ALLOCATED);
525 
526         /* Update our pointers */
527         BaseBlock = NewBaseBlock;
528         HiveSize = NewHiveSize;
529         BcdHive->BaseBlock = BaseBlock;
530         BcdHive->HiveSize = HiveSize;
531     }
532 
533     /* Check if any log stuff needs to happen */
534     if (LogData)
535     {
536         EfiPrintf(L"Log fix not implemented: %lx %lx\r\n");
537         Status = STATUS_REGISTRY_CORRUPT;
538         goto Quickie;
539     }
540 
541     /* Call Hv to setup the hive library */
542     Status = BiInitializeAndValidateHive(BcdHive);
543     if (!NT_SUCCESS(Status))
544     {
545         goto Quickie;
546     }
547 
548     /* Now get the root node */
549     Hive = &BcdHive->Hive.Hive;
550     RootNode = (PCM_KEY_NODE)HvGetCell(Hive, Hive->BaseBlock->RootCell);
551     if (!RootNode)
552     {
553         Status = STATUS_OBJECT_NAME_NOT_FOUND;
554         goto Quickie;
555     }
556 
557     /* Find the Objects subkey under it to see if it's a real BCD hive */
558     RtlInitUnicodeString(&KeyString, L"Objects");
559     CellIndex = CmpFindSubKeyByName(Hive, RootNode, &KeyString);
560     if (CellIndex == HCELL_NIL)
561     {
562         EfiPrintf(L"No OBJECTS subkey found!\r\n");
563         Status = STATUS_OBJECT_NAME_NOT_FOUND;
564         goto Quickie;
565     }
566 
567     /* This is a valid BCD hive, store its root node here */
568     BcdHive->RootNode = RootNode;
569 
570     /* Allocate a copy of the file path */
571     BcdHive->FilePath = BlMmAllocateHeap(FilePath->Length);
572     if (!BcdHive->FilePath)
573     {
574         Status = STATUS_NO_MEMORY;
575         goto Quickie;
576     }
577 
578     /* Make a copy of it */
579     RtlCopyMemory(BcdHive->FilePath, FilePath, FilePath->Length);
580 
581     /* Create a key object to describe the rot */
582     KeyObject = BlMmAllocateHeap(sizeof(*KeyObject));
583     if (!KeyObject)
584     {
585         Status = STATUS_NO_MEMORY;
586         goto Quickie;
587     }
588 
589     /* Fill out the details */
590     KeyObject->KeyNode = RootNode;
591     KeyObject->KeyHive = BcdHive;
592     KeyObject->KeyName = NULL;
593     KeyObject->KeyCell = Hive->BaseBlock->RootCell;
594 
595     /* One reference for the key object, plus one lifetime reference */
596     BcdHive->ReferenceCount = 2;
597 
598     /* This is the hive handle */
599     *HiveHandle  = KeyObject;
600 
601     /* We're all good */
602     Status = STATUS_SUCCESS;
603 
604 Quickie:
605     /* If we had a log name, free it */
606     if (LogName)
607     {
608         BlMmFreeHeap(LogName);
609     }
610 
611     /* If we had logging data, free it */
612     if (LogData)
613     {
614         MmPapFreePages(LogData, BL_MM_INCLUDE_MAPPED_ALLOCATED);
615     }
616 
617     /* Check if this is the failure path */
618     if (!NT_SUCCESS(Status))
619     {
620         /* If we mapped the hive, free it */
621         if (BaseBlock)
622         {
623             MmPapFreePages(BaseBlock, BL_MM_INCLUDE_MAPPED_ALLOCATED);
624         }
625 
626         /* If we opened the device, close it */
627         if (DeviceId != -1)
628         {
629             BlDeviceClose(DeviceId);
630         }
631 
632         /* Did we create a hive object? */
633         if (BcdHive)
634         {
635             /* Free the file path if we made a copy of it */
636             if (BcdHive->FilePath)
637             {
638                 BlMmFreeHeap(BcdHive->FilePath);
639             }
640 
641             /* Free the hive itself */
642             BlMmFreeHeap(BcdHive);
643         }
644 
645         /* Finally, free the root key object if we created one */
646         if (KeyObject)
647         {
648             BlMmFreeHeap(KeyObject);
649         }
650     }
651 
652     /* Return the final status */
653     return Status;
654 }
655 
656 NTSTATUS
657 BiGetRegistryValue (
658     _In_ HANDLE KeyHandle,
659     _In_ PWCHAR ValueName,
660     _In_ ULONG Type,
661     _Out_ PVOID* Buffer,
662     _Out_ PULONG ValueLength
663     )
664 {
665     PCM_KEY_NODE KeyNode;
666     PHHIVE KeyHive;
667     UNICODE_STRING ValueString;
668     PBI_KEY_OBJECT KeyObject;
669     PCM_KEY_VALUE KeyValue;
670     PVOID ValueCopy;
671     ULONG Size;
672     HCELL_INDEX CellIndex;
673     PCELL_DATA ValueData;
674 
675     /* Get the key object, node,and hive */
676     KeyObject = (PBI_KEY_OBJECT)KeyHandle;
677     KeyNode = KeyObject->KeyNode;
678     KeyHive = &KeyObject->KeyHive->Hive.Hive;
679 
680     /* Find the value cell index in the list of values */
681     RtlInitUnicodeString(&ValueString, ValueName);
682     CmpFindNameInList(KeyHive,
683                       &KeyNode->ValueList,
684                       &ValueString,
685                       NULL,
686                       &CellIndex);
687     if (CellIndex == HCELL_NIL)
688     {
689         return STATUS_OBJECT_NAME_NOT_FOUND;
690     }
691 
692     /* Get the cell data for it */
693     KeyValue = (PCM_KEY_VALUE)HvGetCell(KeyHive, CellIndex);
694     if (!KeyValue)
695     {
696         return STATUS_REGISTRY_CORRUPT;
697     }
698 
699     /* Make sure the type matches */
700     if (KeyValue->Type != Type)
701     {
702         return STATUS_OBJECT_TYPE_MISMATCH;
703     }
704 
705     /* Now get the data cell */
706     ValueData = CmpValueToData(KeyHive, KeyValue, &Size);
707 
708     /* Make a copy of it */
709     ValueCopy = BlMmAllocateHeap(Size);
710     if (!ValueCopy)
711     {
712         return STATUS_NO_MEMORY;
713     }
714 
715     /* Copy it in the buffer, and return it and its size */
716     RtlCopyMemory(ValueCopy, ValueData, Size);
717     *Buffer = ValueCopy;
718     *ValueLength = Size;
719     return STATUS_SUCCESS;
720 }
721 
722 NTSTATUS
723 BiEnumerateSubKeys (
724     _In_ HANDLE KeyHandle,
725     _Out_ PWCHAR** SubKeyList,
726     _Out_ PULONG SubKeyCount
727     )
728 {
729     PCM_KEY_NODE KeyNode, Node;
730     PBI_KEY_OBJECT KeyObject;
731     ULONG KeyCount;
732     ULONG NameLength, NewTotalNameLength, FinalLength, TotalNameLength;
733     PHHIVE Hive;
734     PWCHAR KeyName, NameEnd;
735     HCELL_INDEX CellIndex;
736     PWCHAR* SubKeys;
737     NTSTATUS Status;
738     ULONG i;
739 
740     /* Get the key object, node, and hive */
741     KeyObject = (PBI_KEY_OBJECT)KeyHandle;
742     KeyNode = KeyObject->KeyNode;
743     Hive = &KeyObject->KeyHive->Hive.Hive;
744 
745     /* Assume it's empty */
746     *SubKeyList = 0;
747     *SubKeyCount = 0;
748 
749     /* Initialize locals */
750     KeyCount = 0;
751     SubKeys = 0;
752     TotalNameLength = 0;
753 
754     /* Find the first subkey cell index */
755     CellIndex = CmpFindSubKeyByNumber(Hive, KeyNode, KeyCount);
756     while (CellIndex != HCELL_NIL)
757     {
758         /* Move to the next one */
759         KeyCount++;
760 
761         /* Get the cell data for it */
762         Node = (PCM_KEY_NODE)HvGetCell(Hive, CellIndex);
763         if (!Node)
764         {
765             return STATUS_REGISTRY_CORRUPT;
766         }
767 
768         /* Check if the value is compressed */
769         if (Node->Flags & KEY_COMP_NAME)
770         {
771             /* Get the compressed name size */
772             NameLength = CmpCompressedNameSize(Node->Name, Node->NameLength);
773         }
774         else
775         {
776             /* Get the real size */
777             NameLength = Node->NameLength;
778         }
779 
780         /* Add up the new length, protecting against overflow */
781         NewTotalNameLength = TotalNameLength + NameLength + sizeof(UNICODE_NULL);
782         if (NewTotalNameLength < TotalNameLength)
783         {
784             Status = STATUS_NAME_TOO_LONG;
785             goto Quickie;
786         }
787 
788         /* We're good, use the new length */
789         TotalNameLength = NewTotalNameLength;
790 
791         /* Find the next subkey cell index */
792         CellIndex = CmpFindSubKeyByNumber(Hive, KeyNode, KeyCount);
793     }
794 
795     /* Were there no keys? We're done, if so */
796     if (!KeyCount)
797     {
798         return STATUS_SUCCESS;
799     }
800 
801     /* Safely compute the size of the array needed */
802     Status = RtlULongLongToULong(sizeof(PWCHAR) * KeyCount, &FinalLength);
803     if (!NT_SUCCESS(Status))
804     {
805         goto Quickie;
806     }
807 
808     /* Safely add that to the name length */
809     Status = RtlULongAdd(TotalNameLength, FinalLength, &FinalLength);
810     if (!NT_SUCCESS(Status))
811     {
812         goto Quickie;
813     }
814 
815     /* Allocate an array big enough for the names and pointers */
816     SubKeys = BlMmAllocateHeap(FinalLength);
817     if (!SubKeys)
818     {
819         Status = STATUS_NO_MEMORY;
820         goto Quickie;
821     }
822 
823     /* Go over each key again */
824     NameEnd = (PWCHAR)&SubKeys[KeyCount];
825     for (i = 0; i < KeyCount; i++)
826     {
827         /* Get the cell index for this subkey */
828         CellIndex = CmpFindSubKeyByNumber(Hive, KeyNode, i);
829         if (CellIndex == HCELL_NIL)
830         {
831             break;
832         }
833 
834         /* Get the cell data for it */
835         Node = (PCM_KEY_NODE)HvGetCell(Hive, CellIndex);
836         if (!Node)
837         {
838             Status = STATUS_REGISTRY_CORRUPT;
839             goto Quickie;
840         }
841 
842         /* Check if the value is compressed */
843         KeyName = Node->Name;
844         if (Node->Flags & KEY_COMP_NAME)
845         {
846             /* Get the compressed name size */
847             NameLength = CmpCompressedNameSize(KeyName, Node->NameLength);
848             CmpCopyCompressedName(NameEnd, NameLength, KeyName, Node->NameLength);
849         }
850         else
851         {
852             /* Get the real size */
853             NameLength = Node->NameLength;
854             RtlCopyMemory(NameEnd, KeyName, NameLength);
855         }
856 
857         /* Move the name buffer to the next spot, and NULL-terminate */
858         SubKeys[i] = NameEnd;
859         NameEnd += (NameLength / sizeof(WCHAR));
860         *NameEnd = UNICODE_NULL;
861 
862         /* Keep going */
863         NameEnd++;
864     }
865 
866     /* Check if the subkeys were empty */
867     if (i == 0)
868     {
869         /* They disappeared in the middle of enumeration */
870         Status = STATUS_OBJECT_NAME_NOT_FOUND;
871         goto Quickie;
872     }
873 
874     /* Return the count and the array of names */
875     *SubKeyList = SubKeys;
876     *SubKeyCount = i;
877     SubKeys = NULL;
878     Status = STATUS_SUCCESS;
879 
880 Quickie:
881     /* On the failure path, free the subkeys if any exist */
882     if (SubKeys)
883     {
884         BlMmFreeHeap(SubKeys);
885     }
886 
887     /* All done, return the result */
888     return Status;
889 }
890 
891 NTSTATUS
892 BiDeleteKey (
893     _In_ HANDLE KeyHandle
894     )
895 {
896     NTSTATUS Status;
897     PBI_KEY_OBJECT KeyObject;
898     PHHIVE Hive;
899     ULONG SubKeyCount, i;
900     PWCHAR* SubKeyList;
901     HANDLE SubKeyHandle;
902 
903     /* Get the key object and hive */
904     KeyObject = (PBI_KEY_OBJECT)KeyHandle;
905     Hive = &KeyObject->KeyHive->Hive.Hive;
906 
907     /* Make sure the hive is writeable */
908     if (!(KeyObject->KeyHive->Flags & BI_HIVE_WRITEABLE))
909     {
910         return STATUS_MEDIA_WRITE_PROTECTED;
911     }
912 
913     /* Enumerate all of the subkeys */
914     Status = BiEnumerateSubKeys(KeyHandle, &SubKeyList, &SubKeyCount);
915     if ((NT_SUCCESS(Status)) && (SubKeyCount > 0))
916     {
917         /* Loop through each one */
918         for (i = 0; i < SubKeyCount; i++)
919         {
920             /* Open a handle to it */
921             Status = BiOpenKey(KeyHandle, SubKeyList[i], &SubKeyHandle);
922             if (NT_SUCCESS(Status))
923             {
924                 /* Recursively call us to delete it */
925                 Status = BiDeleteKey(SubKeyHandle);
926                 if (Status != STATUS_SUCCESS)
927                 {
928                     /* Close the key on failure */
929                     BiCloseKey(SubKeyHandle);
930                 }
931             }
932         }
933     }
934 
935     /* Check if we had a list of subkeys */
936     if (SubKeyList)
937     {
938         /* Free it */
939         BlMmFreeHeap(SubKeyList);
940     }
941 
942     /* Delete this key cell */
943     Status = CmpFreeKeyByCell(Hive, KeyObject->KeyCell, TRUE);
944     if (NT_SUCCESS(Status))
945     {
946         /* Mark the hive as requiring a flush */
947         KeyObject->KeyHive->Flags |= BI_FLUSH_HIVE;
948         BiCloseKey(KeyHandle);
949     }
950 
951     /* All done */
952     return Status;
953 }
954