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