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 CompareTreeKeys
  85 * @implemented
  86 *
  87 * Compare two B_TREE_KEY's to determine their order in the tree.
  88 *
  89 * @param Key1
  90 * Pointer to a B_TREE_KEY that will be compared.
  91 *
  92 * @param Key2
  93 * Pointer to the other B_TREE_KEY that will be compared.
  94 *
  95 * @param CaseSensitive
  96 * Boolean indicating if the function should operate in case-sensitive mode. This will be TRUE
  97 * if an application created the file with the FILE_FLAG_POSIX_SEMANTICS flag.
  98 *
  99 * @returns
 100 * 0 if the two keys are equal.
 101 * < 0 if key1 is less thank key2
 102 * > 0 if key1 is greater than key2
 103 *
 104 * @remarks
 105 * Any other key is always less than the final (dummy) key in a node. Key1 must not be the dummy node.
 106 */
 107 LONG
 108 CompareTreeKeys(PB_TREE_KEY Key1, PB_TREE_KEY Key2, BOOLEAN CaseSensitive)
 109 {
 110     UNICODE_STRING Key1Name, Key2Name;
 111     LONG Comparison;
 112
 113     // Key1 must not be the final key (AKA the dummy key)
 114     ASSERT(!(Key1->IndexEntry->Flags & NTFS_INDEX_ENTRY_END));
 115
 116     // If Key2 is the "dummy key", key 1 will always come first
 117     if (Key2->NextKey == NULL)
 118         return -1;
 119
 120     Key1Name.Buffer = Key1->IndexEntry->FileName.Name;
 121     Key1Name.Length = Key1Name.MaximumLength
 122         = Key1->IndexEntry->FileName.NameLength * sizeof(WCHAR);
 123
 124     Key2Name.Buffer = Key2->IndexEntry->FileName.Name;
 125     Key2Name.Length = Key2Name.MaximumLength
 126         = Key2->IndexEntry->FileName.NameLength * sizeof(WCHAR);
 127
 128     // Are the two keys the same length?
 129     if (Key1Name.Length == Key2Name.Length)
 130         return RtlCompareUnicodeString(&Key1Name, &Key2Name, !CaseSensitive);
 131
 132     // Is Key1 shorter?
 133     if (Key1Name.Length < Key2Name.Length)
 134     {
 135         // Truncate KeyName2 to be the same length as KeyName1
 136         Key2Name.Length = Key1Name.Length;
 137
 138         // Compare the names of the same length
 139         Comparison = RtlCompareUnicodeString(&Key1Name, &Key2Name, !CaseSensitive);
 140
 141         // If the truncated names are the same length, the shorter one comes first
 142         if (Comparison == 0)
 143             return -1;
 144     }
 145     else
 146     {
 147         // Key2 is shorter
 148         // Truncate KeyName1 to be the same length as KeyName2
 149         Key1Name.Length = Key2Name.Length;
 150
 151         // Compare the names of the same length
 152         Comparison = RtlCompareUnicodeString(&Key1Name, &Key2Name, !CaseSensitive);
 153
 154         // If the truncated names are the same length, the shorter one comes first
 155         if (Comparison == 0)
 156             return 1;
 157     }
 158
 159     return Comparison;
 160 }
 161
 162 PB_TREE_FILENAME_NODE
 163 CreateBTreeNodeFromIndexNode(PDEVICE_EXTENSION Vcb,
 164                              PINDEX_ROOT_ATTRIBUTE IndexRoot,
 165                              PNTFS_ATTR_CONTEXT IndexAllocationAttributeCtx,
 166                              PINDEX_ENTRY_ATTRIBUTE NodeEntry)
 167 {
 168     PB_TREE_FILENAME_NODE NewNode;
 169     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 170     PINDEX_ENTRY_ATTRIBUTE FirstNodeEntry;
 171     ULONG CurrentEntryOffset = 0;
 172     PINDEX_BUFFER NodeBuffer;
 173     ULONG IndexBufferSize = Vcb->NtfsInfo.BytesPerIndexRecord;
 174     PULONGLONG NodeNumber;
 175     PB_TREE_KEY CurrentKey;
 176     NTSTATUS Status;
 177     ULONGLONG IndexNodeOffset;
 178     ULONG BytesRead;
 179
 180     if (IndexAllocationAttributeCtx == NULL)
 181     {
 182         DPRINT1("ERROR: Couldn't find index allocation attribute even though there should be one!\n");
 183         return NULL;
 184     }
 185
 186     // Get the node number from the end of the node entry
 187     NodeNumber = (PULONGLONG)((ULONG_PTR)NodeEntry + NodeEntry->Length - sizeof(ULONGLONG));
 188
 189     // Create the new tree node
 190     DPRINT1("About to allocate %ld for NewNode\n", sizeof(B_TREE_FILENAME_NODE));
 191     NewNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
 192     if (!NewNode)
 193     {
 194         DPRINT1("ERROR: Couldn't allocate memory for new filename node.\n");
 195         return NULL;
 196     }
 197     RtlZeroMemory(NewNode, sizeof(B_TREE_FILENAME_NODE));
 198
 199     // Create the first key
 200     CurrentKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 201     if (!CurrentKey)
 202     {
 203         DPRINT1("ERROR: Failed to allocate memory for key!\n");
 204         ExFreePoolWithTag(NewNode, TAG_NTFS);
 205         return NULL;
 206     }
 207     RtlZeroMemory(CurrentKey, sizeof(B_TREE_KEY));
 208     NewNode->FirstKey = CurrentKey;
 209
 210     // Allocate memory for the node buffer
 211     NodeBuffer = ExAllocatePoolWithTag(NonPagedPool, IndexBufferSize, TAG_NTFS);
 212     if (!NodeBuffer)
 213     {
 214         DPRINT1("ERROR: Couldn't allocate memory for node buffer!\n");
 215         ExFreePoolWithTag(CurrentKey, TAG_NTFS);
 216         ExFreePoolWithTag(NewNode, TAG_NTFS);
 217         return NULL;
 218     }
 219
 220     // Calculate offset into index allocation
 221     IndexNodeOffset = GetAllocationOffsetFromVCN(Vcb, IndexBufferSize, *NodeNumber);
 222
 223     // TODO: Confirm index bitmap has this node marked as in-use
 224
 225     // Read the node
 226     BytesRead = ReadAttribute(Vcb,
 227                               IndexAllocationAttributeCtx,
 228                               IndexNodeOffset,
 229                               (PCHAR)NodeBuffer,
 230                               IndexBufferSize);
 231
 232     ASSERT(BytesRead == IndexBufferSize);
 233     NT_ASSERT(NodeBuffer->Ntfs.Type == NRH_INDX_TYPE);
 234     NT_ASSERT(NodeBuffer->VCN == *NodeNumber);
 235
 236     // Apply the fixup array to the node buffer
 237     Status = FixupUpdateSequenceArray(Vcb, &NodeBuffer->Ntfs);
 238     if (!NT_SUCCESS(Status))
 239     {
 240         DPRINT1("ERROR: Couldn't apply fixup array to index node buffer!\n");
 241         ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 242         ExFreePoolWithTag(CurrentKey, TAG_NTFS);
 243         ExFreePoolWithTag(NewNode, TAG_NTFS);
 244         return NULL;
 245     }
 246
 247     // Walk through the index and create keys for all the entries
 248     FirstNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)(&NodeBuffer->Header)
 249                                                + NodeBuffer->Header.FirstEntryOffset);
 250     CurrentNodeEntry = FirstNodeEntry;
 251     while (CurrentEntryOffset < NodeBuffer->Header.TotalSizeOfEntries)
 252     {
 253         // Allocate memory for the current entry
 254         CurrentKey->IndexEntry = ExAllocatePoolWithTag(NonPagedPool, CurrentNodeEntry->Length, TAG_NTFS);
 255         if (!CurrentKey->IndexEntry)
 256         {
 257             DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 258             DestroyBTreeNode(NewNode);
 259             ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 260             return NULL;
 261         }
 262
 263         NewNode->KeyCount++;
 264
 265         // If this isn't the last entry
 266         if (!(CurrentNodeEntry->Flags & NTFS_INDEX_ENTRY_END))
 267         {
 268             // Create the next key
 269             PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 270             if (!NextKey)
 271             {
 272                 DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 273                 DestroyBTreeNode(NewNode);
 274                 ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 275                 return NULL;
 276             }
 277             RtlZeroMemory(NextKey, sizeof(B_TREE_KEY));
 278
 279             // Add NextKey to the end of the list
 280             CurrentKey->NextKey = NextKey;
 281
 282             // Copy the current entry to its key
 283             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 284
 285             // See if the current key has a sub-node
 286             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 287             {
 288                 DPRINT1("TODO: Only a node with a single-level is supported right now!\n");
 289                 // Needs debugging:
 290                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 291                                                                        IndexRoot,
 292                                                                        IndexAllocationAttributeCtx,
 293                                                                        CurrentKey->IndexEntry);
 294             }
 295
 296             CurrentKey = NextKey;
 297         }
 298         else
 299         {
 300             // Copy the final entry to its key
 301             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 302             CurrentKey->NextKey = NULL;
 303
 304             // See if the current key has a sub-node
 305             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 306             {
 307                 DPRINT1("TODO: Only a node with a single-level is supported right now!\n");
 308                 // Needs debugging:
 309                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 310                                                                        IndexRoot,
 311                                                                        IndexAllocationAttributeCtx,
 312                                                                        CurrentKey->IndexEntry);
 313             }
 314
 315             break;
 316         }
 317
 318         // Advance to the next entry
 319         CurrentEntryOffset += CurrentNodeEntry->Length;
 320         CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 321     }
 322
 323     NewNode->NodeNumber = *NodeNumber;
 324     NewNode->ExistsOnDisk = TRUE;
 325
 326     ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
 327
 328     return NewNode;
 329 }
 330
 331 /**
 332 * @name CreateBTreeFromIndex
 333 * @implemented
 334 *
 335 * Parse an index and create a B-Tree in memory from it.
 336 *
 337 * @param IndexRootContext
 338 * Pointer to an NTFS_ATTR_CONTEXT that describes the location of the index root attribute.
 339 *
 340 * @param NewTree
 341 * Pointer to a PB_TREE that will receive the pointer to a newly-created B-Tree.
 342 *
 343 * @returns
 344 * STATUS_SUCCESS on success.
 345 * STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
 346 *
 347 * @remarks
 348 * Allocates memory for the entire tree. Caller is responsible for destroying the tree with DestroyBTree().
 349 */
 350 NTSTATUS
 351 CreateBTreeFromIndex(PDEVICE_EXTENSION Vcb,
 352                      PFILE_RECORD_HEADER FileRecordWithIndex,
 353                      /*PCWSTR IndexName,*/
 354                      PNTFS_ATTR_CONTEXT IndexRootContext,
 355                      PINDEX_ROOT_ATTRIBUTE IndexRoot,
 356                      PB_TREE *NewTree)
 357 {
 358     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 359     PB_TREE Tree = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE), TAG_NTFS);
 360     PB_TREE_FILENAME_NODE RootNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
 361     PB_TREE_KEY CurrentKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 362     ULONG CurrentOffset = IndexRoot->Header.FirstEntryOffset;
 363     PNTFS_ATTR_CONTEXT IndexAllocationContext = NULL;
 364     NTSTATUS Status;
 365
 366     DPRINT1("CreateBTreeFromIndex(%p, %p)\n", IndexRoot, NewTree);
 367
 368     if (!Tree || !RootNode || !CurrentKey)
 369     {
 370         DPRINT1("Couldn't allocate enough memory for B-Tree!\n");
 371         if (Tree)
 372             ExFreePoolWithTag(Tree, TAG_NTFS);
 373         if (CurrentKey)
 374             ExFreePoolWithTag(CurrentKey, TAG_NTFS);
 375         if (RootNode)
 376             ExFreePoolWithTag(RootNode, TAG_NTFS);
 377         return STATUS_INSUFFICIENT_RESOURCES;
 378     }
 379
 380     RtlZeroMemory(Tree, sizeof(B_TREE));
 381     RtlZeroMemory(RootNode, sizeof(B_TREE_FILENAME_NODE));
 382     RtlZeroMemory(CurrentKey, sizeof(B_TREE_KEY));
 383
 384     // See if the file record has an attribute allocation
 385     Status = FindAttribute(Vcb,
 386                            FileRecordWithIndex,
 387                            AttributeIndexAllocation,
 388                            L"$I30",
 389                            4,
 390                            &IndexAllocationContext,
 391                            NULL);
 392     if (!NT_SUCCESS(Status))
 393         IndexAllocationContext = NULL;
 394     else
 395         PrintAllVCNs(Vcb, IndexAllocationContext, IndexRoot->SizeOfEntry);
 396
 397     // Setup the Tree
 398     RootNode->FirstKey = CurrentKey;
 399     Tree->RootNode = RootNode;
 400
 401     // Make sure we won't try reading past the attribute-end
 402     if (FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header) + IndexRoot->Header.TotalSizeOfEntries > IndexRootContext->pRecord->Resident.ValueLength)
 403     {
 404         DPRINT1("Filesystem corruption detected!\n");
 405         DestroyBTree(Tree);
 406         return STATUS_FILE_CORRUPT_ERROR;
 407     }
 408
 409     // Start at the first node entry
 410     CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)IndexRoot
 411                                                 + FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header)
 412                                                 + IndexRoot->Header.FirstEntryOffset);
 413
 414     // Create a key for each entry in the node
 415     while (CurrentOffset < IndexRoot->Header.TotalSizeOfEntries)
 416     {
 417         // Allocate memory for the current entry
 418         CurrentKey->IndexEntry = ExAllocatePoolWithTag(NonPagedPool, CurrentNodeEntry->Length, TAG_NTFS);
 419         if (!CurrentKey->IndexEntry)
 420         {
 421             DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 422             DestroyBTree(Tree);
 423             return STATUS_INSUFFICIENT_RESOURCES;
 424         }
 425
 426         RootNode->KeyCount++;
 427
 428         // If this isn't the last entry
 429         if (!(CurrentNodeEntry->Flags & NTFS_INDEX_ENTRY_END))
 430         {
 431             // Create the next key
 432             PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 433             if (!NextKey)
 434             {
 435                 DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
 436                 DestroyBTree(Tree);
 437                 return STATUS_INSUFFICIENT_RESOURCES;
 438             }
 439
 440             RtlZeroMemory(NextKey, sizeof(B_TREE_KEY));
 441
 442             // Add NextKey to the end of the list
 443             CurrentKey->NextKey = NextKey;
 444
 445             // Copy the current entry to its key
 446             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 447
 448             // Does this key have a sub-node?
 449             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 450             {
 451                 // Create the child node
 452                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 453                                                                        IndexRoot,
 454                                                                        IndexAllocationContext,
 455                                                                        CurrentKey->IndexEntry);
 456                 if (!CurrentKey->LesserChild)
 457                 {
 458                     DPRINT1("ERROR: Couldn't create child node!\n");
 459                     DestroyBTree(Tree);
 460                     return STATUS_NOT_IMPLEMENTED;
 461                 }
 462             }
 463
 464             // Advance to the next entry
 465             CurrentOffset += CurrentNodeEntry->Length;
 466             CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 467             CurrentKey = NextKey;
 468         }
 469         else
 470         {
 471             // Copy the final entry to its key
 472             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 473             CurrentKey->NextKey = NULL;
 474
 475             // Does this key have a sub-node?
 476             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 477             {
 478                 // Create the child node
 479                 CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
 480                                                                        IndexRoot,
 481                                                                        IndexAllocationContext,
 482                                                                        CurrentKey->IndexEntry);
 483                 if (!CurrentKey->LesserChild)
 484                 {
 485                     DPRINT1("ERROR: Couldn't create child node!\n");
 486                     DestroyBTree(Tree);
 487                     return STATUS_NOT_IMPLEMENTED;
 488                 }
 489             }
 490
 491             break;
 492         }
 493     }
 494
 495     *NewTree = Tree;
 496
 497     if (IndexAllocationContext)
 498         ReleaseAttributeContext(IndexAllocationContext);
 499
 500     return STATUS_SUCCESS;
 501 }
 502
 503 /**
 504 * @name CreateIndexRootFromBTree
 505 * @implemented
 506 *
 507 * Parse a B-Tree in memory and convert it into an index that can be written to disk.
 508 *
 509 * @param DeviceExt
 510 * Pointer to the DEVICE_EXTENSION of the target drive.
 511 *
 512 * @param Tree
 513 * Pointer to a B_TREE that describes the index to be written.
 514 *
 515 * @param MaxIndexSize
 516 * Describes how large the index can be before it will take too much space in the file record.
 517 * After reaching MaxIndexSize, an index can no longer be represented with just an index root
 518 * attribute, and will require an index allocation and $I30 bitmap (TODO).
 519 *
 520 * @param IndexRoot
 521 * Pointer to a PINDEX_ROOT_ATTRIBUTE that will receive a pointer to the newly-created index.
 522 *
 523 * @param Length
 524 * Pointer to a ULONG which will receive the length of the new index root.
 525 *
 526 * @returns
 527 * STATUS_SUCCESS on success.
 528 * STATUS_INSUFFICIENT_RESOURCES if an allocation fails.
 529 * STATUS_NOT_IMPLEMENTED if the new index can't fit within MaxIndexSize.
 530 *
 531 * @remarks
 532 * If the function succeeds, it's the caller's responsibility to free IndexRoot with ExFreePoolWithTag().
 533 */
 534 NTSTATUS
 535 CreateIndexRootFromBTree(PDEVICE_EXTENSION DeviceExt,
 536                          PB_TREE Tree,
 537                          ULONG MaxIndexSize,
 538                          PINDEX_ROOT_ATTRIBUTE *IndexRoot,
 539                          ULONG *Length)
 540 {
 541     ULONG i;
 542     PB_TREE_KEY CurrentKey;
 543     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 544     PINDEX_ROOT_ATTRIBUTE NewIndexRoot = ExAllocatePoolWithTag(NonPagedPool,
 545                                                                DeviceExt->NtfsInfo.BytesPerFileRecord,
 546                                                                TAG_NTFS);
 547
 548     DPRINT1("CreateIndexRootFromBTree(%p, %p, 0x%lx, %p, %p)\n", DeviceExt, Tree, MaxIndexSize, IndexRoot, Length);
 549
 550     if (!NewIndexRoot)
 551     {
 552         DPRINT1("Failed to allocate memory for Index Root!\n");
 553         return STATUS_INSUFFICIENT_RESOURCES;
 554     }
 555
 556     // Setup the new index root
 557     RtlZeroMemory(NewIndexRoot, DeviceExt->NtfsInfo.BytesPerFileRecord);
 558
 559     NewIndexRoot->AttributeType = AttributeFileName;
 560     NewIndexRoot->CollationRule = COLLATION_FILE_NAME;
 561     NewIndexRoot->SizeOfEntry = DeviceExt->NtfsInfo.BytesPerIndexRecord;
 562     // If Bytes per index record is less than cluster size, clusters per index record becomes sectors per index
 563     if (NewIndexRoot->SizeOfEntry < DeviceExt->NtfsInfo.BytesPerCluster)
 564         NewIndexRoot->ClustersPerIndexRecord = NewIndexRoot->SizeOfEntry / DeviceExt->NtfsInfo.BytesPerSector;
 565     else
 566         NewIndexRoot->ClustersPerIndexRecord = NewIndexRoot->SizeOfEntry / DeviceExt->NtfsInfo.BytesPerCluster;
 567
 568     // Setup the Index node header
 569     NewIndexRoot->Header.FirstEntryOffset = sizeof(INDEX_HEADER_ATTRIBUTE);
 570     NewIndexRoot->Header.Flags = INDEX_ROOT_SMALL;
 571
 572     // Start summing the total size of this node's entries
 573     NewIndexRoot->Header.TotalSizeOfEntries = NewIndexRoot->Header.FirstEntryOffset;
 574
 575     // Setup each Node Entry
 576     CurrentKey = Tree->RootNode->FirstKey;
 577     CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)NewIndexRoot
 578                                                 + FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header)
 579                                                 + NewIndexRoot->Header.FirstEntryOffset);
 580     for (i = 0; i < Tree->RootNode->KeyCount; i++)
 581     {
 582         // Would adding the current entry to the index increase the index size beyond the limit we've set?
 583         ULONG IndexSize = FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header)
 584                           + NewIndexRoot->Header.TotalSizeOfEntries
 585                           + CurrentNodeEntry->Length;
 586         if (IndexSize > MaxIndexSize)
 587         {
 588             DPRINT1("TODO: Adding file would require creating an index allocation!\n");
 589             ExFreePoolWithTag(NewIndexRoot, TAG_NTFS);
 590             return STATUS_NOT_IMPLEMENTED;
 591         }
 592
 593         ASSERT(CurrentKey->IndexEntry->Length != 0);
 594
 595         // Copy the index entry
 596         RtlCopyMemory(CurrentNodeEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
 597
 598         DPRINT1("Index Node Entry Stream Length: %u\nIndex Node Entry Length: %u\n",
 599                 CurrentNodeEntry->KeyLength,
 600                 CurrentNodeEntry->Length);
 601
 602         // Does the current key have any sub-nodes?
 603         if (CurrentKey->LesserChild)
 604             NewIndexRoot->Header.Flags = INDEX_ROOT_LARGE;
 605
 606         // Add Length of Current Entry to Total Size of Entries
 607         NewIndexRoot->Header.TotalSizeOfEntries += CurrentNodeEntry->Length;
 608
 609         // Go to the next node entry
 610         CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 611
 612         CurrentKey = CurrentKey->NextKey;
 613     }
 614
 615     NewIndexRoot->Header.AllocatedSize = NewIndexRoot->Header.TotalSizeOfEntries;
 616
 617     *IndexRoot = NewIndexRoot;
 618     *Length = NewIndexRoot->Header.AllocatedSize + FIELD_OFFSET(INDEX_ROOT_ATTRIBUTE, Header);
 619
 620     return STATUS_SUCCESS;
 621 }
 622
 623 NTSTATUS
 624 CreateIndexBufferFromBTreeNode(PDEVICE_EXTENSION DeviceExt,
 625                                PB_TREE_FILENAME_NODE Node,
 626                                ULONG BufferSize,
 627                                PINDEX_BUFFER IndexBuffer)
 628 {
 629     ULONG i;
 630     PB_TREE_KEY CurrentKey;
 631     PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
 632     NTSTATUS Status;
 633
 634     // TODO: Fix magic, do math
 635     RtlZeroMemory(IndexBuffer, BufferSize);
 636     IndexBuffer->Ntfs.Type = NRH_INDX_TYPE;
 637     IndexBuffer->Ntfs.UsaOffset = 0x28;
 638     IndexBuffer->Ntfs.UsaCount = 9;
 639
 640     // TODO: Check bitmap for VCN
 641     ASSERT(Node->ExistsOnDisk);
 642     IndexBuffer->VCN = Node->NodeNumber;
 643
 644     IndexBuffer->Header.FirstEntryOffset = 0x28;
 645     IndexBuffer->Header.AllocatedSize = BufferSize - FIELD_OFFSET(INDEX_BUFFER, Header);
 646
 647     // Start summing the total size of this node's entries
 648     IndexBuffer->Header.TotalSizeOfEntries = IndexBuffer->Header.FirstEntryOffset;
 649
 650     CurrentKey = Node->FirstKey;
 651     CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)&(IndexBuffer->Header)
 652                                                 + IndexBuffer->Header.FirstEntryOffset);
 653     for (i = 0; i < Node->KeyCount; i++)
 654     {
 655         // Would adding the current entry to the index increase the node size beyond the allocation size?
 656         ULONG IndexSize = FIELD_OFFSET(INDEX_BUFFER, Header)
 657             + IndexBuffer->Header.TotalSizeOfEntries
 658             + CurrentNodeEntry->Length;
 659         if (IndexSize > BufferSize)
 660         {
 661             DPRINT1("TODO: Adding file would require creating a new node!\n");
 662             return STATUS_NOT_IMPLEMENTED;
 663         }
 664
 665         ASSERT(CurrentKey->IndexEntry->Length != 0);
 666
 667         // Copy the index entry
 668         RtlCopyMemory(CurrentNodeEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
 669
 670         DPRINT1("Index Node Entry Stream Length: %u\nIndex Node Entry Length: %u\n",
 671                 CurrentNodeEntry->KeyLength,
 672                 CurrentNodeEntry->Length);
 673
 674         // Add Length of Current Entry to Total Size of Entries
 675         IndexBuffer->Header.TotalSizeOfEntries += CurrentNodeEntry->Length;
 676
 677         // TODO: Check for child nodes
 678
 679         // Go to the next node entry
 680         CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
 681         CurrentKey = CurrentKey->NextKey;
 682     }
 683
 684     Status = AddFixupArray(DeviceExt, &IndexBuffer->Ntfs);
 685
 686     return Status;
 687 }
 688
 689 NTSTATUS
 690 UpdateIndexAllocation(PDEVICE_EXTENSION DeviceExt,
 691                       PB_TREE Tree,
 692                       ULONG IndexBufferSize,
 693                       PFILE_RECORD_HEADER FileRecord)
 694 {
 695     // Find the index allocation and bitmap
 696     PNTFS_ATTR_CONTEXT IndexAllocationContext, BitmapContext;
 697     PB_TREE_KEY CurrentKey;
 698     NTSTATUS Status;
 699     BOOLEAN HasIndexAllocation = FALSE;
 700     ULONG i;
 701
 702     DPRINT1("UpdateIndexAllocations() called.\n");
 703
 704     Status = FindAttribute(DeviceExt, FileRecord, AttributeIndexAllocation, L"$I30", 4, &IndexAllocationContext, NULL);
 705     if (NT_SUCCESS(Status))
 706         HasIndexAllocation = TRUE;
 707
 708     // TODO: Handle bitmap
 709     BitmapContext = NULL;
 710
 711     // Walk through the root node and update all the sub-nodes
 712     CurrentKey = Tree->RootNode->FirstKey;
 713     for (i = 0; i < Tree->RootNode->KeyCount; i++)
 714     {
 715         if (CurrentKey->LesserChild)
 716         {
 717             if (!HasIndexAllocation)
 718             {
 719                 DPRINT1("FIXME: Need to add index allocation\n");
 720                 return STATUS_NOT_IMPLEMENTED;
 721             }
 722             else
 723             {
 724                 Status = UpdateIndexNode(DeviceExt, CurrentKey->LesserChild, IndexBufferSize, IndexAllocationContext, BitmapContext);
 725                 if (!NT_SUCCESS(Status))
 726                 {
 727                     DPRINT1("ERROR: Failed to update index node!\n");
 728                     ReleaseAttributeContext(IndexAllocationContext);
 729                     return Status;
 730                 }
 731             }
 732
 733         }
 734         CurrentKey = CurrentKey->NextKey;
 735     }
 736
 737     if(HasIndexAllocation)
 738         ReleaseAttributeContext(IndexAllocationContext);
 739
 740     return STATUS_SUCCESS;
 741 }
 742
 743 NTSTATUS
 744 UpdateIndexNode(PDEVICE_EXTENSION DeviceExt,
 745                 PB_TREE_FILENAME_NODE Node,
 746                 ULONG IndexBufferSize,
 747                 PNTFS_ATTR_CONTEXT IndexAllocationContext,
 748                 PNTFS_ATTR_CONTEXT BitmapContext)
 749 {
 750     ULONG i;
 751     PB_TREE_KEY CurrentKey = Node->FirstKey;
 752     NTSTATUS Status;
 753
 754     DPRINT1("UpdateIndexNode(%p, %p, %lu, %p, %p) called for index node with VCN %I64u\n", DeviceExt, Node, IndexBufferSize, IndexAllocationContext, BitmapContext, Node->NodeNumber);
 755
 756     // Do we need to write this node to disk?
 757     if (Node->DiskNeedsUpdating)
 758     {
 759         ULONGLONG NodeOffset;
 760         ULONG LengthWritten;
 761
 762         // Allocate memory for an index buffer
 763         PINDEX_BUFFER IndexBuffer = ExAllocatePoolWithTag(NonPagedPool, IndexBufferSize, TAG_NTFS);
 764         if (!IndexBuffer)
 765         {
 766             DPRINT1("ERROR: Failed to allocate %lu bytes for index buffer!\n", IndexBufferSize);
 767             return STATUS_INSUFFICIENT_RESOURCES;
 768         }
 769
 770         // Create the index buffer we'll be writing to disk to represent this node
 771         Status = CreateIndexBufferFromBTreeNode(DeviceExt, Node, IndexBufferSize, IndexBuffer);
 772         if (!NT_SUCCESS(Status))
 773         {
 774             DPRINT1("ERROR: Failed to create index buffer from node!\n");
 775             ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
 776             return Status;
 777         }
 778
 779         // Get Offset of index buffer in index allocation
 780         NodeOffset = GetAllocationOffsetFromVCN(DeviceExt, IndexBufferSize, Node->NodeNumber);
 781
 782         // Write the buffer to the index allocation
 783         Status = WriteAttribute(DeviceExt, IndexAllocationContext, NodeOffset, (const PUCHAR)IndexBuffer, IndexBufferSize, &LengthWritten, NULL);
 784         if (!NT_SUCCESS(Status) || LengthWritten != IndexBufferSize)
 785         {
 786             DPRINT1("ERROR: Failed to update index allocation!\n");
 787             ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
 788             if (!NT_SUCCESS(Status))
 789                 return Status;
 790             else
 791                 return STATUS_END_OF_FILE;
 792         }
 793
 794         Node->DiskNeedsUpdating = FALSE;
 795
 796         // Free the index buffer
 797         ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
 798     }
 799
 800     // Walk through the node and look for children to update
 801     for (i = 0; i < Node->KeyCount; i++)
 802     {
 803         ASSERT(CurrentKey);
 804
 805         // If there's a child node
 806         if (CurrentKey->LesserChild)
 807         {
 808             // Update the child node on disk
 809             Status = UpdateIndexNode(DeviceExt, CurrentKey->LesserChild, IndexBufferSize, IndexAllocationContext, BitmapContext);
 810             if (!NT_SUCCESS(Status))
 811             {
 812                 DPRINT1("ERROR: Failed to update child node!\n");
 813                 return Status;
 814             }
 815         }
 816
 817         CurrentKey = CurrentKey->NextKey;
 818     }
 819
 820     return STATUS_SUCCESS;
 821 }
 822
 823 PB_TREE_KEY
 824 CreateBTreeKeyFromFilename(ULONGLONG FileReference, PFILENAME_ATTRIBUTE FileNameAttribute)
 825 {
 826     PB_TREE_KEY NewKey;
 827     ULONG AttributeSize = GetFileNameAttributeLength(FileNameAttribute);
 828     ULONG EntrySize = ALIGN_UP_BY(AttributeSize + FIELD_OFFSET(INDEX_ENTRY_ATTRIBUTE, FileName), 8);
 829
 830     // Create a new Index Entry for the file
 831     PINDEX_ENTRY_ATTRIBUTE NewEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_NTFS);
 832     if (!NewEntry)
 833     {
 834         DPRINT1("ERROR: Failed to allocate memory for Index Entry!\n");
 835         return NULL;
 836     }
 837
 838     // Setup the Index Entry
 839     RtlZeroMemory(NewEntry, EntrySize);
 840     NewEntry->Data.Directory.IndexedFile = FileReference;
 841     NewEntry->Length = EntrySize;
 842     NewEntry->KeyLength = AttributeSize;
 843
 844     // Copy the FileNameAttribute
 845     RtlCopyMemory(&NewEntry->FileName, FileNameAttribute, AttributeSize);
 846
 847     // Setup the New Key
 848     NewKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
 849     if (!NewKey)
 850     {
 851         DPRINT1("ERROR: Failed to allocate memory for new key!\n");
 852         ExFreePoolWithTag(NewEntry, TAG_NTFS);
 853         return NULL;
 854     }
 855     NewKey->IndexEntry = NewEntry;
 856     NewKey->NextKey = NULL;
 857
 858     return NewKey;
 859 }
 860
 861 VOID
 862 DestroyBTreeKey(PB_TREE_KEY Key)
 863 {
 864     if (Key->IndexEntry)
 865         ExFreePoolWithTag(Key->IndexEntry, TAG_NTFS);
 866
 867     if (Key->LesserChild)
 868         DestroyBTreeNode(Key->LesserChild);
 869
 870     ExFreePoolWithTag(Key, TAG_NTFS);
 871 }
 872
 873 VOID
 874 DestroyBTreeNode(PB_TREE_FILENAME_NODE Node)
 875 {
 876     PB_TREE_KEY NextKey;
 877     PB_TREE_KEY CurrentKey = Node->FirstKey;
 878     ULONG i;
 879     for (i = 0; i < Node->KeyCount; i++)
 880     {
 881         NT_ASSERT(CurrentKey);
 882         NextKey = CurrentKey->NextKey;
 883         DestroyBTreeKey(CurrentKey);
 884         CurrentKey = NextKey;
 885     }
 886
 887     NT_ASSERT(NextKey == NULL);
 888
 889     ExFreePoolWithTag(Node, TAG_NTFS);
 890 }
 891
 892 /**
 893 * @name DestroyBTree
 894 * @implemented
 895 *
 896 * Destroys a B-Tree.
 897 *
 898 * @param Tree
 899 * Pointer to the B_TREE which will be destroyed.
 900 *
 901 * @remarks
 902 * Destroys every bit of data stored in the tree.
 903 */
 904 VOID
 905 DestroyBTree(PB_TREE Tree)
 906 {
 907     DestroyBTreeNode(Tree->RootNode);
 908     ExFreePoolWithTag(Tree, TAG_NTFS);
 909 }
 910
 911 VOID
 912 DumpBTreeKey(PB_TREE_KEY Key, ULONG Number, ULONG Depth)
 913 {
 914     ULONG i;
 915     for (i = 0; i < Depth; i++)
 916         DbgPrint(" ");
 917     DbgPrint(" Key #%d", Number);
 918
 919     if (!(Key->IndexEntry->Flags & NTFS_INDEX_ENTRY_END))
 920     {
 921         UNICODE_STRING FileName;
 922         FileName.Length = Key->IndexEntry->FileName.NameLength * sizeof(WCHAR);
 923         FileName.MaximumLength = FileName.Length;
 924         FileName.Buffer = Key->IndexEntry->FileName.Name;
 925         DbgPrint(" '%wZ'\n", &FileName);
 926     }
 927     else
 928     {
 929         DbgPrint(" (Dummy Key)\n");
 930     }
 931
 932     // Is there a child node?
 933     if (Key->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
 934     {
 935         if (Key->LesserChild)
 936             DumpBTreeNode(Key->LesserChild, Number, Depth + 1);
 937         else
 938         {
 939             // This will be an assert once nodes with arbitrary depth are debugged
 940             DPRINT1("DRIVER ERROR: No Key->LesserChild despite Key->IndexEntry->Flags indicating this is a node!\n");
 941         }
 942     }
 943 }
 944
 945 VOID
 946 DumpBTreeNode(PB_TREE_FILENAME_NODE Node, ULONG Number, ULONG Depth)
 947 {
 948     PB_TREE_KEY CurrentKey;
 949     ULONG i;
 950     for (i = 0; i < Depth; i++)
 951         DbgPrint(" ");
 952     DbgPrint("Node #%d, Depth %d, has %d key%s\n", Number, Depth, Node->KeyCount, Node->KeyCount == 1 ? "" : "s");
 953
 954     CurrentKey = Node->FirstKey;
 955     for (i = 1; i <= Node->KeyCount; i++)
 956     {
 957         DumpBTreeKey(CurrentKey, i, Depth);
 958         CurrentKey = CurrentKey->NextKey;
 959     }
 960 }
 961
 962 /**
 963 * @name DumpBTree
 964 * @implemented
 965 *
 966 * Displays a B-Tree.
 967 *
 968 * @param Tree
 969 * Pointer to the B_TREE which will be displayed.
 970 *
 971 * @remarks
 972 * Displays a diagnostic summary of a B_TREE.
 973 */
 974 VOID
 975 DumpBTree(PB_TREE Tree)
 976 {
 977     DbgPrint("B_TREE @ %p\n", Tree);
 978     DumpBTreeNode(Tree->RootNode, 0, 0);
 979 }
 980
 981 // Calculates start of Index Buffer relative to the index allocation, given the node's VCN
 982 ULONGLONG
 983 GetAllocationOffsetFromVCN(PDEVICE_EXTENSION DeviceExt,
 984                            ULONG IndexBufferSize,
 985                            ULONGLONG Vcn)
 986 {
 987     if (IndexBufferSize < DeviceExt->NtfsInfo.BytesPerCluster)
 988         return Vcn * DeviceExt->NtfsInfo.BytesPerSector;
 989
 990     return Vcn * DeviceExt->NtfsInfo.BytesPerCluster;
 991 }
 992
 993 /**
 994 * @name NtfsInsertKey
 995 * @implemented
 996 *
 997 * Inserts a FILENAME_ATTRIBUTE into a B-Tree node.
 998 *
 999 * @param FileReference
1000 * Reference number to the file being added. This will be a combination of the MFT index and update sequence number.
1001 *
1002 * @param FileNameAttribute
1003 * Pointer to a FILENAME_ATTRIBUTE which is the data for the key that will be added to the tree. A copy will be made.
1004 *
1005 * @param Node
1006 * Pointer to a B_TREE_FILENAME_NODE into which a new key will be inserted, in order.
1007 *
1008 * @param CaseSensitive
1009 * Boolean indicating if the function should operate in case-sensitive mode. This will be TRUE
1010 * if an application created the file with the FILE_FLAG_POSIX_SEMANTICS flag.
1011 *
1012 * @remarks
1013 * A node is always sorted, with the least comparable filename stored first and a dummy key to mark the end.
1014 */
1015 NTSTATUS
1016 NtfsInsertKey(ULONGLONG FileReference,
1017               PFILENAME_ATTRIBUTE FileNameAttribute,
1018               PB_TREE_FILENAME_NODE Node,
1019               BOOLEAN CaseSensitive)
1020 {
1021     PB_TREE_KEY NewKey, CurrentKey, PreviousKey;
1022     NTSTATUS Status = STATUS_SUCCESS;
1023     ULONG NodeSize;
1024     ULONG AllocatedNodeSize;
1025     ULONG MaxNodeSizeWithoutHeader;
1026     ULONG i;
1027
1028     DPRINT1("NtfsInsertKey(0x%I64x, %p, %p, %s)\n",
1029             FileReference,
1030             FileNameAttribute,
1031             Node,
1032             CaseSensitive ? "TRUE" : "FALSE");
1033
1034     // Create the key for the filename attribute
1035     NewKey = CreateBTreeKeyFromFilename(FileReference, FileNameAttribute);
1036     if (!NewKey)
1037         return STATUS_INSUFFICIENT_RESOURCES;
1038
1039     // Find where to insert the key
1040     CurrentKey = Node->FirstKey;
1041     PreviousKey = NULL;
1042     for (i = 0; i < Node->KeyCount; i++)
1043     {
1044         // Should the New Key go before the current key?
1045         LONG Comparison = CompareTreeKeys(NewKey, CurrentKey, CaseSensitive);
1046
1047         ASSERT(Comparison != 0);
1048
1049         // Is NewKey < CurrentKey?
1050         if (Comparison < 0)
1051         {
1052
1053             // Does CurrentKey have a sub-node?
1054             if (CurrentKey->LesserChild)
1055             {
1056                 // Insert the key into the child node
1057                 Status = NtfsInsertKey(FileReference, FileNameAttribute, CurrentKey->LesserChild, CaseSensitive);
1058             }
1059             else
1060             {
1061                 // Insert New Key before Current Key
1062                 NewKey->NextKey = CurrentKey;
1063
1064                 // Increase KeyCount and mark node as dirty
1065                 Node->KeyCount++;
1066                 Node->DiskNeedsUpdating = TRUE;
1067
1068                 // was CurrentKey the first key?
1069                 if (CurrentKey == Node->FirstKey)
1070                     Node->FirstKey = NewKey;
1071                 else
1072                     PreviousKey->NextKey = NewKey;
1073                 break;
1074             }
1075         }
1076
1077         PreviousKey = CurrentKey;
1078         CurrentKey = CurrentKey->NextKey;
1079     }
1080
1081     // Is the node larger than its allocated size?
1082     NodeSize = 0;
1083     CurrentKey = Node->FirstKey;
1084     for (i = 0; i < Node->KeyCount; i++)
1085     {
1086         NodeSize += CurrentKey->IndexEntry->Length;
1087         CurrentKey = CurrentKey->NextKey;
1088     }
1089
1090     // TEMPTEMP: TODO: MATH
1091     AllocatedNodeSize = 0xfe8;
1092     MaxNodeSizeWithoutHeader = AllocatedNodeSize - 0x28;
1093
1094     if (NodeSize > MaxNodeSizeWithoutHeader)
1095     {
1096         DPRINT1("FIXME: Splitting a node is still a WIP!\n");
1097         //SplitBTreeNode(NULL, Node);
1098         //DumpBTree(Tree);
1099         return STATUS_NOT_IMPLEMENTED;
1100     }
1101
1102     // NewEntry and NewKey will be destroyed later by DestroyBTree()
1103
1104     return Status;
1105 }