[NTFS] - Add some improvements to B-Tree support. Add preliminary support for trees...
authorTrevor Thompson <tmt256@email.vccs.edu>
Thu, 27 Jul 2017 19:52:33 +0000 (19:52 +0000)
committerThomas Faber <thomas.faber@reactos.org>
Sun, 10 Dec 2017 10:15:08 +0000 (11:15 +0100)
-CreateBTreeFromIndex() - Fix memory allocation; allocate sizeof(B_TREE_KEY) bytes for the key, not the size of the pointer. Add support for child nodes. Update parameter list.
-CreateIndexRootFromBTree() - Update Header->Flags if any of index root's keys have sub-nodes.
+CreateIndexBufferFromBTreeNode() - Converts a B-Tree node to an index buffer to be written to the index allocation.
+UpdateIndexAllocation() - Updates all of the stale nodes in an index allocation based on a PB_TREE.
+UpdateIndexNode() - Writes a B-Tree node into the index allocation.
+CreateBTreeKeyFromFilename() - Creates a PB_TREE_KEY based on a FILENAME_ATTRIBUTE.
-DestroyBTreeKey() - Destroy a child node, if present.
-DumpBTreeKey() - Dump a child node, if present.
-DumpBTreeNode() - Include Node->KeyCount in output.
-NtfsAddFilenameToDirectory() - Call UpdateIndexAllocation() prior to CreateIndexRootFromBTree().
-Update B_TREE_FILENAME_NODE with members to keep track of its existence on disk.

svn path=/branches/GSoC_2016/NTFS/; revision=75428

drivers/filesystems/ntfs/btree.c
drivers/filesystems/ntfs/mft.c
drivers/filesystems/ntfs/ntfs.h

index f4bdf90..c1a907e 100644 (file)
@@ -42,12 +42,13 @@ PrintAllVCNs(PDEVICE_EXTENSION Vcb,
     ULONGLONG CurrentOffset = 0;
     PINDEX_BUFFER CurrentNode, Buffer;
     ULONGLONG BufferSize = AttributeDataLength(&IndexAllocationContext->Record);
+    ULONG BytesRead;
     ULONGLONG i;
     int Count = 0;
 
     Buffer = ExAllocatePoolWithTag(NonPagedPool, BufferSize, TAG_NTFS);
 
-    ULONG BytesRead = ReadAttribute(Vcb, IndexAllocationContext, 0, (PCHAR)Buffer, BufferSize);
+    BytesRead = ReadAttribute(Vcb, IndexAllocationContext, 0, (PCHAR)Buffer, BufferSize);
 
     ASSERT(BytesRead = BufferSize);
 
@@ -152,6 +153,175 @@ CompareTreeKeys(PB_TREE_KEY Key1, PB_TREE_KEY Key2, BOOLEAN CaseSensitive)
     return Comparison;
 }
 
+PB_TREE_FILENAME_NODE
+CreateBTreeNodeFromIndexNode(PDEVICE_EXTENSION Vcb,
+                             PINDEX_ROOT_ATTRIBUTE IndexRoot,
+                             PNTFS_ATTR_CONTEXT IndexAllocationAttributeCtx,
+                             PINDEX_ENTRY_ATTRIBUTE NodeEntry)
+{
+    PB_TREE_FILENAME_NODE NewNode;
+    PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
+    PINDEX_ENTRY_ATTRIBUTE FirstNodeEntry;
+    ULONG CurrentEntryOffset = 0;
+    PINDEX_BUFFER NodeBuffer;
+    ULONG IndexBufferSize = Vcb->NtfsInfo.BytesPerIndexRecord;
+    PULONGLONG NodeNumber;
+    PB_TREE_KEY CurrentKey;
+    NTSTATUS Status;
+    ULONGLONG IndexNodeOffset;
+    ULONG BytesRead;
+
+    if (IndexAllocationAttributeCtx == NULL)
+    {
+        DPRINT1("ERROR: Couldn't find index allocation attribute even though there should be one!\n");
+        return NULL;
+    }
+
+    // Get the node number from the end of the node entry
+    NodeNumber = (PULONGLONG)((ULONG_PTR)NodeEntry + NodeEntry->Length - sizeof(ULONGLONG));
+
+    // Create the new tree node
+    DPRINT1("About to allocate %ld for NewNode\n", sizeof(B_TREE_FILENAME_NODE));
+    NewNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
+    if (!NewNode)
+    {
+        DPRINT1("ERROR: Couldn't allocate memory for new filename node.\n");
+        return NULL;
+    }
+    RtlZeroMemory(NewNode, sizeof(B_TREE_FILENAME_NODE));
+
+    // Create the first key
+    CurrentKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
+    if (!CurrentKey)
+    {
+        DPRINT1("ERROR: Failed to allocate memory for key!\n");
+        ExFreePoolWithTag(NewNode, TAG_NTFS);
+        return NULL;
+    }
+    RtlZeroMemory(CurrentKey, sizeof(B_TREE_KEY));
+    NewNode->FirstKey = CurrentKey;
+
+    // Allocate memory for the node buffer
+    NodeBuffer = ExAllocatePoolWithTag(NonPagedPool, IndexBufferSize, TAG_NTFS);
+    if (!NodeBuffer)
+    {
+        DPRINT1("ERROR: Couldn't allocate memory for node buffer!\n");
+        ExFreePoolWithTag(CurrentKey, TAG_NTFS);
+        ExFreePoolWithTag(NewNode, TAG_NTFS);
+        return NULL;
+    }
+
+    // Calculate offset into index allocation
+    IndexNodeOffset = GetAllocationOffsetFromVCN(Vcb, IndexBufferSize, *NodeNumber);
+
+    // TODO: Confirm index bitmap has this node marked as in-use
+
+    // Read the node
+    BytesRead = ReadAttribute(Vcb,
+                              IndexAllocationAttributeCtx,
+                              IndexNodeOffset,
+                              (PCHAR)NodeBuffer,
+                              IndexBufferSize);
+
+    ASSERT(BytesRead == IndexBufferSize);
+    NT_ASSERT(NodeBuffer->Ntfs.Type == NRH_INDX_TYPE);
+    NT_ASSERT(NodeBuffer->VCN == *NodeNumber);
+
+    // Apply the fixup array to the node buffer
+    Status = FixupUpdateSequenceArray(Vcb, &NodeBuffer->Ntfs);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("ERROR: Couldn't apply fixup array to index node buffer!\n");
+        ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
+        ExFreePoolWithTag(CurrentKey, TAG_NTFS);
+        ExFreePoolWithTag(NewNode, TAG_NTFS);
+        return NULL;
+    }
+
+    // Walk through the index and create keys for all the entries
+    FirstNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)(&NodeBuffer->Header)
+                                               + NodeBuffer->Header.FirstEntryOffset);
+    CurrentNodeEntry = FirstNodeEntry;
+    while (CurrentEntryOffset < NodeBuffer->Header.TotalSizeOfEntries)
+    {
+        // Allocate memory for the current entry
+        CurrentKey->IndexEntry = ExAllocatePoolWithTag(NonPagedPool, CurrentNodeEntry->Length, TAG_NTFS);
+        if (!CurrentKey->IndexEntry)
+        {
+            DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
+            DestroyBTreeNode(NewNode);
+            ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
+            return NULL;
+        }
+
+        NewNode->KeyCount++;
+
+        // If this isn't the last entry
+        if (!(CurrentNodeEntry->Flags & NTFS_INDEX_ENTRY_END))
+        {
+            // Create the next key
+            PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
+            if (!NextKey)
+            {
+                DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
+                DestroyBTreeNode(NewNode);
+                ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
+                return NULL;
+            }
+            RtlZeroMemory(NextKey, sizeof(B_TREE_KEY));
+
+            // Add NextKey to the end of the list
+            CurrentKey->NextKey = NextKey;
+
+            // Copy the current entry to its key
+            RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
+
+            // See if the current key has a sub-node
+            if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
+            {
+                DPRINT1("TODO: Only a node with a single-level is supported right now!\n");
+                // Needs debugging:
+                /*CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
+                                                                       IndexRoot,
+                                                                       IndexAllocationAttributeCtx,
+                                                                       CurrentNodeEntry);*/
+            }
+
+            CurrentKey = NextKey;
+        }
+        else
+        {
+            // Copy the final entry to its key
+            RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
+            CurrentKey->NextKey = NULL;
+
+            // See if the current key has a sub-node
+            if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
+            {
+                DPRINT1("TODO: Only a node with a single-level is supported right now!\n");
+                // Needs debugging:
+                /*CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
+                                                                         IndexRoot,
+                                                                         IndexAllocationAttributeCtx,
+                                                                         CurrentNodeEntry);*/
+            }
+
+            break;
+        }
+
+        // Advance to the next entry
+        CurrentEntryOffset += CurrentNodeEntry->Length;
+        CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
+    }
+
+    NewNode->NodeNumber = *NodeNumber;
+    NewNode->ExistsOnDisk = TRUE;
+
+    ExFreePoolWithTag(NodeBuffer, TAG_NTFS);
+
+    return NewNode;
+}
+
 /**
 * @name CreateBTreeFromIndex
 * @implemented
@@ -172,7 +342,10 @@ CompareTreeKeys(PB_TREE_KEY Key1, PB_TREE_KEY Key2, BOOLEAN CaseSensitive)
 * Allocates memory for the entire tree. Caller is responsible for destroying the tree with DestroyBTree().
 */
 NTSTATUS
-CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
+CreateBTreeFromIndex(PDEVICE_EXTENSION Vcb,
+                     PFILE_RECORD_HEADER FileRecordWithIndex,
+                     /*PCWSTR IndexName,*/
+                     PNTFS_ATTR_CONTEXT IndexRootContext,
                      PINDEX_ROOT_ATTRIBUTE IndexRoot,
                      PB_TREE *NewTree)
 {
@@ -181,8 +354,10 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
     PB_TREE_FILENAME_NODE RootNode = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_FILENAME_NODE), TAG_NTFS);
     PB_TREE_KEY CurrentKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
     ULONG CurrentOffset = IndexRoot->Header.FirstEntryOffset;
+    PNTFS_ATTR_CONTEXT IndexAllocationContext = NULL;
+    NTSTATUS Status;
 
-    DPRINT1("CreateBTreeFromIndex(%p, %p, %p)\n", IndexRootContext, IndexRoot, NewTree);
+    DPRINT1("CreateBTreeFromIndex(%p, %p)\n", IndexRoot, NewTree);
 
     if (!Tree || !RootNode || !CurrentKey)
     {
@@ -200,6 +375,19 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
     RtlZeroMemory(RootNode, sizeof(B_TREE_FILENAME_NODE));
     RtlZeroMemory(CurrentKey, sizeof(B_TREE_KEY));
 
+    // See if the file record has an attribute allocation
+    Status = FindAttribute(Vcb,
+                           FileRecordWithIndex,
+                           AttributeIndexAllocation,
+                           L"$I30",
+                           4,
+                           &IndexAllocationContext,
+                           NULL);
+    if (!NT_SUCCESS(Status))
+        IndexAllocationContext = NULL;
+    else
+        PrintAllVCNs(Vcb, IndexAllocationContext, IndexRoot->SizeOfEntry);
+
     // Setup the Tree
     RootNode->FirstKey = CurrentKey;
     Tree->RootNode = RootNode;
@@ -235,7 +423,7 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
         if (!(CurrentNodeEntry->Flags & NTFS_INDEX_ENTRY_END))
         {
             // Create the next key
-            PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(PB_TREE_KEY), TAG_NTFS);
+            PB_TREE_KEY NextKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
             if (!NextKey)
             {
                 DPRINT1("ERROR: Couldn't allocate memory for next key!\n");
@@ -243,7 +431,7 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
                 return STATUS_INSUFFICIENT_RESOURCES;
             }
 
-            RtlZeroMemory(NextKey, sizeof(PB_TREE_KEY));
+            RtlZeroMemory(NextKey, sizeof(B_TREE_KEY));
 
             // Add NextKey to the end of the list
             CurrentKey->NextKey = NextKey;
@@ -251,12 +439,20 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
             // Copy the current entry to its key
             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
 
-            // Make sure this B-Tree is only one level deep (flat list)
+            // Does this key have a sub-node?
             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
             {
-                DPRINT1("TODO: Only directories with single-level B-Trees are supported right now!\n");
-                DestroyBTree(Tree);
-                return STATUS_NOT_IMPLEMENTED;
+                // Create the child node
+                CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
+                                                                       IndexRoot,
+                                                                       IndexAllocationContext,
+                                                                       CurrentKey->IndexEntry);
+                if (!CurrentKey->LesserChild)
+                {
+                    DPRINT1("ERROR: Couldn't create child node!\n");
+                    DestroyBTree(Tree);
+                    return STATUS_NOT_IMPLEMENTED;
+                }
             }
 
             // Advance to the next entry
@@ -270,12 +466,20 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
             RtlCopyMemory(CurrentKey->IndexEntry, CurrentNodeEntry, CurrentNodeEntry->Length);
             CurrentKey->NextKey = NULL;
 
-            // Make sure this B-Tree is only one level deep (flat list)
+            // Does this key have a sub-node?
             if (CurrentKey->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
             {
-                DPRINT1("TODO: Only directories with single-level B-Trees are supported right now!\n");
-                DestroyBTree(Tree);
-                return STATUS_NOT_IMPLEMENTED;
+                // Create the child node
+                CurrentKey->LesserChild = CreateBTreeNodeFromIndexNode(Vcb,
+                                                                       IndexRoot,
+                                                                       IndexAllocationContext,
+                                                                       CurrentKey->IndexEntry);
+                if (!CurrentKey->LesserChild)
+                {
+                    DPRINT1("ERROR: Couldn't create child node!\n");
+                    DestroyBTree(Tree);
+                    return STATUS_NOT_IMPLEMENTED;
+                }
             }
 
             break;
@@ -284,6 +488,9 @@ CreateBTreeFromIndex(PNTFS_ATTR_CONTEXT IndexRootContext,
 
     *NewTree = Tree;
 
+    if (IndexAllocationContext)
+        ReleaseAttributeContext(IndexAllocationContext);
+
     return STATUS_SUCCESS;
 }
 
@@ -387,6 +594,10 @@ CreateIndexRootFromBTree(PDEVICE_EXTENSION DeviceExt,
                 CurrentNodeEntry->KeyLength,
                 CurrentNodeEntry->Length);
 
+        // Does the current key have any sub-nodes?
+        if (CurrentKey->LesserChild)
+            NewIndexRoot->Header.Flags = INDEX_ROOT_LARGE;
+
         // Add Length of Current Entry to Total Size of Entries
         NewIndexRoot->Header.TotalSizeOfEntries += CurrentNodeEntry->Length;
 
@@ -404,13 +615,253 @@ CreateIndexRootFromBTree(PDEVICE_EXTENSION DeviceExt,
     return STATUS_SUCCESS;
 }
 
+NTSTATUS
+CreateIndexBufferFromBTreeNode(PDEVICE_EXTENSION DeviceExt,
+                               PB_TREE_FILENAME_NODE Node,
+                               ULONG BufferSize,
+                               PINDEX_BUFFER IndexBuffer)
+{
+    ULONG i;
+    PB_TREE_KEY CurrentKey;
+    PINDEX_ENTRY_ATTRIBUTE CurrentNodeEntry;
+    NTSTATUS Status;
+
+    // TODO: Fix magic, do math
+    RtlZeroMemory(IndexBuffer, BufferSize);
+    IndexBuffer->Ntfs.Type = NRH_INDX_TYPE;
+    IndexBuffer->Ntfs.UsaOffset = 0x28;
+    IndexBuffer->Ntfs.UsaCount = 9;
+
+    // TODO: Check bitmap for VCN
+    ASSERT(Node->ExistsOnDisk);
+    IndexBuffer->VCN = Node->NodeNumber;
+
+    IndexBuffer->Header.FirstEntryOffset = 0x28;
+    IndexBuffer->Header.AllocatedSize = BufferSize - FIELD_OFFSET(INDEX_BUFFER, Header);
+
+    // Start summing the total size of this node's entries
+    IndexBuffer->Header.TotalSizeOfEntries = IndexBuffer->Header.FirstEntryOffset;
+
+    CurrentKey = Node->FirstKey;
+    CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)&(IndexBuffer->Header)
+                                                + IndexBuffer->Header.FirstEntryOffset);
+    for (i = 0; i < Node->KeyCount; i++)
+    {
+        // Would adding the current entry to the index increase the node size beyond the allocation size?
+        ULONG IndexSize = FIELD_OFFSET(INDEX_BUFFER, Header)
+            + IndexBuffer->Header.FirstEntryOffset
+            + IndexBuffer->Header.TotalSizeOfEntries
+            + CurrentNodeEntry->Length;
+        if (IndexSize > BufferSize)
+        {
+            DPRINT1("TODO: Adding file would require creating a new node!\n");
+            return STATUS_NOT_IMPLEMENTED;
+        }
+
+        ASSERT(CurrentKey->IndexEntry->Length != 0);
+
+        // Copy the index entry
+        RtlCopyMemory(CurrentNodeEntry, CurrentKey->IndexEntry, CurrentKey->IndexEntry->Length);
+
+        DPRINT1("Index Node Entry Stream Length: %u\nIndex Node Entry Length: %u\n",
+                CurrentNodeEntry->KeyLength,
+                CurrentNodeEntry->Length);
+
+        // Add Length of Current Entry to Total Size of Entries
+        IndexBuffer->Header.TotalSizeOfEntries += CurrentNodeEntry->Length;
+
+        // TODO: Check for child nodes
+
+        // Go to the next node entry
+        CurrentNodeEntry = (PINDEX_ENTRY_ATTRIBUTE)((ULONG_PTR)CurrentNodeEntry + CurrentNodeEntry->Length);
+        CurrentKey = CurrentKey->NextKey;
+    }
+
+    Status = AddFixupArray(DeviceExt, &IndexBuffer->Ntfs);
+
+    return Status;
+}
+
+NTSTATUS
+UpdateIndexAllocation(PDEVICE_EXTENSION DeviceExt,
+                      PB_TREE Tree,
+                      ULONG IndexBufferSize,
+                      PFILE_RECORD_HEADER FileRecord)
+{
+    // Find the index allocation and bitmap
+    PNTFS_ATTR_CONTEXT IndexAllocationContext, BitmapContext;
+    PB_TREE_KEY CurrentKey;
+    NTSTATUS Status;
+    BOOLEAN HasIndexAllocation = FALSE;
+    ULONG i;
+
+    DPRINT1("UpdateIndexAllocations() called.\n");
+
+    Status = FindAttribute(DeviceExt, FileRecord, AttributeIndexAllocation, L"$I30", 4, &IndexAllocationContext, NULL);
+    if (NT_SUCCESS(Status))
+        HasIndexAllocation = TRUE;
+
+    // TODO: Handle bitmap
+    BitmapContext = NULL;
+
+    // Walk through the root node and update all the sub-nodes
+    CurrentKey = Tree->RootNode->FirstKey;
+    for (i = 0; i < Tree->RootNode->KeyCount; i++)
+    {
+        if (CurrentKey->LesserChild)
+        {
+            if (!HasIndexAllocation)
+            {
+                DPRINT1("FIXME: Need to add index allocation\n");
+                return STATUS_NOT_IMPLEMENTED;
+            }
+            else
+            {
+                Status = UpdateIndexNode(DeviceExt, CurrentKey->LesserChild, IndexBufferSize, IndexAllocationContext, BitmapContext);
+                if (!NT_SUCCESS(Status))
+                {
+                    DPRINT1("ERROR: Failed to update index node!\n");
+                    ReleaseAttributeContext(IndexAllocationContext);
+                    return Status;
+                }
+            }
+
+        }
+        CurrentKey = CurrentKey->NextKey;
+    }
+
+    if(HasIndexAllocation)
+        ReleaseAttributeContext(IndexAllocationContext);
+
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS
+UpdateIndexNode(PDEVICE_EXTENSION DeviceExt,
+                PB_TREE_FILENAME_NODE Node,
+                ULONG IndexBufferSize,
+                PNTFS_ATTR_CONTEXT IndexAllocationContext,
+                PNTFS_ATTR_CONTEXT BitmapContext)
+{
+    ULONG i;
+    PB_TREE_KEY CurrentKey = Node->FirstKey;
+    NTSTATUS Status;
+
+    DPRINT1("UpdateIndexNode(%p, %p, %lu, %p, %p) called for index node with VCN %I64u\n", DeviceExt, Node, IndexBufferSize, IndexAllocationContext, BitmapContext, Node->NodeNumber);
+
+    // Do we need to write this node to disk?
+    if (Node->DiskNeedsUpdating)
+    {
+        ULONGLONG NodeOffset;
+        ULONG LengthWritten;
+
+        // Allocate memory for an index buffer
+        PINDEX_BUFFER IndexBuffer = ExAllocatePoolWithTag(NonPagedPool, IndexBufferSize, TAG_NTFS);
+        if (!IndexBuffer)
+        {
+            DPRINT1("ERROR: Failed to allocate %lu bytes for index buffer!\n", IndexBufferSize);
+            return STATUS_INSUFFICIENT_RESOURCES;
+        }
+
+        // Create the index buffer we'll be writing to disk to represent this node
+        Status = CreateIndexBufferFromBTreeNode(DeviceExt, Node, IndexBufferSize, IndexBuffer);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("ERROR: Failed to create index buffer from node!\n");
+            ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
+            return Status;
+        }
+
+        // Get Offset of index buffer in index allocation
+        NodeOffset = GetAllocationOffsetFromVCN(DeviceExt, IndexBufferSize, Node->NodeNumber);
+
+        // Write the buffer to the index allocation
+        Status = WriteAttribute(DeviceExt, IndexAllocationContext, NodeOffset, (const PUCHAR)IndexBuffer, IndexBufferSize, &LengthWritten);
+        if (!NT_SUCCESS(Status) || LengthWritten != IndexBufferSize)
+        {
+            DPRINT1("ERROR: Failed to update index allocation!\n");
+            ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
+            if (!NT_SUCCESS(Status))
+                return Status;
+            else
+                return STATUS_END_OF_FILE;
+        }
+
+        Node->DiskNeedsUpdating = FALSE;
+
+        // Free the index buffer
+        ExFreePoolWithTag(IndexBuffer, TAG_NTFS);
+    }
+
+    // Walk through the node and look for children to update
+    for (i = 0; i < Node->KeyCount; i++)
+    {
+        ASSERT(CurrentKey);
+
+        // If there's a child node
+        if (CurrentKey->LesserChild)
+        {
+            // Update the child node on disk
+            Status = UpdateIndexNode(DeviceExt, CurrentKey->LesserChild, IndexBufferSize, IndexAllocationContext, BitmapContext);
+            if (!NT_SUCCESS(Status))
+            {
+                DPRINT1("ERROR: Failed to update child node!\n");
+                return Status;
+            }
+        }
+
+        CurrentKey = CurrentKey->NextKey;
+    }
+
+    return STATUS_SUCCESS;
+}
+
+PB_TREE_KEY
+CreateBTreeKeyFromFilename(ULONGLONG FileReference, PFILENAME_ATTRIBUTE FileNameAttribute)
+{
+    PB_TREE_KEY NewKey;
+    ULONG AttributeSize = GetFileNameAttributeLength(FileNameAttribute);
+    ULONG EntrySize = ALIGN_UP_BY(AttributeSize + FIELD_OFFSET(INDEX_ENTRY_ATTRIBUTE, FileName), 8);
+
+    // Create a new Index Entry for the file
+    PINDEX_ENTRY_ATTRIBUTE NewEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_NTFS);
+    if (!NewEntry)
+    {
+        DPRINT1("ERROR: Failed to allocate memory for Index Entry!\n");
+        return NULL;
+    }
+
+    // Setup the Index Entry
+    RtlZeroMemory(NewEntry, EntrySize);
+    NewEntry->Data.Directory.IndexedFile = FileReference;
+    NewEntry->Length = EntrySize;
+    NewEntry->KeyLength = AttributeSize;
+
+    // Copy the FileNameAttribute
+    RtlCopyMemory(&NewEntry->FileName, FileNameAttribute, AttributeSize);
+
+    // Setup the New Key
+    NewKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
+    if (!NewKey)
+    {
+        DPRINT1("ERROR: Failed to allocate memory for new key!\n");
+        ExFreePoolWithTag(NewEntry, TAG_NTFS);
+        return NULL;
+    }
+    NewKey->IndexEntry = NewEntry;
+    NewKey->NextKey = NULL;
+
+    return NewKey;
+}
+
 VOID
 DestroyBTreeKey(PB_TREE_KEY Key)
 {
     if (Key->IndexEntry)
         ExFreePoolWithTag(Key->IndexEntry, TAG_NTFS);
 
-    // We'll destroy Key->LesserChild here after we start using it
+    if (Key->LesserChild)
+        DestroyBTreeNode(Key->LesserChild);
 
     ExFreePoolWithTag(Key, TAG_NTFS);
 }
@@ -473,6 +924,18 @@ DumpBTreeKey(PB_TREE_KEY Key, ULONG Number, ULONG Depth)
     {
         DbgPrint(" (Dummy Key)\n");
     }
+
+    // Is there a child node?
+    if (Key->IndexEntry->Flags & NTFS_INDEX_ENTRY_NODE)
+    {
+        if (Key->LesserChild)
+            DumpBTreeNode(Key->LesserChild, Number, Depth + 1);
+        else
+        {
+            // This will be an assert once nodes with arbitrary depth are debugged
+            DPRINT1("DRIVER ERROR: No Key->LesserChild despite Key->IndexEntry->Flags indicating this is a node!\n");
+        }
+    }
 }
 
 VOID
@@ -482,7 +945,7 @@ DumpBTreeNode(PB_TREE_FILENAME_NODE Node, ULONG Number, ULONG Depth)
     ULONG i;
     for (i = 0; i < Depth; i++)
         DbgPrint(" ");
-    DbgPrint("Node #%d, Depth %d\n", Number, Depth);
+    DbgPrint("Node #%d, Depth %d, has %d keys\n", Number, Depth, Node->KeyCount);
 
     CurrentKey = Node->FirstKey;
     for (i = 0; i < Node->KeyCount; i++)
@@ -551,46 +1014,23 @@ NtfsInsertKey(ULONGLONG FileReference,
               PB_TREE_FILENAME_NODE Node,
               BOOLEAN CaseSensitive)
 {
-    // Calculate size of Attribute and Index Entry
-    ULONG AttributeSize = GetFileNameAttributeLength(FileNameAttribute);
-    ULONG EntrySize = ALIGN_UP_BY(AttributeSize + FIELD_OFFSET(INDEX_ENTRY_ATTRIBUTE, FileName), 8);
-    PINDEX_ENTRY_ATTRIBUTE NewEntry;
     PB_TREE_KEY NewKey, CurrentKey, PreviousKey;
+    NTSTATUS Status = STATUS_SUCCESS; 
+    ULONG NodeSize;
+    ULONG AllocatedNodeSize;
+    ULONG MaxNodeSizeWithoutHeader;
     ULONG i;
 
-    DPRINT1("NtfsInsertKey(0x%02I64, %p, %p, %s)\n",
+    DPRINT1("NtfsInsertKey(0x%I64x, %p, %p, %s)\n",
             FileReference,
             FileNameAttribute,
             Node,
             CaseSensitive ? "TRUE" : "FALSE");
 
-    // Create a new Index Entry for the file
-    NewEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_NTFS);
-    if (!NewEntry)
-    {
-        DPRINT1("ERROR: Failed to allocate memory for Index Entry!\n");
-        return STATUS_INSUFFICIENT_RESOURCES;
-    }
-
-    // Setup the Index Entry
-    RtlZeroMemory(NewEntry, EntrySize);
-    NewEntry->Data.Directory.IndexedFile = FileReference;
-    NewEntry->Length = EntrySize;
-    NewEntry->KeyLength = AttributeSize;
-
-    // Copy the FileNameAttribute
-    RtlCopyMemory(&NewEntry->FileName, FileNameAttribute, AttributeSize);
-
-    // Setup the New Key
-    NewKey = ExAllocatePoolWithTag(NonPagedPool, sizeof(B_TREE_KEY), TAG_NTFS);
+    // Create the key for the filename attribute
+    NewKey = CreateBTreeKeyFromFilename(FileReference, FileNameAttribute);
     if (!NewKey)
-    {
-        DPRINT1("ERROR: Failed to allocate memory for new key!\n");
-        ExFreePoolWithTag(NewEntry, TAG_NTFS);
         return STATUS_INSUFFICIENT_RESOURCES;
-    }
-    NewKey->IndexEntry = NewEntry;
-    NewKey->NextKey = NULL;
 
     // Find where to insert the key
     CurrentKey = Node->FirstKey;
@@ -605,24 +1045,57 @@ NtfsInsertKey(ULONGLONG FileReference,
         // Is NewKey < CurrentKey?
         if (Comparison < 0)
         {
-            // Insert New Key before Current Key
-            NewKey->NextKey = CurrentKey;
 
-            // was CurrentKey the first key?
-            if (CurrentKey == Node->FirstKey)
-                Node->FirstKey = NewKey;
+            // Does CurrentKey have a sub-node?
+            if (CurrentKey->LesserChild)
+            {
+                // Insert the key into the child node
+                Status = NtfsInsertKey(FileReference, FileNameAttribute, CurrentKey->LesserChild, CaseSensitive);
+            }
             else
-                PreviousKey->NextKey = NewKey;
-            break;
+            {
+                // Insert New Key before Current Key
+                NewKey->NextKey = CurrentKey;
+
+                // Increase KeyCount and mark node as dirty
+                Node->KeyCount++;
+                Node->DiskNeedsUpdating = TRUE;
+
+                // was CurrentKey the first key?
+                if (CurrentKey == Node->FirstKey)
+                    Node->FirstKey = NewKey;
+                else
+                    PreviousKey->NextKey = NewKey;
+                break;
+            }
         }
 
         PreviousKey = CurrentKey;
         CurrentKey = CurrentKey->NextKey;
     }
 
-    Node->KeyCount++;
+    // Is the node larger than its allocated size?
+    NodeSize = 0; 
+    CurrentKey = Node->FirstKey;
+    for (i = 0; i < Node->KeyCount; i++)
+    {
+        NodeSize += CurrentKey->IndexEntry->Length;
+        CurrentKey = CurrentKey->NextKey;
+    }
+
+    // TEMPTEMP: TODO: MATH
+    AllocatedNodeSize = 0xfe8;
+    MaxNodeSizeWithoutHeader = AllocatedNodeSize - 0x28;
+
+    if (NodeSize > MaxNodeSizeWithoutHeader)
+    {
+        DPRINT1("FIXME: Splitting a node is still a WIP!\n");
+        //SplitBTreeNode(NULL, Node);
+        //DumpBTree(Tree);
+        return STATUS_NOT_IMPLEMENTED;
+    }
 
     // NewEntry and NewKey will be destroyed later by DestroyBTree()
 
-    return STATUS_SUCCESS;
+    return Status;
 }
\ No newline at end of file
index 928839e..1df31ae 100644 (file)
@@ -1969,7 +1969,11 @@ NtfsAddFilenameToDirectory(PDEVICE_EXTENSION DeviceExt,
     }
 
     // Convert the index to a B*Tree
-    Status = CreateBTreeFromIndex(IndexRootContext, I30IndexRoot, &NewTree);
+    Status = CreateBTreeFromIndex(DeviceExt,
+                                  ParentFileRecord,
+                                  IndexRootContext,
+                                  I30IndexRoot,
+                                  &NewTree);
     if (!NT_SUCCESS(Status))
     {
         DPRINT1("ERROR: Failed to create B-Tree from Index!\n");
@@ -1995,7 +1999,19 @@ NtfsAddFilenameToDirectory(PDEVICE_EXTENSION DeviceExt,
 
     DumpBTree(NewTree);
 
-    // Convert B*Tree back to Index Root
+    // Convert B*Tree back to Index
+    Status = UpdateIndexAllocation(DeviceExt, NewTree, I30IndexRoot->SizeOfEntry, ParentFileRecord);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("ERROR: Failed to update index allocation from B-Tree!\n");
+        DestroyBTree(NewTree);
+        ReleaseAttributeContext(IndexRootContext);
+        ExFreePoolWithTag(I30IndexRoot, TAG_NTFS);
+        ExFreePoolWithTag(ParentFileRecord, TAG_NTFS);
+        return Status;
+    }
+
+    // Create the Index Root from the B*Tree
     Status = CreateIndexRootFromBTree(DeviceExt, NewTree, MaxIndexSize, &NewIndexRoot, &BtreeIndexLength);
     if (!NT_SUCCESS(Status))
     {
index 39c3321..42168be 100644 (file)
@@ -406,20 +406,26 @@ typedef struct
     FILENAME_ATTRIBUTE    FileName;
 } INDEX_ENTRY_ATTRIBUTE, *PINDEX_ENTRY_ATTRIBUTE;
 
+struct _B_TREE_FILENAME_NODE;
+typedef struct _B_TREE_FILENAME_NODE B_TREE_FILENAME_NODE;
+
 // Keys are arranged in nodes as an ordered, linked list
 typedef struct _B_TREE_KEY
 {
     struct _B_TREE_KEY *NextKey;
-    // PB_TREE_FILENAME_NODE LesserChild; // we aren't worried about multi-level trees yet
-    PINDEX_ENTRY_ATTRIBUTE IndexEntry;   // must be last member for FIELD_OFFSET
+    B_TREE_FILENAME_NODE *LesserChild;  // Child-Node. All the keys in this node will be sorted before IndexEntry
+    PINDEX_ENTRY_ATTRIBUTE IndexEntry;  // must be last member for FIELD_OFFSET
 }B_TREE_KEY, *PB_TREE_KEY;
 
 // Every Node is just an ordered list of keys.
 // Sub-nodes can be found attached to a key (if they exist).
 // A key's sub-node precedes that key in the ordered list.
-typedef struct
+typedef struct _B_TREE_FILENAME_NODE
 {
     ULONG KeyCount;
+    BOOLEAN ExistsOnDisk;
+    BOOLEAN DiskNeedsUpdating;
+    ULONGLONG NodeNumber;
     PB_TREE_KEY FirstKey;
 } B_TREE_FILENAME_NODE, *PB_TREE_FILENAME_NODE;
 
@@ -688,7 +694,9 @@ CompareTreeKeys(PB_TREE_KEY Key1,
                 BOOLEAN CaseSensitive);
 
 NTSTATUS
-CreateBTreeFromIndex(/*PDEVICE_EXTENSION Vcb,*/
+CreateBTreeFromIndex(PDEVICE_EXTENSION Vcb,
+                     PFILE_RECORD_HEADER FileRecordWithIndex,
+                     /*PCWSTR IndexName,*/
                      PNTFS_ATTR_CONTEXT IndexRootContext,
                      PINDEX_ROOT_ATTRIBUTE IndexRoot,
                      PB_TREE *NewTree);
@@ -725,6 +733,19 @@ NtfsInsertKey(ULONGLONG FileReference,
               PB_TREE_FILENAME_NODE Node,
               BOOLEAN CaseSensitive);
 
+NTSTATUS
+UpdateIndexAllocation(PDEVICE_EXTENSION DeviceExt,
+                      PB_TREE Tree,
+                      ULONG IndexBufferSize,
+                      PFILE_RECORD_HEADER FileRecord);
+
+NTSTATUS
+UpdateIndexNode(PDEVICE_EXTENSION DeviceExt,
+                PB_TREE_FILENAME_NODE Node,
+                ULONG IndexBufferSize,
+                PNTFS_ATTR_CONTEXT IndexAllocationContext,
+                PNTFS_ATTR_CONTEXT BitmapContext);
+
 /* close.c */
 
 NTSTATUS