Add `NtfsIndex` and `NtfsIndexEntries` to traverse an index in-order.

This introduces a hard dependency on alloc for a core feature.
It is required as long as we don't know the B-tree depth in advance, and need to maintain a stack with elements per layer during traversal.

4 files changed, 178 insertions, 0 deletions

diff --git a/src/error.rs b/src/error.rs
index b7f753a..8015920 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -75,6 +75,8 @@ pub enum NtfsError {
     Io(binread::io::Error),
     /// The Logical Cluster Number (LCN) {lcn} is too big to be processed
     LcnTooBig { lcn: Lcn },
+    /// The index root at byte position {position:#010x} is a large index, but no matching index allocation attribute was provided.
+    MissingIndexAllocation { position: u64 },
     /// The cluster size is {actual} bytes, but the maximum supported one is {expected}
     UnsupportedClusterSize { expected: u32, actual: u32 },
     /// The type of the NTFS attribute at byte position {position:#010x} is {actual:#010x}, which is not supported
diff --git a/src/index.rs b/src/index.rs
new file mode 100644
index 0000000..03c5830
--- /dev/null
+++ b/src/index.rs
@@ -0,0 +1,168 @@
+// Copyright 2021 Colin Finck <colin@reactos.org>
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+use crate::error::{NtfsError, Result};
+use crate::index_entry::{NtfsIndexEntry, NtfsIndexNodeEntries};
+use crate::structured_values::{NewNtfsStructuredValue, NtfsIndexAllocation, NtfsIndexRoot};
+use alloc::vec::Vec;
+use binread::io::{Read, Seek};
+use core::marker::PhantomData;
+
+pub struct NtfsIndex<'n, 'a, K>
+where
+    K: NewNtfsStructuredValue<'n>,
+{
+    index_root: &'a NtfsIndexRoot<'n>,
+    index_allocation: Option<&'a NtfsIndexAllocation<'n>>,
+    key_type: PhantomData<K>,
+}
+
+impl<'n, 'a, K> NtfsIndex<'n, 'a, K>
+where
+    K: NewNtfsStructuredValue<'n>,
+{
+    pub fn new(
+        index_root: &'a NtfsIndexRoot<'n>,
+        index_allocation: Option<&'a NtfsIndexAllocation<'n>>,
+    ) -> Result<Self> {
+        if index_root.is_large_index() && index_allocation.is_none() {
+            return Err(NtfsError::MissingIndexAllocation {
+                position: index_root.position(),
+            });
+        }
+
+        let key_type = PhantomData;
+
+        Ok(Self {
+            index_root,
+            index_allocation,
+            key_type,
+        })
+    }
+
+    pub fn iter<T>(&self, fs: &mut T) -> Result<NtfsIndexEntries<'n, 'a, K>>
+    where
+        K: NewNtfsStructuredValue<'n>,
+        T: Read + Seek,
+    {
+        NtfsIndexEntries::new(fs, self.index_root, self.index_allocation)
+    }
+}
+
+enum StackEntry<'n, K>
+where
+    K: NewNtfsStructuredValue<'n>,
+{
+    EntryToExplore(NtfsIndexEntry<'n, K>),
+    EntryToReturn(NtfsIndexEntry<'n, K>),
+    Iter(NtfsIndexNodeEntries<'n, K>),
+}
+
+/// Iterator over
+///   all index entries of an index,
+///   sorted ascending by the index key,
+///   returning an [`NtfsIndexEntry`] for each entry.
+///
+/// See [`NtfsIndexEntriesAttached`] for an iterator that implements [`Iterator`] and [`FusedIterator`].
+pub struct NtfsIndexEntries<'n, 'a, K>
+where
+    K: NewNtfsStructuredValue<'n>,
+{
+    index_root: &'a NtfsIndexRoot<'n>,
+    index_allocation: Option<&'a NtfsIndexAllocation<'n>>,
+    stack: Vec<StackEntry<'n, K>>,
+}
+
+impl<'n, 'a, K> NtfsIndexEntries<'n, 'a, K>
+where
+    K: NewNtfsStructuredValue<'n>,
+{
+    fn new<T>(
+        fs: &mut T,
+        index_root: &'a NtfsIndexRoot<'n>,
+        index_allocation: Option<&'a NtfsIndexAllocation<'n>>,
+    ) -> Result<Self>
+    where
+        K: NewNtfsStructuredValue<'n>,
+        T: Read + Seek,
+    {
+        // Start with the entries of the top-most node of the B-tree.
+        // This is given by the `NtfsIndexNodeEntries` iterator over the Index Root entries.
+        let stack = vec![StackEntry::Iter(index_root.entries(fs)?)];
+
+        Ok(Self {
+            index_root,
+            index_allocation,
+            stack,
+        })
+    }
+
+    pub fn next<T>(&mut self, fs: &mut T) -> Option<Result<NtfsIndexEntry<'n, K>>>
+    where
+        T: Read + Seek,
+    {
+        // NTFS B-tree indexes are composed out of nodes, with multiple entries per node.
+        // Each entry may have a reference to a subnode.
+        // If that is the case, the subnode comes before the parent node lexicographically.
+        //
+        // An example for an unbalanced, but otherwise valid and sorted tree:
+        //
+        //                                   -------------
+        // INDEX ROOT NODE:                  | 4 | 5 | 6 |
+        //                                   -------------
+        //                                     |
+        //                                 ---------
+        // INDEX ALLOCATION SUBNODE:       | 1 | 3 |
+        //                                 ---------
+        //                                       |
+        //                                     -----
+        // INDEX ALLOCATION SUBNODE:           | 2 |
+        //                                     -----
+        //
+        loop {
+            match self.stack.pop()? {
+                StackEntry::EntryToExplore(entry) => {
+                    // We got an `NtfsIndexEntry` from a previous iteration, which we haven't explored yet.
+                    // Check if it has a subnode that needs to be returned first. In that case, push us on the
+                    // stack to be returned later and push the `NtfsIndexNodeEntries` iterator from the subnode
+                    // to iterate it first.
+                    // If this entry has no subnode, just return and forget about it.
+                    if let Some(subnode_vcn) = entry.subnode_vcn(fs) {
+                        let subnode_vcn = iter_try!(subnode_vcn);
+                        let index_allocation = iter_try!(self.index_allocation.ok_or_else(|| {
+                            NtfsError::MissingIndexAllocation {
+                                position: self.index_root.position(),
+                            }
+                        }));
+                        let subnode = iter_try!(index_allocation.record_from_vcn(
+                            fs,
+                            &self.index_root,
+                            subnode_vcn
+                        ));
+                        let iter = iter_try!(subnode.entries(fs));
+                        self.stack.push(StackEntry::EntryToReturn(entry));
+                        self.stack.push(StackEntry::Iter(iter));
+                    } else {
+                        return Some(Ok(entry));
+                    }
+                }
+                StackEntry::EntryToReturn(entry) => {
+                    // We got an `NtfsIndexEntry` that we have already explored, hence all elements before it
+                    // have already been returned.
+                    // Now it's our turn.
+                    return Some(Ok(entry));
+                }
+                StackEntry::Iter(mut iter) => {
+                    // We got an `NtfsIndexNodeEntries` iterator over the entries of a node.
+                    // Get the next entry from it, and push the updated iterator and the entry back on the stack.
+                    // If this iterator yields no more entries, we are done with this node and can just forget about it.
+                    if let Some(entry) = iter.next(fs) {
+                        let entry = iter_try!(entry);
+                        self.stack.push(StackEntry::Iter(iter));
+                        self.stack.push(StackEntry::EntryToExplore(entry));
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index f086aca..30548c3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -14,6 +14,7 @@ mod attribute_value;
 mod boot_sector;
 mod error;
 mod guid;
+mod index;
 mod index_entry;
 mod index_record;
 mod ntfs;
@@ -29,6 +30,7 @@ pub use crate::attribute::*;
 pub use crate::attribute_value::*;
 pub use crate::error::*;
 pub use crate::guid::*;
+pub use crate::index::*;
 pub use crate::ntfs::*;
 pub use crate::ntfs_file::*;
 pub use crate::string::*;
diff --git a/src/structured_values/index_root.rs b/src/structured_values/index_root.rs
index 0a57631..0355fc0 100644
--- a/src/structured_values/index_root.rs
+++ b/src/structured_values/index_root.rs
@@ -67,6 +67,12 @@ impl<'n> NtfsIndexRoot<'n> {
     pub fn is_large_index(&self) -> bool {
         (self.index_node_header.flags & LARGE_INDEX_FLAG) != 0
     }
+
+    pub fn position(&self) -> u64 {
+        // A structured value is always created from a valid seek position,
+        // and therefore we can safely unwrap here.
+        self.value.data_position().unwrap()
+    }
 }
 
 impl<'n> NewNtfsStructuredValue<'n> for NtfsIndexRoot<'n> {