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();