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