drivers/filesystems/ntfs/btree.c

   1 /*
   2 *  ReactOS kernel
   3 *  Copyright (C) 2002, 2017 ReactOS Team
   4 *
   5 *  This program is free software; you can redistribute it and/or modify
   6 *  it under the terms of the GNU General Public License as published by
   7 *  the Free Software Foundation; either version 2 of the License, or
   8 *  (at your option) any later version.
   9 *
  10 *  This program is distributed in the hope that it will be useful,
  11 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 *  GNU General Public License for more details.
  14 *
  15 *  You should have received a copy of the GNU General Public License
  16 *  along with this program; if not, write to the Free Software
  17 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.
  18 *
  19 * COPYRIGHT:        See COPYING in the top level directory
  20 * PROJECT:          ReactOS kernel
  21 * FILE:             drivers/filesystem/ntfs/btree.c
  22 * PURPOSE:          NTFS filesystem driver
  23 * PROGRAMMERS:      Trevor Thompson
  24 */
  25
  26 /* INCLUDES *****************************************************************/
  27
  28 #include "ntfs.h"
  29
  30 #define NDEBUG
  31 #include <debug.h>
  32
  33 /* FUNCTIONS ****************************************************************/
  34
  35 // TEMP FUNCTION for diagnostic purposes.
  36 // Prints VCN of every node in an index allocation
  37 VOID
  38 PrintAllVCNs(PDEVICE_EXTENSION Vcb,
  39              PNTFS_ATTR_CONTEXT IndexAllocationContext,
  40              ULONG NodeSize)
  41 {
  42     ULONGLONG CurrentOffset = 0;
  43     PINDEX_BUFFER CurrentNode, Buffer;
  44     ULONGLONG BufferSize = AttributeDataLength(IndexAllocationContext->pRecord);
  45     ULONG BytesRead;
  46     ULONGLONG i;
  47     int Count = 0;
  48
  49     if (BufferSize == 0)
  50     {
  51         DPRINT1("Index Allocation is empty.\n");
  52         return;
  53     }
  54
  55     Buffer = ExAllocatePoolWithTag(NonPagedPool, BufferSize, TAG_NTFS);
  56
  57     BytesRead = ReadAttribute(Vcb, IndexAllocationContext, 0, (PCHAR)Buffer, BufferSize);
  58
  59     ASSERT(BytesRead = BufferSize);
  60
  61     CurrentNode = Buffer;
  62
  63     // loop through all the nodes
  64     for (i = 0; i < BufferSize; i += NodeSize)
  65     {
  66         NTSTATUS Status = FixupUpdateSequenceArray(Vcb, &CurrentNode->Ntfs);
  67         if (!NT_SUCCESS(Status))
  68         {
  69             DPRINT1("ERROR: Fixing fixup failed!\n");
  70             continue;
  71         }
  72
  73         DPRINT1("Node #%d, VCN: %I64u\n", Count, CurrentNode->VCN);
  74
  75         CurrentNode = (PINDEX_BUFFER)((ULONG_PTR)CurrentNode + NodeSize);
  76         CurrentOffset += NodeSize;
  77         Count++;
  78     }
  79
  80     ExFreePoolWithTag(Buffer, TAG_NTFS);
  81 }
  82
  83 /**
  84 * @name AllocateIndexNode
  85 * @implemented
  86 *
  87 * Allocates a new index record in an index allocation.
  88 *
  89 * @param DeviceExt
  90 * Pointer to the target DEVICE_EXTENSION describing the volume the node will be created on.
  91 *
  92 * @param FileRecord
  93 * Pointer to a copy of the file record containing the index.
  94 *
  95 * @param IndexBufferSize
  96 * Size of an index record for this index, in bytes. Commonly defined as 4096.
  97 *
  98 * @param IndexAllocationCtx
  99 * Pointer to an NTFS_ATTR_CONTEXT describing the index allocation attribute the node will be assigned to.
 100 *
 101 * @param IndexAllocationOffset
 102 * Offset of the index allocation attribute relative to the file record.
 103 *
 104 * @param NewVCN
 105 * Pointer to a ULONGLONG which will receive the VCN of the newly-assigned index record
 106 *
 107 * @returns
 108 * STATUS_SUCCESS in case of success.
 109 * STATUS_NOT_IMPLEMENTED if there's no $I30 bitmap attribute in the file record.
 110 *
 111 * @remarks
 112 * AllocateIndexNode() doesn't write any data to the index record it creates. Called by UpdateIndexNode().
 113 * Don't call PrintAllVCNs() or NtfsDumpFileRecord() after calling AllocateIndexNode() before UpdateIndexNode() finishes.
 114 * Possible TODO: Create an empty node and write it to the allocated index node, so the index allocation is always valid.
 115 */
 116 NTSTATUS
 117 AllocateIndexNode(PDEVICE_EXTENSION DeviceExt,
 118                   PFILE_RECORD_HEADER FileRecord,
 119                   ULONG IndexBufferSize,
 120                   PNTFS_ATTR_CONTEXT IndexAllocationCtx,
 121                   ULONG IndexAllocationOffset,
 122                   PULONGLONG NewVCN)
 123 {
 124     NTSTATUS Status;
 125     PNTFS_ATTR_CONTEXT BitmapCtx;
 126     ULONGLONG IndexAllocationLength, BitmapLength;
 127     ULONG BitmapOffset;
 128     ULONGLONG NextNodeNumber;
 129     PCHAR *BitmapMem;
 130     ULONG *BitmapPtr;
 131     RTL_BITMAP Bitmap;
 132     ULONG BytesWritten;
 133     ULONG BytesNeeded;
 134     LARGE_INTEGER DataSize;
 135
 136     DPRINT1("AllocateIndexNode(%p, %p, %lu, %p, %lu, %p) called.\n", DeviceExt,
 137             FileRecord,
 138             IndexBufferSize,
 139             IndexAllocationCtx,
 140             IndexAllocationOffset,
 141             NewVCN);
 142
 143     // Get the length of the attribute allocation
 144     IndexAllocationLength = AttributeDataLength(IndexAllocationCtx->pRecord);
 145
 146     // Find the bitmap attribute for the index
 147     Status = FindAttribute(DeviceExt,
 148                            FileRecord,
 149                            AttributeBitmap,
 150                            L"$I30",
 151                            4,
 152                            &BitmapCtx,
 153                            &BitmapOffset);
 154     if (!NT_SUCCESS(Status))
 155     {
 156         DPRINT1("FIXME: Need to add bitmap attribute!\n");
 157         return STATUS_NOT_IMPLEMENTED;
 158     }
 159
 160     // Get the length of the bitmap attribute
 161     BitmapLength = AttributeDataLength(BitmapCtx->pRecord);
 162
 163     NextNodeNumber = IndexAllocationLength / DeviceExt->NtfsInfo.BytesPerIndexRecord;
 164
 165     // TODO: Find unused allocation in bitmap and use that space first
 166
 167     // Add another bit to bitmap
 168
 169     // See how many bytes we need to store the amount of bits we'll have
 170     BytesNeeded = NextNodeNumber / 8;
 171     BytesNeeded++;
 172
 173     // Windows seems to allocate the bitmap in 8-byte chunks to keep any bytes from being wasted on padding
 174     BytesNeeded = ALIGN_UP(BytesNeeded, ATTR_RECORD_ALIGNMENT);
 175
 176     // Allocate memory for the bitmap, including some padding; RtlInitializeBitmap() wants a pointer
 177     // that's ULONG-aligned, and it wants the size of the memory allocated for it to be a ULONG-multiple.
 178     BitmapMem = ExAllocatePoolWithTag(NonPagedPool, BytesNeeded + sizeof(ULONG), TAG_NTFS);
 179     if (!BitmapMem)
 180     {
 181         DPRINT1("Error: failed to allocate bitmap!");
 182         ReleaseAttributeContext(BitmapCtx);
 183         return STATUS_INSUFFICIENT_RESOURCES;
 184     }
 185     // RtlInitializeBitmap() wants a pointer that's ULONG-aligned.
 186     BitmapPtr = (PULONG)ALIGN_UP_BY((ULONG_PTR)BitmapMem, sizeof(ULONG));
 187
 188     RtlZeroMemory(BitmapPtr, BytesNeeded);
 189
 190     // Read the existing bitmap data
 191     Status = ReadAttribute(DeviceExt, BitmapCtx, 0, (PCHAR)BitmapPtr, BitmapLength);
 192
 193     // Initialize bitmap
 194     RtlInitializeBitMap(&Bitmap, BitmapPtr, NextNodeNumber);
 195
 196     // Do we need to enlarge the bitmap?
 197     if (BytesNeeded > BitmapLength)
 198     {
 199         // TODO: handle synchronization issues that could occur from changing the directory's file record
 200         // Change bitmap size
 201         DataSize.QuadPart = BytesNeeded;
 202         if (BitmapCtx->pRecord->IsNonResident)
 203         {
 204             Status = SetNonResidentAttributeDataLength(DeviceExt,
 205                                                        BitmapCtx,
 206                                                        BitmapOffset,
 207                                                        FileRecord,
 208                                                        &DataSize);
 209         }
 210         else
 211         {
 212             Status = SetResidentAttributeDataLength(DeviceExt,
 213                                                     BitmapCtx,
 214                                                     BitmapOffset,
 215                                                     FileRecord,
 216                                                     &DataSize);
 217         }
 218         if (!NT_SUCCESS(Status))
 219         {
 220             DPRINT1("ERROR: Failed to set length of bitmap attribute!\n");
 221             ReleaseAttributeContext(BitmapCtx);
 222             return Status;
 223         }
 224     }
 225
 226     // Enlarge Index Allocation attribute
 227     DataSize.QuadPart = IndexAllocationLength + IndexBufferSize;
 228     Status = SetNonResidentAttributeDataLength(DeviceExt,
 229                                                IndexAllocationCtx,
 230                                                IndexAllocationOffset,
 231                                                FileRecord,
 232                                                &DataSize);
 233     if (!NT_SUCCESS(Status))
 234     {
 235         DPRINT1("ERROR: Failed to set length of index allocation!\n");
 236         ReleaseAttributeContext(BitmapCtx);
 237         return Status;
 238     }
 239
 240     // Update file record on disk
 241     Status = UpdateFileRecord(DeviceExt, IndexAllocationCtx->FileMFTIndex, FileRecord);
 242     if (!NT_SUCCESS(Status))
 243     {
 244         DPRINT1("ERROR: Failed to update file record!\n");
 245         ReleaseAttributeContext(BitmapCtx);
 246         return Status;
 247     }
 248
 249     // Set the bit for the new index record
 250     RtlSetBits(&Bitmap, NextNodeNumber, 1);
 251
 252     // Write the new bitmap attribute
 253     Status = WriteAttribute(DeviceExt,
 254                             BitmapCtx,
 255                             0,
 256                             (const PUCHAR)BitmapPtr,
 257                             BytesNeeded,
 258                             &BytesWritten,
 259                             FileRecord);
 260     if (!NT_SUCCESS(Status))
 261     {
 262         DPRINT1("ERROR: Unable to write to $I30 bitmap attribute!\n");
 263     }
 264
 265     // Calculate VCN of new node number
 266     *NewVCN = NextNodeNumber * (IndexBufferSize / DeviceExt->NtfsInfo.BytesPerCluster);
 267
 268     DPRINT("New VCN: %I64u\n", *NewVCN);
 269
 270     ExFreePoolWithTag(BitmapMem, TAG_NTFS);
 271     ReleaseAttributeContext(BitmapCtx);
 272
 273     return Status;
 274 }
 275
 276 /**
 277 * @name CreateDummyKey
 278 * @implemented
 279 *
 280 * Creates the final B_TREE_KEY for a B_TREE_FILENAME_NODE. Also creates the associated index entry.
 281 *
 282 * @param HasChildNode
 283 * BOOLEAN to indicate if this key will have a LesserChild.
 284 *
 285 * @return
 286 * The newly-created key.
 287 */
 288 PB_TREE_KEY
 289 CreateDummyKey(BOOLEAN HasChildNode)
 290 {
 291     PINDEX_ENTRY_ATTRIBUTE NewIndexEntry;
 292     PB_TREE_KEY NewDummyKey;
 293
 294     // Calculate max size of a dummy key
 295     ULONG EntrySize = ALIGN_UP_BY(FIELD_OFFSET(INDEX_ENTRY_ATTRIBUTE, FileName), 8);
 296     EntrySize += sizeof(ULONGLONG); // for VCN
 297
 298     // Create the index entry for the key
 299     NewIndexEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_NTFS);
 300     if (!NewIndexEntry)
 301     {
 302         DPRINT1("Couldn't allocate memory for dummy key index entry!\n");
 303         return NULL;
 304     }
 305
 306     RtlZeroMemory(NewIndexEntry, EntrySize);
 307
 308     if (HasChildNode)
 309     {
 310         NewIndexEntry->Flags = NTFS_INDEX_ENTRY_NODE | NTFS_INDEX_ENTRY_END;
 311     }
 312     else
 313     {
 314         NewIndexEntry->Flags = NTFS_INDEX_ENTRY_END;
 315         EntrySize -= sizeof(ULONGLONG); // no VCN
 316     }
 317
 318     NewIndexEntry->Length = EntrySize;
 319
 320     // Create the key
 321     NewDummyKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 322     if (!NewDummyKey)
 323     {
 324         DPRINT1("Unable to allocate dummy key!\n");
 325         ExFreePoolWithTag(NewIndexEntry, TAG_NTFS);
 326         return NULL;
 327     }
 328     RtlZeroMemory(NewDummyKey, sizeof(B_TREE_KEY));
 329
 330     NewDummyKey->IndexEntry = NewIndexEntry;
 331
 332     return NewDummyKey;
 333 }
 334
 335 /**
 336 * @name CreateEmptyBTree
 337 * @implemented
 338 *
 339 * Creates an empty B-Tree, which will contain a single root node which will contain a single dummy key.
 340 *
 341 * @param NewTree
 342 * Pointer to a PB_TREE that will receive the pointer of the newly-created B-Tree.
 343 *
 344 * @return
 345 * STATUS_SUCCESS on success. STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
 346 */
 347 NTSTATUS
 348 CreateEmptyBTree(PB_TREE *NewTree)
 349 {
 350     PB_TREE Tree = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE), TAG_NTFS);
 351     PB_TREE_FILENAME_NODE RootNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
 352     PB_TREE_KEY DummyKey;
 353
 354     DPRINT1("CreateEmptyBTree(%p) called\n", NewTree);
 355
 356     if (!Tree || !RootNode)
 357     {
 358         DPRINT1("Couldn't allocate enough memory for B-Tree!\n");
 359         if (Tree)
 360             ExFreePoolWithTag(Tree, TAG_NTFS);
 361         if (RootNode)
 362             ExFreePoolWithTag(RootNode, TAG_NTFS);
 363         return STATUS_INSUFFICIENT_RESOURCES;
 364     }
 365
 366     // Create the dummy key
 367     DummyKey = CreateDummyKey(FALSE);
 368     if (!DummyKey)
 369     {
 370         DPRINT1("ERROR: Failed to create dummy key!\n");
 371         ExFreePoolWithTag(Tree, TAG_NTFS);
 372         ExFreePoolWithTag(RootNode, TAG_NTFS);
 373         return STATUS_INSUFFICIENT_RESOURCES;
 374     }
 375
 376     RtlZeroMemory(Tree, sizeof(B_TREE));
 377     RtlZeroMemory(RootNode, sizeof(B_TREE_FILENAME_NODE));
 378
 379     // Setup the Tree
 380     RootNode->FirstKey = DummyKey;
 381     RootNode->KeyCount = 1;
 382     RootNode->DiskNeedsUpdating = TRUE;
 383     Tree->RootNode = RootNode;
 384
 385     *NewTree = Tree;
 386
 387     // Memory will be freed when DestroyBTree() is called
 388
 389     return STATUS_SUCCESS;
 390 }
 391
 392 /**
 393 * @name CompareTreeKeys
 394 * @implemented
 395 *
 396 * Compare two B_TREE_KEY's to determine their order in the tree.
 397 *
 398 * @param Key1
 399 * Pointer to a B_TREE_KEY that will be compared.
 400 *
 401 * @param Key2
 402 * Pointer to the other B_TREE_KEY that will be compared.
 403 *
 404 * @param CaseSensitive
 405 * Boolean indicating if the function should operate in case-sensitive mode. This will be TRUE
 406 * if an application created the file with the FILE_FLAG_POSIX_SEMANTICS flag.
 407 *
 408 * @returns
 409 * 0 if the two keys are equal.
 410 * < 0 if key1 is less thank key2
 411 * > 0 if key1 is greater than key2
 412 *
 413 * @remarks
 414 * Any other key is always less than the final (dummy) key in a node. Key1 must not be the dummy node.
 415 */
 416 LONG
 417 CompareTreeKeys(PB_TREE_KEY Key1, PB_TREE_KEY Key2, BOOLEAN CaseSensitive)
 418 {
 419     UNICODE_STRING Key1Name, Key2Name;
 420     LONG Comparison;
 421
 422     // Key1 must not be the final key (AKA the dummy key)
 423     ASSERT(!(Key1->IndexEntry->Flags & NTFS_INDEX_ENTRY_END));
 424
 425     // If Key2 is the "dummy key", key 1 will always come first
 426     if (Key2->NextKey == NULL)
 427         return -1;
 428
 429     Key1Name.Buffer = Key1->IndexEntry->FileName.Name;
 430     Key1Name.Length = Key1Name.MaximumLength
 431         = Key1->IndexEntry->FileName.NameLength * sizeof(WCHAR);
 432
 433     Key2Name.Buffer = Key2->IndexEntry->FileName.Name;
 434     Key2Name.Length = Key2Name.MaximumLength
 435         = Key2->IndexEntry->FileName.NameLength * sizeof(WCHAR);
 436
 437     // Are the two keys the same length?
 438     if (Key1Name.Length == Key2Name.Length)
 439         return RtlCompareUnicodeString(&Key1Name, &Key2Name, !CaseSensitive);
 440
 441     // Is Key1 shorter?
 442     if (Key1Name.Length < Key2Name.Length)
 443     {
 444         // Truncate KeyName2 to be the same length as KeyName1
 445         Key2Name.Length = Key1Name.Length;
 446
 447         // Compare the names of the same length
 448         Comparison = RtlCompareUnicodeString(&Key1Name, &Key2Name, !CaseSensitive);
 449
 450         // If the truncated names are the same length, the shorter one comes first
 451         if (Comparison == 0)
 452             return -1;
 453     }
 454     else
 455     {
 456         // Key2 is shorter
 457         // Truncate KeyName1 to be the same length as KeyName2
 458         Key1Name.Length = Key2Name.Length;
 459
 460         // Compare the names of the same length
 461         Comparison = RtlCompareUnicodeString(&Key1Name, &Key2Name, !CaseSensitive);
 462
 463         // If the truncated names are the same length, the shorter one comes first
 464         if (Comparison == 0)
 465             return 1;
 466     }
 467
 468     return Comparison;
 469 }
 470
 471 /**
 472 * @name CountBTreeKeys
 473 * @implemented
 474 *
 475 * Counts the number of linked B-Tree keys, starting with FirstKey.
 476 *
 477 * @param FirstKey
 478 * Pointer to a B_TREE_KEY that will be the first key to be counted.
 479 *
 480 * @return
 481 * The number of keys in a linked-list, including FirstKey and the final dummy key.
 482 */
 483 ULONG
 484 CountBTreeKeys(PB_TREE_KEY FirstKey)
 485 {
 486     ULONG Count = 0;
 487     PB_TREE_KEY Current = FirstKey;
 488
 489     while (Current != NULL)
 490     {
 491         Count++;
 492         Current = Current->NextKey;
 493     }
 494
 495     return Count;
 496 }
 497
 498 PB_TREE_FILENAME_NODE
 499 CreateBTreeNodeFromIndexNode(PDEVICE_EXTENSION Vcb,
 500                              PINDEX_ROOT_ATTRIBUTE IndexRoot,
 501                              PNTFS_ATTR_CONTEXT IndexAllocationAttributeCtx,
 502                              PINDEX_ENTRY_ATTRIBUTE NodeEntry)
 503 {
 504     PB_TREE_FILENAME_NODE NewNode;
 505     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 506     PINDEX_ENTRY_ATTRIBUTE FirstNodeEntry;
 507     ULONG CurrentEntryOffset = 0;
 508     PINDEX_BUFFER NodeBuffer;
 509     ULONG IndexBufferSize = Vcb->NtfsInfo.BytesPerIndexRecord;
 510     PULONGLONG VCN;
 511     PB_TREE_KEY CurrentKey;
 512     NTSTATUS Status;
 513     ULONGLONG IndexNodeOffset;
 514     ULONG BytesRead;
 515
 516     if (IndexAllocationAttributeCtx == NULL)
 517     {
 518         DPRINT1("ERROR: Couldn't find index allocation attribute even though there should be one!\n");
 519         return NULL;
 520     }
 521
 522     // Get the node number from the end of the node entry
 523     VCN = (PULONGLONG)((ULONG_PTR)NodeEntry + NodeEntry->Length - sizeof(ULONGLONG));
 524
 525     // Create the new tree node
 526     NewNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
 527     if (!NewNode)
 528     {
 529         DPRINT1("ERROR: Couldn't allocate memory for new filename node.\n");
 530         return NULL;
 531     }
 532     RtlZeroMemory(NewNode, sizeof(B_TREE_FILENAME_NODE));
 533
 534     // Create the first key
 535     CurrentKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 536     if (!CurrentKey)
 537     {
 538         DPRINT1("ERROR: Failed to allocate memory for key!\n");
 539         ExFreePoolWithTag(NewNode, TAG_NTFS);
 540         return NULL;
 541     }
 542     RtlZeroMemory(CurrentKey, sizeof(B_TREE_KEY));
 543     NewNode->FirstKey = CurrentKey;
 544
 545     // Allocate memory for the node buffer
 546     NodeBuffer = ExAllocatePoolWithTag(NonPagedPool, IndexBufferSize, TAG_NTFS);
 547     if (!NodeBuffer)
 548     {
 549         DPRINT1("ERROR: Couldn't allocate memory for node buffer!\n");
 550         ExFreePoolWithTag(CurrentKey, TAG_NTFS);
 551         ExFreePoolWithTag(NewNode, TAG_NTFS);
 552         return NULL;
 553     }
 554
 555     // Calculate offset into index allocation
 556     IndexNodeOffset = GetAllocationOffsetFromVCN(Vcb, IndexBufferSize, *VCN);
 557
 558     // TODO: Confirm index bitmap has this node marked as in-use
 559
 560     // Read the node
 561     BytesRead = ReadAttribute(Vcb,
 562                               IndexAllocationAttributeCtx,
 563                               IndexNodeOffset,
 564                               (PCHAR)NodeBuffer,
 565                               IndexBufferSize);
 566
 567     ASSERT(BytesRead == IndexBufferSize);
 568     NT_ASSERT(NodeBuffer->Ntfs.Type == NRH_INDX_TYPE);
 569     NT_ASSERT(NodeBuffer->VCN == *VCN);
 570
 571     // Apply the fixup array to the node buffer
 572     Status = FixupUpdateSequenceArray(Vcb, &NodeBuffer->Ntfs);
 573     if (!NT_SUCCESS(Status))
 574     {
 575         DPRINT1("ERROR: Couldn't apply fixup array to index node buffer!\n");
 576         ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 577         ExFreePoolWithTag(CurrentKey, TAG_NTFS);
 578         ExFreePoolWithTag(NewNode, TAG_NTFS);
 579         return NULL;
 580     }
 581
 582     // Walk through the index and create keys for all the entries
 583     FirstNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)(&NodeBuffer->Header)
 584                                                + NodeBuffer->Header.FirstEntryOffset);
 585     CurrentNodeEntry = FirstNodeEntry;
 586     while (CurrentEntryOffset < NodeBuffer->Header.TotalSizeOfEntries)
 587     {
 588         // Allocate memory for the current entry
 589         CurrentKey->IndexEntry = ExAllocatePoolWithTag(NonPagedPool, CurrentNodeEntry->Length, TAG_NTFS);
 590         if (!CurrentKey->IndexEntry)
 591         {
 592             DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 593             DestroyBTreeNode(NewNode);
 594             ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 595             return NULL;
 596         }
 597
 598         NewNode->KeyCount++;
 599
 600         // If this isn't the last entry
 601         if (!(CurrentNodeEntry->Flags & NTFS_INDEX_ENTRY_END))
 602         {
 603             // Create the next key
 604             PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 605             if (!NextKey)
 606             {
 607                 DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 608                 DestroyBTreeNode(NewNode);
 609                 ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 610                 return NULL;
 611             }
 612             RtlZeroMemory(NextKey, sizeof(B_TREE_KEY));
 613
 614             // Add NextKey to the end of the list
 615             CurrentKey->NextKey = NextKey;
 616
 617             // Copy the current entry to its key
 618             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 619
 620             // See if the current key has a sub-node
 621             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 622             {
 623                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 624                                                                        IndexRoot,
 625                                                                        IndexAllocationAttributeCtx,
 626                                                                        CurrentKey->IndexEntry);
 627             }
 628
 629             CurrentKey = NextKey;
 630         }
 631         else
 632         {
 633             // Copy the final entry to its key
 634             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 635             CurrentKey->NextKey = NULL;
 636
 637             // See if the current key has a sub-node
 638             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 639             {
 640                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 641                                                                        IndexRoot,
 642                                                                        IndexAllocationAttributeCtx,
 643                                                                        CurrentKey->IndexEntry);
 644             }
 645
 646             break;
 647         }
 648
 649         // Advance to the next entry
 650         CurrentEntryOffset += CurrentNodeEntry->Length;
 651         CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 652     }
 653
 654     NewNode->VCN = *VCN;
 655     NewNode->HasValidVCN = TRUE;
 656
 657     ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 658
 659     return NewNode;
 660 }
 661
 662 /**
 663 * @name CreateBTreeFromIndex
 664 * @implemented
 665 *
 666 * Parse an index and create a B-Tree in memory from it.
 667 *
 668 * @param IndexRootContext
 669 * Pointer to an NTFS_ATTR_CONTEXT that describes the location of the index root attribute.
 670 *
 671 * @param NewTree
 672 * Pointer to a PB_TREE that will receive the pointer to a newly-created B-Tree.
 673 *
 674 * @returns
 675 * STATUS_SUCCESS on success.
 676 * STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
 677 *
 678 * @remarks
 679 * Allocates memory for the entire tree. Caller is responsible for destroying the tree with DestroyBTree().
 680 */
 681 NTSTATUS
 682 CreateBTreeFromIndex(PDEVICE_EXTENSION Vcb,
 683                      PFILE_RECORD_HEADER FileRecordWithIndex,
 684                      /*PCWSTR IndexName,*/
 685                      PNTFS_ATTR_CONTEXT IndexRootContext,
 686                      PINDEX_ROOT_ATTRIBUTE IndexRoot,
 687                      PB_TREE *NewTree)
 688 {
 689     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 690     PB_TREE Tree = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE), TAG_NTFS);
 691     PB_TREE_FILENAME_NODE RootNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
 692     PB_TREE_KEY CurrentKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 693     ULONG CurrentOffset = IndexRoot->Header.FirstEntryOffset;
 694     PNTFS_ATTR_CONTEXT IndexAllocationContext = NULL;
 695     NTSTATUS Status;
 696
 697     DPRINT1("CreateBTreeFromIndex(%p, %p)\n", IndexRoot, NewTree);
 698
 699     if (!Tree || !RootNode || !CurrentKey)
 700     {
 701         DPRINT1("Couldn't allocate enough memory for B-Tree!\n");
 702         if (Tree)
 703             ExFreePoolWithTag(Tree, TAG_NTFS);
 704         if (CurrentKey)
 705             ExFreePoolWithTag(CurrentKey, TAG_NTFS);
 706         if (RootNode)
 707             ExFreePoolWithTag(RootNode, TAG_NTFS);
 708         return STATUS_INSUFFICIENT_RESOURCES;
 709     }
 710
 711     RtlZeroMemory(Tree, sizeof(B_TREE));
 712     RtlZeroMemory(RootNode, sizeof(B_TREE_FILENAME_NODE));
 713     RtlZeroMemory(CurrentKey, sizeof(B_TREE_KEY));
 714
 715     // See if the file record has an attribute allocation
 716     Status = FindAttribute(Vcb,
 717                            FileRecordWithIndex,
 718                            AttributeIndexAllocation,
 719                            L"$I30",
 720                            4,
 721                            &IndexAllocationContext,
 722                            NULL);
 723     if (!NT_SUCCESS(Status))
 724         IndexAllocationContext = NULL;
 725
 726     // Setup the Tree
 727     RootNode->FirstKey = CurrentKey;
 728     Tree->RootNode = RootNode;
 729
 730     // Make sure we won't try reading past the attribute-end
 731     if (FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header) + IndexRoot->Header.TotalSizeOfEntries > IndexRootContext->pRecord->Resident.ValueLength)
 732     {
 733         DPRINT1("Filesystem corruption detected!\n");
 734         DestroyBTree(Tree);
 735         return STATUS_FILE_CORRUPT_ERROR;
 736     }
 737
 738     // Start at the first node entry
 739     CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)IndexRoot
 740                                                 + FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header)
 741                                                 + IndexRoot->Header.FirstEntryOffset);
 742
 743     // Create a key for each entry in the node
 744     while (CurrentOffset < IndexRoot->Header.TotalSizeOfEntries)
 745     {
 746         // Allocate memory for the current entry
 747         CurrentKey->IndexEntry = ExAllocatePoolWithTag(NonPagedPool, CurrentNodeEntry->Length, TAG_NTFS);
 748         if (!CurrentKey->IndexEntry)
 749         {
 750             DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 751             DestroyBTree(Tree);
 752             return STATUS_INSUFFICIENT_RESOURCES;
 753         }
 754
 755         RootNode->KeyCount++;
 756
 757         // If this isn't the last entry
 758         if (!(CurrentNodeEntry->Flags & NTFS_INDEX_ENTRY_END))
 759         {
 760             // Create the next key
 761             PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 762             if (!NextKey)
 763             {
 764                 DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 765                 DestroyBTree(Tree);
 766                 return STATUS_INSUFFICIENT_RESOURCES;
 767             }
 768
 769             RtlZeroMemory(NextKey, sizeof(B_TREE_KEY));
 770
 771             // Add NextKey to the end of the list
 772             CurrentKey->NextKey = NextKey;
 773
 774             // Copy the current entry to its key
 775             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 776
 777             // Does this key have a sub-node?
 778             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 779             {
 780                 // Create the child node
 781                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 782                                                                        IndexRoot,
 783                                                                        IndexAllocationContext,
 784                                                                        CurrentKey->IndexEntry);
 785                 if (!CurrentKey->LesserChild)
 786                 {
 787                     DPRINT1("ERROR: Couldn't create child node!\n");
 788                     DestroyBTree(Tree);
 789                     return STATUS_NOT_IMPLEMENTED;
 790                 }
 791             }
 792
 793             // Advance to the next entry
 794             CurrentOffset += CurrentNodeEntry->Length;
 795             CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 796             CurrentKey = NextKey;
 797         }
 798         else
 799         {
 800             // Copy the final entry to its key
 801             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 802             CurrentKey->NextKey = NULL;
 803
 804             // Does this key have a sub-node?
 805             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 806             {
 807                 // Create the child node
 808                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 809                                                                        IndexRoot,
 810                                                                        IndexAllocationContext,
 811                                                                        CurrentKey->IndexEntry);
 812                 if (!CurrentKey->LesserChild)
 813                 {
 814                     DPRINT1("ERROR: Couldn't create child node!\n");
 815                     DestroyBTree(Tree);
 816                     return STATUS_NOT_IMPLEMENTED;
 817                 }
 818             }
 819
 820             break;
 821         }
 822     }
 823
 824     *NewTree = Tree;
 825
 826     if (IndexAllocationContext)
 827         ReleaseAttributeContext(IndexAllocationContext);
 828
 829     return STATUS_SUCCESS;
 830 }
 831
 832 /**
 833 * @name GetSizeOfIndexEntries
 834 * @implemented
 835 *
 836 * Sums the size of each index entry in every key in a B-Tree node.
 837 *
 838 * @param Node
 839 * Pointer to a B_TREE_FILENAME_NODE. The size of this node's index entries will be returned.
 840 *
 841 * @returns
 842 * The sum of the sizes of every index entry for each key in the B-Tree node.
 843 *
 844 * @remarks
 845 * Gets only the size of the index entries; doesn't include the size of any headers that would be added to an index record.
 846 */
 847 ULONG
 848 GetSizeOfIndexEntries(PB_TREE_FILENAME_NODE Node)
 849 {
 850     // Start summing the total size of this node's entries
 851     ULONG NodeSize = 0;
 852
 853     // Walk through the list of Node Entries
 854     PB_TREE_KEY CurrentKey = Node->FirstKey;
 855     ULONG i;
 856     for (i = 0; i < Node->KeyCount; i++)
 857     {
 858         ASSERT(CurrentKey->IndexEntry->Length != 0);
 859
 860         // Add the length of the current node
 861         NodeSize += CurrentKey->IndexEntry->Length;
 862         CurrentKey = CurrentKey->NextKey;
 863     }
 864
 865     return NodeSize;
 866 }
 867
 868 /**
 869 * @name CreateIndexRootFromBTree
 870 * @implemented
 871 *
 872 * Parse a B-Tree in memory and convert it into an index that can be written to disk.
 873 *
 874 * @param DeviceExt
 875 * Pointer to the DEVICE_EXTENSION of the target drive.
 876 *
 877 * @param Tree
 878 * Pointer to a B_TREE that describes the index to be written.
 879 *
 880 * @param MaxIndexSize
 881 * Describes how large the index can be before it will take too much space in the file record.
 882 * This is strictly the sum of the sizes of all index entries; it does not include the space
 883 * required by the index root header (INDEX_ROOT_ATTRIBUTE), since that size will be constant.
 884 *
 885 * After reaching MaxIndexSize, an index can no longer be represented with just an index root
 886 * attribute, and will require an index allocation and $I30 bitmap (TODO).
 887 *
 888 * @param IndexRoot
 889 * Pointer to a PINDEX_ROOT_ATTRIBUTE that will receive a pointer to the newly-created index.
 890 *
 891 * @param Length
 892 * Pointer to a ULONG which will receive the length of the new index root.
 893 *
 894 * @returns
 895 * STATUS_SUCCESS on success.
 896 * STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
 897 * STATUS_NOT_IMPLEMENTED if the new index can't fit within MaxIndexSize.
 898 *
 899 * @remarks
 900 * If the function succeeds, it's the caller's responsibility to free IndexRoot with ExFreePoolWithTag().
 901 */
 902 NTSTATUS
 903 CreateIndexRootFromBTree(PDEVICE_EXTENSION DeviceExt,
 904                          PB_TREE Tree,
 905                          ULONG MaxIndexSize,
 906                          PINDEX_ROOT_ATTRIBUTE *IndexRoot,
 907                          ULONG *Length)
 908 {
 909     ULONG i;
 910     PB_TREE_KEY CurrentKey;
 911     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 912     PINDEX_ROOT_ATTRIBUTE NewIndexRoot = ExAllocatePoolWithTag(NonPagedPool,
 913                                                                DeviceExt->NtfsInfo.BytesPerFileRecord,
 914                                                                TAG_NTFS);
 915
 916     DPRINT1("CreateIndexRootFromBTree(%p, %p, 0x%lx, %p, %p)\n", DeviceExt, Tree, MaxIndexSize, IndexRoot, Length);
 917
 918 #ifndef NDEBUG
 919     DumpBTree(Tree);
 920 #endif
 921
 922     if (!NewIndexRoot)
 923     {
 924         DPRINT1("Failed to allocate memory for Index Root!\n");
 925         return STATUS_INSUFFICIENT_RESOURCES;
 926     }
 927
 928     // Setup the new index root
 929     RtlZeroMemory(NewIndexRoot, DeviceExt->NtfsInfo.BytesPerFileRecord);
 930
 931     NewIndexRoot->AttributeType = AttributeFileName;
 932     NewIndexRoot->CollationRule = COLLATION_FILE_NAME;
 933     NewIndexRoot->SizeOfEntry = DeviceExt->NtfsInfo.BytesPerIndexRecord;
 934     // If Bytes per index record is less than cluster size, clusters per index record becomes sectors per index
 935     if (NewIndexRoot->SizeOfEntry < DeviceExt->NtfsInfo.BytesPerCluster)
 936         NewIndexRoot->ClustersPerIndexRecord = NewIndexRoot->SizeOfEntry / DeviceExt->NtfsInfo.BytesPerSector;
 937     else
 938         NewIndexRoot->ClustersPerIndexRecord = NewIndexRoot->SizeOfEntry / DeviceExt->NtfsInfo.BytesPerCluster;
 939
 940     // Setup the Index node header
 941     NewIndexRoot->Header.FirstEntryOffset = sizeof(INDEX_HEADER_ATTRIBUTE);
 942     NewIndexRoot->Header.Flags = INDEX_ROOT_SMALL;
 943
 944     // Start summing the total size of this node's entries
 945     NewIndexRoot->Header.TotalSizeOfEntries = NewIndexRoot->Header.FirstEntryOffset;
 946
 947     // Setup each Node Entry
 948     CurrentKey = Tree->RootNode->FirstKey;
 949     CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)NewIndexRoot
 950                                                 + FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header)
 951                                                 + NewIndexRoot->Header.FirstEntryOffset);
 952     for (i = 0; i < Tree->RootNode->KeyCount; i++)
 953     {
 954         // Would adding the current entry to the index increase the index size beyond the limit we've set?
 955         ULONG IndexSize = NewIndexRoot->Header.TotalSizeOfEntries - NewIndexRoot->Header.FirstEntryOffset + CurrentKey->IndexEntry->Length;
 956         if (IndexSize > MaxIndexSize)
 957         {
 958             DPRINT1("TODO: Adding file would require creating an attribute list!\n");
 959             ExFreePoolWithTag(NewIndexRoot, TAG_NTFS);
 960             return STATUS_NOT_IMPLEMENTED;
 961         }
 962
 963         ASSERT(CurrentKey->IndexEntry->Length != 0);
 964
 965         // Copy the index entry
 966         RtlCopyMemory(CurrentNodeEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
 967
 968         DPRINT1("Index Node Entry Stream Length: %u\nIndex Node Entry Length: %u\n",
 969                 CurrentNodeEntry->KeyLength,
 970                 CurrentNodeEntry->Length);
 971
 972         // Does the current key have any sub-nodes?
 973         if (CurrentKey->LesserChild)
 974             NewIndexRoot->Header.Flags = INDEX_ROOT_LARGE;
 975
 976         // Add Length of Current Entry to Total Size of Entries
 977         NewIndexRoot->Header.TotalSizeOfEntries += CurrentKey->IndexEntry->Length;
 978
 979         // Go to the next node entry
 980         CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 981
 982         CurrentKey = CurrentKey->NextKey;
 983     }
 984
 985     NewIndexRoot->Header.AllocatedSize = NewIndexRoot->Header.TotalSizeOfEntries;
 986
 987     *IndexRoot = NewIndexRoot;
 988     *Length = NewIndexRoot->Header.AllocatedSize + FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header);
 989
 990     return STATUS_SUCCESS;
 991 }
 992
 993 NTSTATUS
 994 CreateIndexBufferFromBTreeNode(PDEVICE_EXTENSION DeviceExt,
 995                                PB_TREE_FILENAME_NODE Node,
 996                                ULONG BufferSize,
 997                                BOOLEAN HasChildren,
 998                                PINDEX_BUFFER IndexBuffer)
 999 {
1000     ULONG i;
1001     PB_TREE_KEY CurrentKey;
1002     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
1003     NTSTATUS Status;
1004
1005     // TODO: Fix magic, do math
1006     RtlZeroMemory(IndexBuffer, BufferSize);
1007     IndexBuffer->Ntfs.Type = NRH_INDX_TYPE;
1008     IndexBuffer->Ntfs.UsaOffset = 0x28;
1009     IndexBuffer->Ntfs.UsaCount = 9;
1010
1011     // TODO: Check bitmap for VCN
1012     ASSERT(Node->HasValidVCN);
1013     IndexBuffer->VCN = Node->VCN;
1014
1015     // Windows seems to alternate between using 0x28 and 0x40 for the first entry offset of each index buffer.
1016     // Interestingly, neither Windows nor chkdsk seem to mind if we just use 0x28 for every index record.
1017     IndexBuffer->Header.FirstEntryOffset = 0x28;
1018     IndexBuffer->Header.AllocatedSize = BufferSize - FIELD_OFFSET(INDEX_BUFFER, Header);
1019
1020     // Start summing the total size of this node's entries
1021     IndexBuffer->Header.TotalSizeOfEntries = IndexBuffer->Header.FirstEntryOffset;
1022
1023     CurrentKey = Node->FirstKey;
1024     CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)&(IndexBuffer->Header)
1025                                                 + IndexBuffer->Header.FirstEntryOffset);
1026     for (i = 0; i < Node->KeyCount; i++)
1027     {
1028         // Would adding the current entry to the index increase the node size beyond the allocation size?
1029         ULONG IndexSize = FIELD_OFFSET(INDEX_BUFFER, Header)
1030             + IndexBuffer->Header.TotalSizeOfEntries
1031             + CurrentNodeEntry->Length;
1032         if (IndexSize > BufferSize)
1033         {
1034             DPRINT1("TODO: Adding file would require creating a new node!\n");
1035             return STATUS_NOT_IMPLEMENTED;
1036         }
1037
1038         ASSERT(CurrentKey->IndexEntry->Length != 0);
1039
1040         // Copy the index entry
1041         RtlCopyMemory(CurrentNodeEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
1042
1043         DPRINT("Index Node Entry Stream Length: %u\nIndex Node Entry Length: %u\n",
1044                CurrentNodeEntry->KeyLength,
1045                CurrentNodeEntry->Length);
1046
1047         // Add Length of Current Entry to Total Size of Entries
1048         IndexBuffer->Header.TotalSizeOfEntries += CurrentNodeEntry->Length;
1049
1050         // Check for child nodes
1051         if (HasChildren)
1052             IndexBuffer->Header.Flags = INDEX_NODE_LARGE;
1053
1054         // Go to the next node entry
1055         CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
1056         CurrentKey = CurrentKey->NextKey;
1057     }
1058
1059     Status = AddFixupArray(DeviceExt, &IndexBuffer->Ntfs);
1060
1061     return Status;
1062 }
1063
1064 /**
1065 * @name DemoteBTreeRoot
1066 * @implemented
1067 *
1068 * Demoting the root means first putting all the keys in the root node into a new node, and making
1069 * the new node a child of a dummy key. The dummy key then becomes the sole contents of the root node.
1070 * The B-Tree gets one level deeper. This operation is needed when an index root grows too large for its file record.
1071 * Demotion is my own term; I might change the name later if I think of something more descriptive or can find
1072 * an appropriate name for this operation in existing B-Tree literature.
1073 *
1074 * @param Tree
1075 * Pointer to the B_TREE whose root is being demoted
1076 *
1077 * @returns
1078 * STATUS_SUCCESS on success.
1079 * STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
1080 */
1081 NTSTATUS
1082 DemoteBTreeRoot(PB_TREE Tree)
1083 {
1084     PB_TREE_FILENAME_NODE NewSubNode, NewIndexRoot;
1085     PB_TREE_KEY DummyKey;
1086
1087     DPRINT1("Collapsing Index Root into sub-node.\n");
1088
1089 #ifndef NDEBUG
1090     DumpBTree(Tree);
1091 #endif
1092
1093     // Create a new node that will hold the keys currently in index root
1094     NewSubNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
1095     if (!NewSubNode)
1096     {
1097         DPRINT1("ERROR: Couldn't allocate memory for new sub-node.\n");
1098         return STATUS_INSUFFICIENT_RESOURCES;
1099     }
1100     RtlZeroMemory(NewSubNode, sizeof(B_TREE_FILENAME_NODE));
1101
1102     // Copy the applicable data from the old index root node
1103     NewSubNode->KeyCount = Tree->RootNode->KeyCount;
1104     NewSubNode->FirstKey = Tree->RootNode->FirstKey;
1105     NewSubNode->DiskNeedsUpdating = TRUE;
1106
1107     // Create a new dummy key, and make the new node it's child
1108     DummyKey = CreateDummyKey(TRUE);
1109     if (!DummyKey)
1110     {
1111         DPRINT1("ERROR: Couldn't allocate memory for new root node.\n");
1112         ExFreePoolWithTag(NewSubNode, TAG_NTFS);
1113         return STATUS_INSUFFICIENT_RESOURCES;
1114     }
1115
1116     // Make the new node a child of the dummy key
1117     DummyKey->LesserChild = NewSubNode;
1118
1119     // Create a new index root node
1120     NewIndexRoot = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
1121     if (!NewIndexRoot)
1122     {
1123         DPRINT1("ERROR: Couldn't allocate memory for new index root.\n");
1124         ExFreePoolWithTag(NewSubNode, TAG_NTFS);
1125         ExFreePoolWithTag(DummyKey, TAG_NTFS);
1126         return STATUS_INSUFFICIENT_RESOURCES;
1127     }
1128     RtlZeroMemory(NewIndexRoot, sizeof(B_TREE_FILENAME_NODE));
1129
1130     NewIndexRoot->DiskNeedsUpdating = TRUE;
1131
1132     // Insert the dummy key into the new node
1133     NewIndexRoot->FirstKey = DummyKey;
1134     NewIndexRoot->KeyCount = 1;
1135     NewIndexRoot->DiskNeedsUpdating = TRUE;
1136
1137     // Make the new node the Tree's root node
1138     Tree->RootNode = NewIndexRoot;
1139
1140 #ifndef NDEBUG
1141     DumpBTree(Tree);
1142 #endif
1143
1144     return STATUS_SUCCESS;
1145 }
1146
1147 /**
1148 * @name SetIndexEntryVCN
1149 * @implemented
1150 *
1151 * Sets the VCN of a given IndexEntry.
1152 *
1153 * @param IndexEntry
1154 * Pointer to an INDEX_ENTRY_ATTRIBUTE structure that will have its VCN set.
1155 *
1156 * @param VCN
1157 * VCN to store in the index entry.
1158 *
1159 * @remarks
1160 * The index entry must have enough memory allocated to store the VCN, and must have the NTFS_INDEX_ENTRY_NODE flag set.
1161 * The VCN of an index entry is stored at the very end of the structure, after the filename attribute. Since the filename
1162 * attribute can be a variable size, this function makes setting this member easy.
1163 */
1164 VOID
1165 SetIndexEntryVCN(PINDEX_ENTRY_ATTRIBUTE IndexEntry, ULONGLONG VCN)
1166 {
1167     PULONGLONG Destination = (PULONGLONG)((ULONG_PTR)IndexEntry + IndexEntry->Length - sizeof(ULONGLONG));
1168
1169     ASSERT(IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE);
1170
1171     *Destination = VCN;
1172 }
1173
1174 NTSTATUS
1175 UpdateIndexAllocation(PDEVICE_EXTENSION DeviceExt,
1176                       PB_TREE Tree,
1177                       ULONG IndexBufferSize,
1178                       PFILE_RECORD_HEADER FileRecord)
1179 {
1180     // Find the index allocation and bitmap
1181     PNTFS_ATTR_CONTEXT IndexAllocationContext;
1182     PB_TREE_KEY CurrentKey;
1183     NTSTATUS Status;
1184     BOOLEAN HasIndexAllocation = FALSE;
1185     ULONG i;
1186     ULONG IndexAllocationOffset;
1187
1188     DPRINT1("UpdateIndexAllocation() called.\n");
1189
1190     Status = FindAttribute(DeviceExt, FileRecord, AttributeIndexAllocation, L"$I30", 4, &IndexAllocationContext, &IndexAllocationOffset);
1191     if (NT_SUCCESS(Status))
1192     {
1193         HasIndexAllocation = TRUE;
1194
1195 #ifndef NDEBUG
1196         PrintAllVCNs(DeviceExt,
1197                      IndexAllocationContext,
1198                      IndexBufferSize);
1199 #endif
1200     }
1201     // Walk through the root node and update all the sub-nodes
1202     CurrentKey = Tree->RootNode->FirstKey;
1203     for (i = 0; i < Tree->RootNode->KeyCount; i++)
1204     {
1205         if (CurrentKey->LesserChild)
1206         {
1207             if (!HasIndexAllocation)
1208             {
1209                 // We need to add an index allocation to the file record
1210                 PNTFS_ATTR_RECORD EndMarker = (PNTFS_ATTR_RECORD)((ULONG_PTR)FileRecord + FileRecord->BytesInUse - (sizeof(ULONG) * 2));
1211                 DPRINT1("Adding index allocation...\n");
1212
1213                 // Add index allocation to the very end of the file record
1214                 Status = AddIndexAllocation(DeviceExt,
1215                                             FileRecord,
1216                                             EndMarker,
1217                                             L"$I30",
1218                                             4);
1219                 if (!NT_SUCCESS(Status))
1220                 {
1221                     DPRINT1("ERROR: Failed to add index allocation!\n");
1222                     return Status;
1223                 }
1224
1225                 // Find the new attribute
1226                 Status = FindAttribute(DeviceExt, FileRecord, AttributeIndexAllocation, L"$I30", 4, &IndexAllocationContext, &IndexAllocationOffset);
1227                 if (!NT_SUCCESS(Status))
1228                 {
1229                     DPRINT1("ERROR: Couldn't find newly-created index allocation!\n");
1230                     return Status;
1231                 }
1232
1233                 // Advance end marker
1234                 EndMarker = (PNTFS_ATTR_RECORD)((ULONG_PTR)EndMarker + EndMarker->Length);
1235
1236                 // Add index bitmap to the very end of the file record
1237                 Status = AddBitmap(DeviceExt,
1238                                    FileRecord,
1239                                    EndMarker,
1240                                    L"$I30",
1241                                    4);
1242                 if (!NT_SUCCESS(Status))
1243                 {
1244                     DPRINT1("ERROR: Failed to add index bitmap!\n");
1245                     return Status;
1246                 }
1247
1248                 HasIndexAllocation = TRUE;
1249             }
1250
1251             // Is the Index Entry large enough to store the VCN?
1252             if (!CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
1253             {
1254                 // Allocate memory for the larger index entry
1255                 PINDEX_ENTRY_ATTRIBUTE NewEntry = ExAllocatePoolWithTag(NonPagedPool,
1256                                                                         CurrentKey->IndexEntry->Length + sizeof(ULONGLONG),
1257                                                                         TAG_NTFS);
1258                 if (!NewEntry)
1259                 {
1260                     DPRINT1("ERROR: Unable to allocate memory for new index entry!\n");
1261                     if (HasIndexAllocation)
1262                         ReleaseAttributeContext(IndexAllocationContext);
1263                     return STATUS_INSUFFICIENT_RESOURCES;
1264                 }
1265
1266                 // Copy the old entry to the new one
1267                 RtlCopyMemory(NewEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
1268
1269                 NewEntry->Length += sizeof(ULONGLONG);
1270
1271                 // Free the old memory
1272                 ExFreePoolWithTag(CurrentKey->IndexEntry, TAG_NTFS);
1273
1274                 CurrentKey->IndexEntry = NewEntry;
1275                 CurrentKey->IndexEntry->Flags |= NTFS_INDEX_ENTRY_NODE;
1276             }
1277
1278             // Update the sub-node
1279             Status = UpdateIndexNode(DeviceExt, FileRecord, CurrentKey->LesserChild, IndexBufferSize, IndexAllocationContext, IndexAllocationOffset);
1280             if (!NT_SUCCESS(Status))
1281             {
1282                 DPRINT1("ERROR: Failed to update index node!\n");
1283                 ReleaseAttributeContext(IndexAllocationContext);
1284                 return Status;
1285             }
1286
1287             // Update the VCN stored in the index entry of CurrentKey
1288             SetIndexEntryVCN(CurrentKey->IndexEntry, CurrentKey->LesserChild->VCN);
1289         }
1290         CurrentKey = CurrentKey->NextKey;
1291     }
1292
1293 #ifndef NDEBUG
1294     DumpBTree(Tree);
1295 #endif
1296
1297     if (HasIndexAllocation)
1298     {
1299 #ifndef NDEBUG
1300         PrintAllVCNs(DeviceExt,
1301                      IndexAllocationContext,
1302                      IndexBufferSize);
1303 #endif
1304         ReleaseAttributeContext(IndexAllocationContext);
1305     }
1306
1307     return STATUS_SUCCESS;
1308 }
1309
1310 NTSTATUS
1311 UpdateIndexNode(PDEVICE_EXTENSION DeviceExt,
1312                 PFILE_RECORD_HEADER FileRecord,
1313                 PB_TREE_FILENAME_NODE Node,
1314                 ULONG IndexBufferSize,
1315                 PNTFS_ATTR_CONTEXT IndexAllocationContext,
1316                 ULONG IndexAllocationOffset)
1317 {
1318     ULONG i;
1319     PB_TREE_KEY CurrentKey = Node->FirstKey;
1320     BOOLEAN HasChildren = FALSE;
1321     NTSTATUS Status;
1322
1323
1324     DPRINT("UpdateIndexNode(%p, %p, %p, %lu, %p, %lu) called for index node with VCN %I64u\n",
1325            DeviceExt,
1326            FileRecord,
1327            Node,
1328            IndexBufferSize,
1329            IndexAllocationContext,
1330            IndexAllocationOffset,
1331            Node->VCN);
1332
1333     // Walk through the node and look for children to update
1334     for (i = 0; i < Node->KeyCount; i++)
1335     {
1336         ASSERT(CurrentKey);
1337
1338         // If there's a child node
1339         if (CurrentKey->LesserChild)
1340         {
1341             HasChildren = TRUE;
1342
1343             // Update the child node on disk
1344             Status = UpdateIndexNode(DeviceExt, FileRecord, CurrentKey->LesserChild, IndexBufferSize, IndexAllocationContext, IndexAllocationOffset);
1345             if (!NT_SUCCESS(Status))
1346             {
1347                 DPRINT1("ERROR: Failed to update child node!\n");
1348                 return Status;
1349             }
1350
1351             // Is the Index Entry large enough to store the VCN?
1352             if (!CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
1353             {
1354                 // Allocate memory for the larger index entry
1355                 PINDEX_ENTRY_ATTRIBUTE NewEntry = ExAllocatePoolWithTag(NonPagedPool,
1356                                                                         CurrentKey->IndexEntry->Length + sizeof(ULONGLONG),
1357                                                                         TAG_NTFS);
1358                 if (!NewEntry)
1359                 {
1360                     DPRINT1("ERROR: Unable to allocate memory for new index entry!\n");
1361                     return STATUS_INSUFFICIENT_RESOURCES;
1362                 }
1363
1364                 // Copy the old entry to the new one
1365                 RtlCopyMemory(NewEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
1366
1367                 NewEntry->Length += sizeof(ULONGLONG);
1368
1369                 // Free the old memory
1370                 ExFreePoolWithTag(CurrentKey->IndexEntry, TAG_NTFS);
1371
1372                 CurrentKey->IndexEntry = NewEntry;
1373             }
1374
1375             // Update the VCN stored in the index entry of CurrentKey
1376             SetIndexEntryVCN(CurrentKey->IndexEntry, CurrentKey->LesserChild->VCN);
1377
1378             CurrentKey->IndexEntry->Flags |= NTFS_INDEX_ENTRY_NODE;
1379         }
1380
1381         CurrentKey = CurrentKey->NextKey;
1382     }
1383
1384
1385     // Do we need to write this node to disk?
1386     if (Node->DiskNeedsUpdating)
1387     {
1388         ULONGLONG NodeOffset;
1389         ULONG LengthWritten;
1390         PINDEX_BUFFER IndexBuffer;
1391
1392         // Does the node need to be assigned a VCN?
1393         if (!Node->HasValidVCN)
1394         {
1395             // Allocate the node
1396             Status = AllocateIndexNode(DeviceExt,
1397                                        FileRecord,
1398                                        IndexBufferSize,
1399                                        IndexAllocationContext,
1400                                        IndexAllocationOffset,
1401                                        &Node->VCN);
1402             if (!NT_SUCCESS(Status))
1403             {
1404                 DPRINT1("ERROR: Failed to allocate index record in index allocation!\n");
1405                 return Status;
1406             }
1407
1408             Node->HasValidVCN = TRUE;
1409         }
1410
1411         // Allocate memory for an index buffer
1412         IndexBuffer = ExAllocatePoolWithTag(NonPagedPool, IndexBufferSize, TAG_NTFS);
1413         if (!IndexBuffer)
1414         {
1415             DPRINT1("ERROR: Failed to allocate %lu bytes for index buffer!\n", IndexBufferSize);
1416             return STATUS_INSUFFICIENT_RESOURCES;
1417         }
1418
1419         // Create the index buffer we'll be writing to disk to represent this node
1420         Status = CreateIndexBufferFromBTreeNode(DeviceExt, Node, IndexBufferSize, HasChildren, IndexBuffer);
1421         if (!NT_SUCCESS(Status))
1422         {
1423             DPRINT1("ERROR: Failed to create index buffer from node!\n");
1424             ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
1425             return Status;
1426         }
1427
1428         // Get Offset of index buffer in index allocation
1429         NodeOffset = GetAllocationOffsetFromVCN(DeviceExt, IndexBufferSize, Node->VCN);
1430
1431         // Write the buffer to the index allocation
1432         Status = WriteAttribute(DeviceExt, IndexAllocationContext, NodeOffset, (const PUCHAR)IndexBuffer, IndexBufferSize, &LengthWritten, FileRecord);
1433         if (!NT_SUCCESS(Status) || LengthWritten != IndexBufferSize)
1434         {
1435             DPRINT1("ERROR: Failed to update index allocation!\n");
1436             ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
1437             if (!NT_SUCCESS(Status))
1438                 return Status;
1439             else
1440                 return STATUS_END_OF_FILE;
1441         }
1442
1443         Node->DiskNeedsUpdating = FALSE;
1444
1445         // Free the index buffer
1446         ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
1447     }
1448
1449     return STATUS_SUCCESS;
1450 }
1451
1452 PB_TREE_KEY
1453 CreateBTreeKeyFromFilename(ULONGLONG FileReference, PFILENAME_ATTRIBUTE FileNameAttribute)
1454 {
1455     PB_TREE_KEY NewKey;
1456     ULONG AttributeSize = GetFileNameAttributeLength(FileNameAttribute);
1457     ULONG EntrySize = ALIGN_UP_BY(AttributeSize + FIELD_OFFSET(INDEX_ENTRY_ATTRIBUTE, FileName), 8);
1458
1459     // Create a new Index Entry for the file
1460     PINDEX_ENTRY_ATTRIBUTE NewEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_NTFS);
1461     if (!NewEntry)
1462     {
1463         DPRINT1("ERROR: Failed to allocate memory for Index Entry!\n");
1464         return NULL;
1465     }
1466
1467     // Setup the Index Entry
1468     RtlZeroMemory(NewEntry, EntrySize);
1469     NewEntry->Data.Directory.IndexedFile = FileReference;
1470     NewEntry->Length = EntrySize;
1471     NewEntry->KeyLength = AttributeSize;
1472
1473     // Copy the FileNameAttribute
1474     RtlCopyMemory(&NewEntry->FileName, FileNameAttribute, AttributeSize);
1475
1476     // Setup the New Key
1477     NewKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
1478     if (!NewKey)
1479     {
1480         DPRINT1("ERROR: Failed to allocate memory for new key!\n");
1481         ExFreePoolWithTag(NewEntry, TAG_NTFS);
1482         return NULL;
1483     }
1484     NewKey->IndexEntry = NewEntry;
1485     NewKey->NextKey = NULL;
1486
1487     return NewKey;
1488 }
1489
1490 VOID
1491 DestroyBTreeKey(PB_TREE_KEY Key)
1492 {
1493     if (Key->IndexEntry)
1494         ExFreePoolWithTag(Key->IndexEntry, TAG_NTFS);
1495
1496     if (Key->LesserChild)
1497         DestroyBTreeNode(Key->LesserChild);
1498
1499     ExFreePoolWithTag(Key, TAG_NTFS);
1500 }
1501
1502 VOID
1503 DestroyBTreeNode(PB_TREE_FILENAME_NODE Node)
1504 {
1505     PB_TREE_KEY NextKey;
1506     PB_TREE_KEY CurrentKey = Node->FirstKey;
1507     ULONG i;
1508     for (i = 0; i < Node->KeyCount; i++)
1509     {
1510         NT_ASSERT(CurrentKey);
1511         NextKey = CurrentKey->NextKey;
1512         DestroyBTreeKey(CurrentKey);
1513         CurrentKey = NextKey;
1514     }
1515
1516     NT_ASSERT(NextKey == NULL);
1517
1518     ExFreePoolWithTag(Node, TAG_NTFS);
1519 }
1520
1521 /**
1522 * @name DestroyBTree
1523 * @implemented
1524 *
1525 * Destroys a B-Tree.
1526 *
1527 * @param Tree
1528 * Pointer to the B_TREE which will be destroyed.
1529 *
1530 * @remarks
1531 * Destroys every bit of data stored in the tree.
1532 */
1533 VOID
1534 DestroyBTree(PB_TREE Tree)
1535 {
1536     DestroyBTreeNode(Tree->RootNode);
1537     ExFreePoolWithTag(Tree, TAG_NTFS);
1538 }
1539
1540 VOID
1541 DumpBTreeKey(PB_TREE Tree, PB_TREE_KEY Key, ULONG Number, ULONG Depth)
1542 {
1543     ULONG i;
1544     for (i = 0; i < Depth; i++)
1545         DbgPrint(" ");
1546     DbgPrint(" Key #%d", Number);
1547
1548     if (!(Key->IndexEntry->Flags & NTFS_INDEX_ENTRY_END))
1549     {
1550         UNICODE_STRING FileName;
1551         FileName.Length = Key->IndexEntry->FileName.NameLength * sizeof(WCHAR);
1552         FileName.MaximumLength = FileName.Length;
1553         FileName.Buffer = Key->IndexEntry->FileName.Name;
1554         DbgPrint(" '%wZ'\n", &FileName);
1555     }
1556     else
1557     {
1558         DbgPrint(" (Dummy Key)\n");
1559     }
1560
1561     // Is there a child node?
1562     if (Key->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
1563     {
1564         if (Key->LesserChild)
1565             DumpBTreeNode(Tree, Key->LesserChild, Number, Depth + 1);
1566         else
1567         {
1568             // This will be an assert once nodes with arbitrary depth are debugged
1569             DPRINT1("DRIVER ERROR: No Key->LesserChild despite Key->IndexEntry->Flags indicating this is a node!\n");
1570         }
1571     }
1572 }
1573
1574 VOID
1575 DumpBTreeNode(PB_TREE Tree,
1576               PB_TREE_FILENAME_NODE Node,
1577               ULONG Number,
1578               ULONG Depth)
1579 {
1580     PB_TREE_KEY CurrentKey;
1581     ULONG i;
1582     for (i = 0; i < Depth; i++)
1583         DbgPrint(" ");
1584     DbgPrint("Node #%d, Depth %d, has %d key%s", Number, Depth, Node->KeyCount, Node->KeyCount == 1 ? "" : "s");
1585
1586     if (Node->HasValidVCN)
1587         DbgPrint(" VCN: %I64u\n", Node->VCN);
1588     else if (Tree->RootNode == Node)
1589         DbgPrint(" Index Root");
1590     else
1591         DbgPrint(" NOT ASSIGNED VCN YET\n");
1592
1593     CurrentKey = Node->FirstKey;
1594     for (i = 0; i < Node->KeyCount; i++)
1595     {
1596         DumpBTreeKey(Tree, CurrentKey, i, Depth);
1597         CurrentKey = CurrentKey->NextKey;
1598     }
1599 }
1600
1601 /**
1602 * @name DumpBTree
1603 * @implemented
1604 *
1605 * Displays a B-Tree.
1606 *
1607 * @param Tree
1608 * Pointer to the B_TREE which will be displayed.
1609 *
1610 * @remarks
1611 * Displays a diagnostic summary of a B_TREE.
1612 */
1613 VOID
1614 DumpBTree(PB_TREE Tree)
1615 {
1616     DbgPrint("B_TREE @ %p\n", Tree);
1617     DumpBTreeNode(Tree, Tree->RootNode, 0, 0);
1618 }
1619
1620 // Calculates start of Index Buffer relative to the index allocation, given the node's VCN
1621 ULONGLONG
1622 GetAllocationOffsetFromVCN(PDEVICE_EXTENSION DeviceExt,
1623                            ULONG IndexBufferSize,
1624                            ULONGLONG Vcn)
1625 {
1626     if (IndexBufferSize < DeviceExt->NtfsInfo.BytesPerCluster)
1627         return Vcn * DeviceExt->NtfsInfo.BytesPerSector;
1628
1629     return Vcn * DeviceExt->NtfsInfo.BytesPerCluster;
1630 }
1631
1632 ULONGLONG
1633 GetIndexEntryVCN(PINDEX_ENTRY_ATTRIBUTE IndexEntry)
1634 {
1635     PULONGLONG Destination = (PULONGLONG)((ULONG_PTR)IndexEntry + IndexEntry->Length - sizeof(ULONGLONG));
1636
1637     ASSERT(IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE);
1638
1639     return *Destination;
1640 }
1641
1642 /**
1643 * @name NtfsInsertKey
1644 * @implemented
1645 *
1646 * Inserts a FILENAME_ATTRIBUTE into a B-Tree node.
1647 *
1648 * @param Tree
1649 * Pointer to the B_TREE the key (filename attribute) is being inserted into.
1650 *
1651 * @param FileReference
1652 * Reference number to the file being added. This will be a combination of the MFT index and update sequence number.
1653 *
1654 * @param FileNameAttribute
1655 * Pointer to a FILENAME_ATTRIBUTE which is the data for the key that will be added to the tree. A copy will be made.
1656 *
1657 * @param Node
1658 * Pointer to a B_TREE_FILENAME_NODE into which a new key will be inserted, in order.
1659 *
1660 * @param CaseSensitive
1661 * Boolean indicating if the function should operate in case-sensitive mode. This will be TRUE
1662 * if an application created the file with the FILE_FLAG_POSIX_SEMANTICS flag.
1663 *
1664 * @param MaxIndexRootSize
1665 * The maximum size, in bytes, of node entries that can be stored in the index root before it will grow too large for
1666 * the file record. This number is just the size of the entries, without any headers for the attribute or index root.
1667 *
1668 * @param IndexRecordSize
1669 * The size, in bytes, of an index record for this index. AKA an index buffer. Usually set to 4096.
1670 *
1671 * @param MedianKey
1672 * Pointer to a PB_TREE_KEY that will receive a pointer to the median key, should the node grow too large and need to be split.
1673 * Will be set to NULL if the node isn't split.
1674 *
1675 * @param NewRightHandSibling
1676 * Pointer to a PB_TREE_FILENAME_NODE that will receive a pointer to a newly-created right-hand sibling node,
1677 * should the node grow too large and need to be split. Will be set to NULL if the node isn't split.
1678 *
1679 * @remarks
1680 * A node is always sorted, with the least comparable filename stored first and a dummy key to mark the end.
1681 */
1682 NTSTATUS
1683 NtfsInsertKey(PB_TREE Tree,
1684               ULONGLONG FileReference,
1685               PFILENAME_ATTRIBUTE FileNameAttribute,
1686               PB_TREE_FILENAME_NODE Node,
1687               BOOLEAN CaseSensitive,
1688               ULONG MaxIndexRootSize,
1689               ULONG IndexRecordSize,
1690               PB_TREE_KEY *MedianKey,
1691               PB_TREE_FILENAME_NODE *NewRightHandSibling)
1692 {
1693     PB_TREE_KEY NewKey, CurrentKey, PreviousKey;
1694     NTSTATUS Status = STATUS_SUCCESS;
1695     ULONG NodeSize;
1696     ULONG AllocatedNodeSize;
1697     ULONG MaxNodeSizeWithoutHeader;
1698     ULONG i;
1699
1700     *MedianKey = NULL;
1701     *NewRightHandSibling = NULL;
1702
1703     DPRINT1("NtfsInsertKey(%p, 0x%I64x, %p, %p, %s, %lu, %lu, %p, %p)\n",
1704             Tree,
1705             FileReference,
1706             FileNameAttribute,
1707             Node,
1708             CaseSensitive ? "TRUE" : "FALSE",
1709             MaxIndexRootSize,
1710             IndexRecordSize,
1711             MedianKey,
1712             NewRightHandSibling);
1713
1714     // Create the key for the filename attribute
1715     NewKey = CreateBTreeKeyFromFilename(FileReference, FileNameAttribute);
1716     if (!NewKey)
1717         return STATUS_INSUFFICIENT_RESOURCES;
1718
1719     // Find where to insert the key
1720     CurrentKey = Node->FirstKey;
1721     PreviousKey = NULL;
1722     for (i = 0; i < Node->KeyCount; i++)
1723     {
1724         // Should the New Key go before the current key?
1725         LONG Comparison = CompareTreeKeys(NewKey, CurrentKey, CaseSensitive);
1726
1727         if (Comparison == 0)
1728         {
1729             DPRINT1("\t\tComparison == 0: %.*S\n", NewKey->IndexEntry->FileName.NameLength, NewKey->IndexEntry->FileName.Name);
1730             DPRINT1("\t\tComparison == 0: %.*S\n", CurrentKey->IndexEntry->FileName.NameLength, CurrentKey->IndexEntry->FileName.Name);
1731         }
1732         ASSERT(Comparison != 0);
1733
1734         // Is NewKey < CurrentKey?
1735         if (Comparison < 0)
1736         {
1737             // Does CurrentKey have a sub-node?
1738             if (CurrentKey->LesserChild)
1739             {
1740                 PB_TREE_KEY NewLeftKey;
1741                 PB_TREE_FILENAME_NODE NewChild;
1742
1743                 // Insert the key into the child node
1744                 Status = NtfsInsertKey(Tree,
1745                                        FileReference,
1746                                        FileNameAttribute,
1747                                        CurrentKey->LesserChild,
1748                                        CaseSensitive,
1749                                        MaxIndexRootSize,
1750                                        IndexRecordSize,
1751                                        &NewLeftKey,
1752                                        &NewChild);
1753                 if (!NT_SUCCESS(Status))
1754                 {
1755                     DPRINT1("ERROR: Failed to insert key.\n");
1756                     ExFreePoolWithTag(NewKey, TAG_NTFS);
1757                     return Status;
1758                 }
1759
1760                 // Did the child node get split?
1761                 if (NewLeftKey)
1762                 {
1763                     ASSERT(NewChild != NULL);
1764
1765                     // Insert the new left key to the left of the current key
1766                     NewLeftKey->NextKey = CurrentKey;
1767
1768                     // Is CurrentKey the first key?
1769                     if (!PreviousKey)
1770                         Node->FirstKey = NewLeftKey;
1771                     else
1772                         PreviousKey->NextKey = NewLeftKey;
1773
1774                     // CurrentKey->LesserChild will be the right-hand sibling
1775                     CurrentKey->LesserChild = NewChild;
1776
1777                     Node->KeyCount++;
1778                     Node->DiskNeedsUpdating = TRUE;
1779
1780 #ifndef NDEBUG
1781                     DumpBTree(Tree);
1782 #endif
1783                 }
1784             }
1785             else
1786             {
1787                 // Insert New Key before Current Key
1788                 NewKey->NextKey = CurrentKey;
1789
1790                 // Increase KeyCount and mark node as dirty
1791                 Node->KeyCount++;
1792                 Node->DiskNeedsUpdating = TRUE;
1793
1794                 // was CurrentKey the first key?
1795                 if (CurrentKey == Node->FirstKey)
1796                     Node->FirstKey = NewKey;
1797                 else
1798                     PreviousKey->NextKey = NewKey;
1799             }
1800             break;
1801         }
1802
1803         PreviousKey = CurrentKey;
1804         CurrentKey = CurrentKey->NextKey;
1805     }
1806
1807     // Determine how much space the index entries will need
1808     NodeSize = GetSizeOfIndexEntries(Node);
1809
1810     // Is Node not the root node?
1811     if (Node != Tree->RootNode)
1812     {
1813         // Calculate maximum size of index entries without any headers
1814         AllocatedNodeSize = IndexRecordSize - FIELD_OFFSET(INDEX_BUFFER, Header);
1815
1816         // TODO: Replace magic with math
1817         MaxNodeSizeWithoutHeader = AllocatedNodeSize - 0x28;
1818
1819         // Has the node grown larger than its allocated size?
1820         if (NodeSize > MaxNodeSizeWithoutHeader)
1821         {
1822             NTSTATUS Status;
1823
1824             Status = SplitBTreeNode(Tree, Node, MedianKey, NewRightHandSibling, CaseSensitive);
1825             if (!NT_SUCCESS(Status))
1826             {
1827                 DPRINT1("ERROR: Failed to split B-Tree node!\n");
1828                 return Status;
1829             }
1830
1831             return Status;
1832         }
1833     }
1834
1835     // NewEntry and NewKey will be destroyed later by DestroyBTree()
1836
1837     return Status;
1838 }
1839
1840
1841
1842 /**
1843 * @name SplitBTreeNode
1844 * @implemented
1845 *
1846 * Splits a B-Tree node that has grown too large. Finds the median key and sets up a right-hand-sibling
1847 * node to contain the keys to the right of the median key.
1848 *
1849 * @param Tree
1850 * Pointer to the B_TREE which contains the node being split
1851 *
1852 * @param Node
1853 * Pointer to the B_TREE_FILENAME_NODE that needs to be split
1854 *
1855 * @param MedianKey
1856 * Pointer a PB_TREE_KEY that will receive the pointer to the key in the middle of the node being split
1857 *
1858 * @param NewRightHandSibling
1859 * Pointer to a PB_TREE_FILENAME_NODE that will receive a pointer to a newly-created B_TREE_FILENAME_NODE
1860 * containing the keys to the right of MedianKey.
1861 *
1862 * @param CaseSensitive
1863 * Boolean indicating if the function should operate in case-sensitive mode. This will be TRUE
1864 * if an application created the file with the FILE_FLAG_POSIX_SEMANTICS flag.
1865 *
1866 * @return
1867 * STATUS_SUCCESS on success.
1868 * STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
1869 *
1870 * @remarks
1871 * It's the responsibility of the caller to insert the new median key into the parent node, as well as making the
1872 * NewRightHandSibling the lesser child of the node that is currently Node's parent.
1873 */
1874 NTSTATUS
1875 SplitBTreeNode(PB_TREE Tree,
1876                PB_TREE_FILENAME_NODE Node,
1877                PB_TREE_KEY *MedianKey,
1878                PB_TREE_FILENAME_NODE *NewRightHandSibling,
1879                BOOLEAN CaseSensitive)
1880 {
1881     ULONG MedianKeyIndex;
1882     PB_TREE_KEY LastKeyBeforeMedian, FirstKeyAfterMedian;
1883     ULONG KeyCount;
1884     ULONG HalfSize;
1885     ULONG SizeSum;
1886     ULONG i;
1887
1888     DPRINT1("SplitBTreeNode(%p, %p, %p, %p, %s) called\n",
1889             Tree,
1890             Node,
1891             MedianKey,
1892             NewRightHandSibling,
1893             CaseSensitive ? "TRUE" : "FALSE");
1894
1895 #ifndef NDEBUG
1896     DumpBTreeNode(Node, 0, 0);
1897 #endif
1898
1899     // Create the right hand sibling
1900     *NewRightHandSibling = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
1901     if (*NewRightHandSibling == NULL)
1902     {
1903         DPRINT1("Error: Failed to allocate memory for right hand sibling!\n");
1904         return STATUS_INSUFFICIENT_RESOURCES;
1905     }
1906     RtlZeroMemory(*NewRightHandSibling, sizeof(B_TREE_FILENAME_NODE));
1907     (*NewRightHandSibling)->DiskNeedsUpdating = TRUE;
1908
1909
1910     // Find the last key before the median
1911
1912     // This is roughly how NTFS-3G calculates median, and it's not congruent with what Windows does:
1913     /*
1914     // find the median key index
1915     MedianKeyIndex = (Node->KeyCount + 1) / 2;
1916     MedianKeyIndex--;
1917
1918     LastKeyBeforeMedian = Node->FirstKey;
1919     for (i = 0; i < MedianKeyIndex - 1; i++)
1920         LastKeyBeforeMedian = LastKeyBeforeMedian->NextKey;*/
1921
1922     // The method we'll use is a little bit closer to how Windows determines the median but it's not identical.
1923     // What Windows does is actually more complicated than this, I think because Windows allocates more slack space to Odd-numbered
1924     // Index Records, leaving less room for index entries in these records (I haven't discovered why this is done).
1925     // (Neither Windows nor chkdsk complain if we choose a different median than Windows would have chosen, as our median will be in the ballpark)
1926
1927     // Use size to locate the median key / index
1928     LastKeyBeforeMedian = Node->FirstKey;
1929     MedianKeyIndex = 0;
1930     HalfSize = 2016; // half the allocated size after subtracting the first index entry offset (TODO: MATH)
1931     SizeSum = 0;
1932     for (i = 0; i < Node->KeyCount; i++)
1933     {
1934         SizeSum += LastKeyBeforeMedian->IndexEntry->Length;
1935
1936         if (SizeSum > HalfSize)
1937             break;
1938
1939         MedianKeyIndex++;
1940         LastKeyBeforeMedian = LastKeyBeforeMedian->NextKey;
1941     }
1942
1943     // Now we can get the median key and the key that follows it
1944     *MedianKey = LastKeyBeforeMedian->NextKey;
1945     FirstKeyAfterMedian = (*MedianKey)->NextKey;
1946
1947     DPRINT1("%lu keys, %lu median\n", Node->KeyCount, MedianKeyIndex);
1948     DPRINT1("\t\tMedian: %.*S\n", (*MedianKey)->IndexEntry->FileName.NameLength, (*MedianKey)->IndexEntry->FileName.Name);
1949
1950     // "Node" will be the left hand sibling after the split, containing all keys prior to the median key
1951
1952     // We need to create a dummy pointer at the end of the LHS. The dummy's child will be the median's child.
1953     LastKeyBeforeMedian->NextKey = CreateDummyKey(BooleanFlagOn((*MedianKey)->IndexEntry->Flags, NTFS_INDEX_ENTRY_NODE));
1954     if (LastKeyBeforeMedian->NextKey == NULL)
1955     {
1956         DPRINT1("Error: Couldn't allocate dummy key!\n");
1957         LastKeyBeforeMedian->NextKey = *MedianKey;
1958         ExFreePoolWithTag(*NewRightHandSibling, TAG_NTFS);
1959         return STATUS_INSUFFICIENT_RESOURCES;
1960     }
1961
1962     // Did the median key have a child node?
1963     if ((*MedianKey)->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
1964     {
1965         // Set the child of the new dummy key
1966         LastKeyBeforeMedian->NextKey->LesserChild = (*MedianKey)->LesserChild;
1967
1968         // Give the dummy key's index entry the same sub-node VCN the median
1969         SetIndexEntryVCN(LastKeyBeforeMedian->NextKey->IndexEntry, GetIndexEntryVCN((*MedianKey)->IndexEntry));
1970     }
1971     else
1972     {
1973         // Median key didn't have a child node, but it will. Create a new index entry large enough to store a VCN.
1974         PINDEX_ENTRY_ATTRIBUTE NewIndexEntry = ExAllocatePoolWithTag(NonPagedPool,
1975                                                                      (*MedianKey)->IndexEntry->Length + sizeof(ULONGLONG),
1976                                                                      TAG_NTFS);
1977         if (!NewIndexEntry)
1978         {
1979             DPRINT1("Unable to allocate memory for new index entry!\n");
1980             LastKeyBeforeMedian->NextKey = *MedianKey;
1981             ExFreePoolWithTag(*NewRightHandSibling, TAG_NTFS);
1982             return STATUS_INSUFFICIENT_RESOURCES;
1983         }
1984
1985         // Copy the old index entry to the new one
1986         RtlCopyMemory(NewIndexEntry, (*MedianKey)->IndexEntry, (*MedianKey)->IndexEntry->Length);
1987
1988         // Use the new index entry after freeing the old one
1989         ExFreePoolWithTag((*MedianKey)->IndexEntry, TAG_NTFS);
1990         (*MedianKey)->IndexEntry = NewIndexEntry;
1991
1992         // Update the length for the VCN
1993         (*MedianKey)->IndexEntry->Length += sizeof(ULONGLONG);
1994
1995         // Set the node flag
1996         (*MedianKey)->IndexEntry->Flags |= NTFS_INDEX_ENTRY_NODE;
1997     }
1998
1999     // "Node" will become the child of the median key
2000     (*MedianKey)->LesserChild = Node;
2001     SetIndexEntryVCN((*MedianKey)->IndexEntry, Node->VCN);
2002
2003     // Update Node's KeyCount (remember to add 1 for the new dummy key)
2004     Node->KeyCount = MedianKeyIndex + 2;
2005
2006     KeyCount = CountBTreeKeys(Node->FirstKey);
2007     ASSERT(Node->KeyCount == KeyCount);
2008
2009     // everything to the right of MedianKey becomes the right hand sibling of Node
2010     (*NewRightHandSibling)->FirstKey = FirstKeyAfterMedian;
2011     (*NewRightHandSibling)->KeyCount = CountBTreeKeys(FirstKeyAfterMedian);
2012
2013 #ifndef NDEBUG
2014     DPRINT1("Left-hand node after split:\n");
2015     DumpBTreeNode(Node, 0, 0);
2016
2017     DPRINT1("Right-hand sibling node after split:\n");
2018     DumpBTreeNode(*NewRightHandSibling, 0, 0);
2019 #endif
2020
2021     return STATUS_SUCCESS;
2022 }