[NTFS] - Add some improvements to B-Tree support. Add preliminary support for trees of arbitrary depth. Add support for creating files in directories which have an index allocation, provided there aren't too many files and the node doesn't need to be split.

-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

@@ -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;
 }

drivers/filesystems/ntfs/mft.c

@@ -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))
     {

drivers/filesystems/ntfs/ntfs.h

@@ -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