Introduce typed indexes to support more than just file name indexes.

Index entry keys do not necessarily correspond to the structured values. In fact, file name indexes are the exception and not the rule. All other index types have dedicated key structures.
Likewise, some index types (e.g. SDH and SII) also have extra data structures instead of a file reference.
Finally, some other index types (like R over Reparse points) only have a key structure, but neither data nor a file reference.

8 files changed, 241 insertions, 48 deletions

diff --git a/src/index.rs b/src/index.rs
index 0af1a8b..21d5b7f 100644
--- a/src/index.rs
+++ b/src/index.rs
@@ -3,16 +3,25 @@
 
 use crate::error::{NtfsError, Result};
 use crate::index_entry::{IndexEntryRange, IndexNodeEntryRanges, NtfsIndexEntry};
+use crate::indexes::NtfsIndexEntryType;
 use crate::structured_values::{NtfsIndexAllocation, NtfsIndexRoot};
 use alloc::vec::Vec;
 use binread::io::{Read, Seek};
+use core::marker::PhantomData;
 
-pub struct NtfsIndex<'n, 'f> {
+pub struct NtfsIndex<'n, 'f, E>
+where
+    E: NtfsIndexEntryType,
+{
     index_root: NtfsIndexRoot<'f>,
     index_allocation: Option<NtfsIndexAllocation<'n, 'f>>,
+    entry_type: PhantomData<E>,
 }
 
-impl<'n, 'f> NtfsIndex<'n, 'f> {
+impl<'n, 'f, E> NtfsIndex<'n, 'f, E>
+where
+    E: NtfsIndexEntryType,
+{
     pub fn new(
         index_root: NtfsIndexRoot<'f>,
         index_allocation: Option<NtfsIndexAllocation<'n, 'f>>,
@@ -23,13 +32,16 @@ impl<'n, 'f> NtfsIndex<'n, 'f> {
             });
         }
 
+        let entry_type = PhantomData;
+
         Ok(Self {
             index_root,
             index_allocation,
+            entry_type,
         })
     }
 
-    pub fn iter<'i>(&'i self) -> Result<NtfsIndexEntries<'n, 'f, 'i>> {
+    pub fn iter<'i>(&'i self) -> Result<NtfsIndexEntries<'n, 'f, 'i, E>> {
         NtfsIndexEntries::new(self)
     }
 }
@@ -40,14 +52,20 @@ impl<'n, 'f> NtfsIndex<'n, 'f> {
 ///   returning an [`NtfsIndexEntry`] for each entry.
 ///
 /// See [`NtfsIndexEntriesAttached`] for an iterator that implements [`Iterator`] and [`FusedIterator`].
-pub struct NtfsIndexEntries<'n, 'f, 'i> {
-    index: &'i NtfsIndex<'n, 'f>,
-    inner_iterators: Vec<IndexNodeEntryRanges>,
-    following_entries: Vec<IndexEntryRange>,
+pub struct NtfsIndexEntries<'n, 'f, 'i, E>
+where
+    E: NtfsIndexEntryType,
+{
+    index: &'i NtfsIndex<'n, 'f, E>,
+    inner_iterators: Vec<IndexNodeEntryRanges<E>>,
+    following_entries: Vec<IndexEntryRange<E>>,
 }
 
-impl<'n, 'f, 'i> NtfsIndexEntries<'n, 'f, 'i> {
-    fn new(index: &'i NtfsIndex<'n, 'f>) -> Result<Self> {
+impl<'n, 'f, 'i, E> NtfsIndexEntries<'n, 'f, 'i, E>
+where
+    E: NtfsIndexEntryType,
+{
+    fn new(index: &'i NtfsIndex<'n, 'f, E>) -> Result<Self> {
         let inner_iterators = vec![index.index_root.entry_ranges()];
         let following_entries = Vec::new();
 
@@ -58,7 +76,7 @@ impl<'n, 'f, 'i> NtfsIndexEntries<'n, 'f, 'i> {
         })
     }
 
-    pub fn next<'a, T>(&'a mut self, fs: &mut T) -> Option<Result<NtfsIndexEntry<'a>>>
+    pub fn next<'a, T>(&'a mut self, fs: &mut T) -> Option<Result<NtfsIndexEntry<'a, E>>>
     where
         T: Read + Seek,
     {
diff --git a/src/index_entry.rs b/src/index_entry.rs
index a3e54f4..1c84768 100644
--- a/src/index_entry.rs
+++ b/src/index_entry.rs
@@ -2,11 +2,16 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 use crate::error::Result;
-use crate::structured_values::NtfsStructuredValueFromSlice;
+use crate::file_reference::NtfsFileReference;
+use crate::indexes::{
+    NtfsIndexEntryData, NtfsIndexEntryHasData, NtfsIndexEntryHasFileReference, NtfsIndexEntryKey,
+    NtfsIndexEntryType,
+};
 use crate::types::Vcn;
 use bitflags::bitflags;
 use byteorder::{ByteOrder, LittleEndian};
 use core::iter::FusedIterator;
+use core::marker::PhantomData;
 use core::mem;
 use core::ops::Range;
 use memoffset::offset_of;
@@ -16,7 +21,14 @@ const INDEX_ENTRY_HEADER_SIZE: i64 = 16;
 
 #[repr(C, packed)]
 struct IndexEntryHeader {
-    file_ref: u64,
+    // The following three fields are used for the u64 file reference if the entry type
+    // has no data, but a file reference instead.
+    // This is indicated by the entry type implementing `NtfsIndexEntryHasFileReference`.
+    // Currently, only `NtfsFileNameIndex` has such a file reference.
+    data_offset: u16,
+    data_length: u16,
+    padding: u32,
+
     index_entry_length: u16,
     key_length: u16,
     flags: u8,
@@ -31,30 +43,99 @@ bitflags! {
     }
 }
 
-pub(crate) struct IndexEntryRange {
+pub(crate) struct IndexEntryRange<E>
+where
+    E: NtfsIndexEntryType,
+{
     range: Range<usize>,
     position: u64,
+    entry_type: PhantomData<E>,
 }
 
-impl IndexEntryRange {
-    pub(crate) const fn new(range: Range<usize>, position: u64) -> Self {
-        Self { range, position }
+impl<E> IndexEntryRange<E>
+where
+    E: NtfsIndexEntryType,
+{
+    pub(crate) fn new(range: Range<usize>, position: u64) -> Self {
+        let entry_type = PhantomData;
+        Self {
+            range,
+            position,
+            entry_type,
+        }
     }
 
-    pub(crate) fn to_entry<'s>(&self, slice: &'s [u8]) -> NtfsIndexEntry<'s> {
+    pub(crate) fn to_entry<'s>(&self, slice: &'s [u8]) -> NtfsIndexEntry<'s, E> {
         NtfsIndexEntry::new(&slice[self.range.clone()], self.position)
     }
 }
 
 #[derive(Clone, Debug)]
-pub struct NtfsIndexEntry<'s> {
+pub struct NtfsIndexEntry<'s, E>
+where
+    E: NtfsIndexEntryType,
+{
     slice: &'s [u8],
     position: u64,
+    entry_type: PhantomData<E>,
 }
 
-impl<'s> NtfsIndexEntry<'s> {
-    pub(crate) const fn new(slice: &'s [u8], position: u64) -> Self {
-        Self { slice, position }
+impl<'s, E> NtfsIndexEntry<'s, E>
+where
+    E: NtfsIndexEntryType,
+{
+    pub(crate) fn new(slice: &'s [u8], position: u64) -> Self {
+        let entry_type = PhantomData;
+        Self {
+            slice,
+            position,
+            entry_type,
+        }
+    }
+
+    pub fn data(&self) -> Option<Result<E::DataType>>
+    where
+        E: NtfsIndexEntryHasData,
+    {
+        if self.data_offset() == 0 || self.data_length() == 0 {
+            return None;
+        }
+
+        let start = self.data_offset() as usize;
+        let end = start + self.data_length() as usize;
+        let position = self.position + start as u64;
+
+        let data = iter_try!(E::DataType::data_from_slice(
+            &self.slice[start..end],
+            position
+        ));
+        Some(Ok(data))
+    }
+
+    fn data_offset(&self) -> u16
+    where
+        E: NtfsIndexEntryHasData,
+    {
+        let start = offset_of!(IndexEntryHeader, data_offset);
+        LittleEndian::read_u16(&self.slice[start..])
+    }
+
+    pub fn data_length(&self) -> u16
+    where
+        E: NtfsIndexEntryHasData,
+    {
+        let start = offset_of!(IndexEntryHeader, data_length);
+        LittleEndian::read_u16(&self.slice[start..])
+    }
+
+    pub fn file_reference(&self) -> NtfsFileReference
+    where
+        E: NtfsIndexEntryHasFileReference,
+    {
+        // The "file_reference_data" is at the same position as the `data_offset`, `data_length`, and `padding` fields.
+        // There can either be extra data or a file reference!
+        let file_reference_data = LittleEndian::read_u64(self.slice);
+        NtfsFileReference::new(file_reference_data)
     }
 
     pub fn flags(&self) -> NtfsIndexEntryFlags {
@@ -67,18 +148,10 @@ impl<'s> NtfsIndexEntry<'s> {
         LittleEndian::read_u16(&self.slice[start..])
     }
 
-    pub fn key_length(&self) -> u16 {
-        let start = offset_of!(IndexEntryHeader, key_length);
-        LittleEndian::read_u16(&self.slice[start..])
-    }
-
     /// Returns the structured value of the key of this Index Entry,
     /// or `None` if this Index Entry has no key.
     /// The last Index Entry never has a key.
-    pub fn key_structured_value<K>(&self) -> Option<Result<K>>
-    where
-        K: NtfsStructuredValueFromSlice<'s>,
-    {
+    pub fn key(&self) -> Option<Result<E::KeyType>> {
         // The key/stream is only set when the last entry flag is not set.
         // https://flatcap.org/linux-ntfs/ntfs/concepts/index_entry.html
         if self.key_length() == 0 || self.flags().contains(NtfsIndexEntryFlags::LAST_ENTRY) {
@@ -89,8 +162,16 @@ impl<'s> NtfsIndexEntry<'s> {
         let end = start + self.key_length() as usize;
         let position = self.position + start as u64;
 
-        let structured_value = iter_try!(K::from_slice(&self.slice[start..end], position));
-        Some(Ok(structured_value))
+        let key = iter_try!(E::KeyType::key_from_slice(
+            &self.slice[start..end],
+            position
+        ));
+        Some(Ok(key))
+    }
+
+    pub fn key_length(&self) -> u16 {
+        let start = offset_of!(IndexEntryHeader, key_length);
+        LittleEndian::read_u16(&self.slice[start..])
     }
 
     /// Returns the Virtual Cluster Number (VCN) of the subnode of this Index Entry,
@@ -108,20 +189,29 @@ impl<'s> NtfsIndexEntry<'s> {
     }
 }
 
-pub(crate) struct IndexNodeEntryRanges {
+pub(crate) struct IndexNodeEntryRanges<E>
+where
+    E: NtfsIndexEntryType,
+{
     data: Vec<u8>,
     range: Range<usize>,
     position: u64,
+    entry_type: PhantomData<E>,
 }
 
-impl IndexNodeEntryRanges {
+impl<E> IndexNodeEntryRanges<E>
+where
+    E: NtfsIndexEntryType,
+{
     pub(crate) fn new(data: Vec<u8>, range: Range<usize>, position: u64) -> Self {
         debug_assert!(range.end <= data.len());
+        let entry_type = PhantomData;
 
         Self {
             data,
             range,
             position,
+            entry_type,
         }
     }
 
@@ -130,8 +220,11 @@ impl IndexNodeEntryRanges {
     }
 }
 
-impl Iterator for IndexNodeEntryRanges {
-    type Item = IndexEntryRange;
+impl<E> Iterator for IndexNodeEntryRanges<E>
+where
+    E: NtfsIndexEntryType,
+{
+    type Item = IndexEntryRange<E>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.range.is_empty() {
@@ -141,7 +234,7 @@ impl Iterator for IndexNodeEntryRanges {
         // Get the current entry.
         let start = self.range.start;
         let position = self.position + self.range.start as u64;
-        let entry = NtfsIndexEntry::new(&self.data[start..], position);
+        let entry = NtfsIndexEntry::<'_, E>::new(&self.data[start..], position);
         let end = start + entry.index_entry_length() as usize;
 
         if entry.flags().contains(NtfsIndexEntryFlags::LAST_ENTRY) {
@@ -158,22 +251,37 @@ impl Iterator for IndexNodeEntryRanges {
     }
 }
 
-impl FusedIterator for IndexNodeEntryRanges {}
+impl<E> FusedIterator for IndexNodeEntryRanges<E> where E: NtfsIndexEntryType {}
 
 #[derive(Clone, Debug)]
-pub struct NtfsIndexNodeEntries<'s> {
+pub struct NtfsIndexNodeEntries<'s, E>
+where
+    E: NtfsIndexEntryType,
+{
     slice: &'s [u8],
     position: u64,
+    entry_type: PhantomData<E>,
 }
 
-impl<'s> NtfsIndexNodeEntries<'s> {
+impl<'s, E> NtfsIndexNodeEntries<'s, E>
+where
+    E: NtfsIndexEntryType,
+{
     pub(crate) fn new(slice: &'s [u8], position: u64) -> Self {
-        Self { slice, position }
+        let entry_type = PhantomData;
+        Self {
+            slice,
+            position,
+            entry_type,
+        }
     }
 }
 
-impl<'s> Iterator for NtfsIndexNodeEntries<'s> {
-    type Item = NtfsIndexEntry<'s>;
+impl<'s, E> Iterator for NtfsIndexNodeEntries<'s, E>
+where
+    E: NtfsIndexEntryType,
+{
+    type Item = NtfsIndexEntry<'s, E>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.slice.is_empty() {
@@ -199,4 +307,4 @@ impl<'s> Iterator for NtfsIndexNodeEntries<'s> {
     }
 }
 
-impl<'s> FusedIterator for NtfsIndexNodeEntries<'s> {}
+impl<'s, E> FusedIterator for NtfsIndexNodeEntries<'s, E> where E: NtfsIndexEntryType {}
diff --git a/src/index_record.rs b/src/index_record.rs
index e6fccd3..e0fab4f 100644
--- a/src/index_record.rs
+++ b/src/index_record.rs
@@ -4,6 +4,7 @@
 use crate::attribute_value::NtfsNonResidentAttributeValue;
 use crate::error::{NtfsError, Result};
 use crate::index_entry::{IndexNodeEntryRanges, NtfsIndexNodeEntries};
+use crate::indexes::NtfsIndexEntryType;
 use crate::record::Record;
 use crate::record::RecordHeader;
 use crate::traits::NtfsReadSeek;
@@ -66,7 +67,10 @@ impl<'n> NtfsIndexRecord<'n> {
         Ok(index_record)
     }
 
-    pub fn entries<'r>(&'r self) -> Result<NtfsIndexNodeEntries<'r>> {
+    pub fn entries<'r, E>(&'r self) -> Result<NtfsIndexNodeEntries<'r, E>>
+    where
+        E: NtfsIndexEntryType,
+    {
         let (entries_range, position) = self.entries_range_and_position();
         let data = &self.record.data()[entries_range];
 
@@ -104,7 +108,10 @@ impl<'n> NtfsIndexRecord<'n> {
         LittleEndian::read_u32(&self.record.data()[start..])
     }
 
-    pub(crate) fn into_entry_ranges(self) -> IndexNodeEntryRanges {
+    pub(crate) fn into_entry_ranges<E>(self) -> IndexNodeEntryRanges<E>
+    where
+        E: NtfsIndexEntryType,
+    {
         let (entries_range, position) = self.entries_range_and_position();
         IndexNodeEntryRanges::new(self.record.into_data(), entries_range, position)
     }
diff --git a/src/indexes/file_name.rs b/src/indexes/file_name.rs
new file mode 100644
index 0000000..069836a
--- /dev/null
+++ b/src/indexes/file_name.rs
@@ -0,0 +1,14 @@
+// Copyright 2021 Colin Finck <colin@reactos.org>
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+use crate::indexes::{NtfsIndexEntryHasFileReference, NtfsIndexEntryType};
+use crate::structured_values::NtfsFileName;
+
+#[derive(Debug)]
+pub struct NtfsFileNameIndex {}
+
+impl NtfsIndexEntryType for NtfsFileNameIndex {
+    type KeyType = NtfsFileName;
+}
+
+impl NtfsIndexEntryHasFileReference for NtfsFileNameIndex {}
diff --git a/src/indexes/mod.rs b/src/indexes/mod.rs
new file mode 100644
index 0000000..1e83146
--- /dev/null
+++ b/src/indexes/mod.rs
@@ -0,0 +1,30 @@
+// Copyright 2021 Colin Finck <colin@reactos.org>
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+mod file_name;
+
+pub use file_name::*;
+
+use crate::error::Result;
+
+pub trait NtfsIndexEntryType {
+    type KeyType: NtfsIndexEntryKey;
+}
+
+pub trait NtfsIndexEntryKey: Sized {
+    fn key_from_slice(slice: &[u8], position: u64) -> Result<Self>;
+}
+
+/// Indicates that the index entry type has additional data.
+// This would benefit from negative trait bounds, as this trait and `NtfsIndexEntryHasFileReference` are mutually exclusive!
+pub trait NtfsIndexEntryHasData: NtfsIndexEntryType {
+    type DataType: NtfsIndexEntryData;
+}
+
+pub trait NtfsIndexEntryData: Sized {
+    fn data_from_slice(slice: &[u8], position: u64) -> Result<Self>;
+}
+
+/// Indicates that the index entry type has a file reference.
+// This would benefit from negative trait bounds, as this trait and `NtfsIndexEntryHasData` are mutually exclusive!
+pub trait NtfsIndexEntryHasFileReference: NtfsIndexEntryType {}
diff --git a/src/lib.rs b/src/lib.rs
index 487ac50..5db12fb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -18,6 +18,7 @@ mod guid;
 mod index;
 mod index_entry;
 mod index_record;
+pub mod indexes;
 mod ntfs;
 mod ntfs_file;
 mod record;
diff --git a/src/structured_values/file_name.rs b/src/structured_values/file_name.rs
index 32e40dc..0fdac97 100644
--- a/src/structured_values/file_name.rs
+++ b/src/structured_values/file_name.rs
@@ -4,6 +4,7 @@
 use crate::attribute::NtfsAttributeType;
 use crate::error::{NtfsError, Result};
 use crate::file_reference::NtfsFileReference;
+use crate::indexes::NtfsIndexEntryKey;
 use crate::string::NtfsString;
 use crate::structured_values::{
     NtfsFileAttributeFlags, NtfsStructuredValue, NtfsStructuredValueFromSlice,
@@ -163,6 +164,13 @@ impl<'s> NtfsStructuredValueFromSlice<'s> for NtfsFileName {
     }
 }
 
+// `NtfsFileName` is special in the regard that the index entry key has the same structure as the structured value.
+impl NtfsIndexEntryKey for NtfsFileName {
+    fn key_from_slice(slice: &[u8], position: u64) -> Result<Self> {
+        Self::from_slice(slice, position)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/structured_values/index_root.rs b/src/structured_values/index_root.rs
index 797fd56..cae761a 100644
--- a/src/structured_values/index_root.rs
+++ b/src/structured_values/index_root.rs
@@ -5,6 +5,7 @@ use crate::attribute::NtfsAttributeType;
 use crate::error::{NtfsError, Result};
 use crate::index_entry::{IndexNodeEntryRanges, NtfsIndexNodeEntries};
 use crate::index_record::{IndexNodeHeader, INDEX_NODE_HEADER_SIZE};
+use crate::indexes::NtfsIndexEntryType;
 use crate::structured_values::{NtfsStructuredValue, NtfsStructuredValueFromSlice};
 use byteorder::{ByteOrder, LittleEndian};
 use core::ops::Range;
@@ -30,7 +31,10 @@ pub struct NtfsIndexRoot<'f> {
 const LARGE_INDEX_FLAG: u8 = 0x01;
 
 impl<'f> NtfsIndexRoot<'f> {
-    pub fn entries(&self) -> Result<NtfsIndexNodeEntries<'f>> {
+    pub fn entries<E>(&self) -> Result<NtfsIndexNodeEntries<'f, E>>
+    where
+        E: NtfsIndexEntryType,
+    {
         let (entries_range, position) = self.entries_range_and_position();
         let slice = &self.slice[entries_range];
 
@@ -45,7 +49,10 @@ impl<'f> NtfsIndexRoot<'f> {
         (start..end, position)
     }
 
-    pub(crate) fn entry_ranges(&self) -> IndexNodeEntryRanges {
+    pub(crate) fn entry_ranges<E>(&self) -> IndexNodeEntryRanges<E>
+    where
+        E: NtfsIndexEntryType,
+    {
         let (entries_range, position) = self.entries_range_and_position();
         let entries_data = self.slice[entries_range].to_vec();
         let range = 0..entries_data.len();