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