fast_ntfs 1.0.2

// Copyright 2021-2026 Colin Finck <colin@reactos.org>
// SPDX-License-Identifier: MIT OR Apache-2.0

use core::iter::FusedIterator;

use crate::attribute::NtfsAttributeType;
use crate::attribute_value::{NtfsAttributeValue, NtfsDataRun};
use crate::error::{NtfsError, Result};
use crate::index_record::{NtfsIndexRecord, validate_index_record_size};
use crate::io::{Read, Seek};
use crate::mapped_data_runs::MappedDataRuns;
use crate::structured_values::NtfsStructuredValue;
use crate::traits::NtfsReadSeek;
use crate::types::{NtfsPosition, Vcn};
use alloc::vec::Vec;

/// Structure of an $INDEX_ALLOCATION attribute.
///
/// This attribute describes the sub-nodes of a B-tree.
/// The top-level nodes are managed via [`NtfsIndexRoot`].
///
/// NTFS uses B-trees for describing directories (as indexes of [`NtfsFileName`]s), looking up Object IDs,
/// Reparse Points, and Security Descriptors, to just name a few.
///
/// An $INDEX_ALLOCATION attribute can be resident or non-resident.
///
/// Reference: <https://flatcap.github.io/linux-ntfs/ntfs/attributes/index_allocation.html>
///
/// [`NtfsFileName`]: crate::structured_values::NtfsFileName
/// [`NtfsIndexRoot`]: crate::structured_values::NtfsIndexRoot
#[derive(Clone, Debug)]
pub struct NtfsIndexAllocation<'n, 'f> {
    value: NtfsAttributeValue<'n, 'f>,
    data_runs: NtfsIndexAllocationDataRuns,
}

#[derive(Clone, Debug)]
pub(crate) struct NtfsIndexAllocationDataRuns {
    cluster_size: u32,
    value_position: NtfsPosition,
    runs: MappedDataRuns,
}

impl NtfsIndexAllocationDataRuns {
    pub(crate) fn record_from_vcn<T>(
        &self,
        fs: &mut T,
        index_record_size: u32,
        vcn: Vcn,
    ) -> Result<NtfsIndexRecord>
    where
        T: Read + Seek,
    {
        self.record_from_vcn_with_buffer(fs, index_record_size, vcn, Vec::new())
    }

    pub(crate) fn record_from_vcn_with_buffer<T>(
        &self,
        fs: &mut T,
        index_record_size: u32,
        vcn: Vcn,
        mut data: Vec<u8>,
    ) -> Result<NtfsIndexRecord>
    where
        T: Read + Seek,
    {
        validate_index_record_size(index_record_size, self.value_position)?;
        let offset = vcn
            .value()
            .checked_mul(self.cluster_size as i64)
            .ok_or(NtfsError::VcnTooBig { vcn })?;
        let offset =
            u64::try_from(offset).map_err(|_| NtfsError::VcnOutOfBoundsInIndexAllocation {
                position: self.value_position,
                vcn,
            })?;
        data.resize(index_record_size as usize, 0);
        let position = self.runs.read_exact_at(fs, offset, &mut data, || {
            NtfsError::VcnOutOfBoundsInIndexAllocation {
                position: self.value_position,
                vcn,
            }
        })?;

        let record = NtfsIndexRecord::from_data(data, position)?;
        validate_record_vcn(record, self.value_position, vcn)
    }
}

impl<'n, 'f> NtfsIndexAllocation<'n, 'f> {
    /// Returns the [`NtfsIndexRecord`] located at the given Virtual Cluster Number (VCN).
    ///
    /// The record is fully read, fixed up, and validated.
    ///
    /// This function is usually called on the return value of [`NtfsIndexEntry::subnode_vcn`] to move further
    /// down in the B-tree.
    ///
    /// [`NtfsIndexEntry::subnode_vcn`]: crate::NtfsIndexEntry::subnode_vcn
    pub fn record_from_vcn<T>(
        &self,
        fs: &mut T,
        index_record_size: u32,
        vcn: Vcn,
    ) -> Result<NtfsIndexRecord>
    where
        T: Read + Seek,
    {
        self.data_runs.record_from_vcn(fs, index_record_size, vcn)
    }

    /// Returns an iterator over all Index Records of this $INDEX_ALLOCATION attribute (cf. [`NtfsIndexRecord`]).
    ///
    /// Each Index Record is fully read, fixed up, and validated.
    pub fn records(&self, index_record_size: u32) -> NtfsIndexRecords<'n, 'f> {
        NtfsIndexRecords::new(self.clone(), index_record_size)
    }

    pub(crate) fn data_runs(&self) -> &NtfsIndexAllocationDataRuns {
        &self.data_runs
    }

    pub(crate) fn into_data_runs(self) -> NtfsIndexAllocationDataRuns {
        self.data_runs
    }
}

impl<'n, 'f> NtfsStructuredValue<'n, 'f> for NtfsIndexAllocation<'n, 'f> {
    const TY: NtfsAttributeType = NtfsAttributeType::IndexAllocation;

    fn from_attribute_value<T>(fs: &mut T, value: NtfsAttributeValue<'n, 'f>) -> Result<Self>
    where
        T: Read + Seek,
    {
        let ntfs = match &value {
            NtfsAttributeValue::AttributeListNonResident(value) => value.ntfs(),
            NtfsAttributeValue::NonResident(value) => value.ntfs(),
            NtfsAttributeValue::Resident(_) => {
                let position = value.data_position();
                return Err(NtfsError::UnexpectedResidentAttribute { position });
            }
        };

        let data_runs = {
            let mut runs = MappedDataRuns::new(value.len());
            let mut push_data_run = |data_run: NtfsDataRun| {
                if !runs.push(data_run.allocated_size(), data_run.data_position()) {
                    return Err(NtfsError::VcnOutOfBoundsInIndexAllocation {
                        position: value.data_position(),
                        vcn: Vcn::from(0),
                    });
                }
                Ok(())
            };
            match &value {
                NtfsAttributeValue::NonResident(non_resident) => {
                    for data_run in non_resident.data_runs() {
                        push_data_run(data_run?)?;
                    }
                }
                NtfsAttributeValue::AttributeListNonResident(attribute_list) => {
                    attribute_list.for_each_data_run(fs, &mut push_data_run)?;
                }
                NtfsAttributeValue::Resident(_) => unreachable!(),
            }
            runs.validate_complete(value.data_position())?;
            NtfsIndexAllocationDataRuns {
                cluster_size: ntfs.cluster_size(),
                value_position: value.data_position(),
                runs,
            }
        };

        Ok(Self { value, data_runs })
    }
}

fn validate_record_vcn(
    record: NtfsIndexRecord,
    position: NtfsPosition,
    vcn: Vcn,
) -> Result<NtfsIndexRecord> {
    if record.vcn() != vcn {
        return Err(NtfsError::VcnMismatchInIndexAllocation {
            position,
            expected: vcn,
            actual: record.vcn(),
        });
    }

    Ok(record)
}

/// Iterator over
///   all index records of an [`NtfsIndexAllocation`],
///   returning an [`NtfsIndexRecord`] for each record.
///
/// This iterator is returned from the [`NtfsIndexAllocation::records`] function.
///
/// See [`NtfsIndexRecordsAttached`] for an iterator that implements [`Iterator`] and [`FusedIterator`].
#[derive(Clone, Debug)]
pub struct NtfsIndexRecords<'n, 'f> {
    index_allocation: NtfsIndexAllocation<'n, 'f>,
    index_record_size: u32,
    finished: bool,
}

impl<'n, 'f> NtfsIndexRecords<'n, 'f> {
    fn new(index_allocation: NtfsIndexAllocation<'n, 'f>, index_record_size: u32) -> Self {
        Self {
            index_allocation,
            index_record_size,
            finished: false,
        }
    }

    /// Returns a variant of this iterator that implements [`Iterator`] and [`FusedIterator`]
    /// by mutably borrowing the filesystem reader.
    pub fn attach<'a, T>(self, fs: &'a mut T) -> NtfsIndexRecordsAttached<'n, 'f, 'a, T>
    where
        T: Read + Seek,
    {
        NtfsIndexRecordsAttached::new(fs, self)
    }

    /// See [`Iterator::next`].
    pub fn next<T>(&mut self, fs: &mut T) -> Option<Result<NtfsIndexRecord>>
    where
        T: Read + Seek,
    {
        if self.finished {
            return None;
        }

        if let Err(error) = validate_index_record_size(
            self.index_record_size,
            self.index_allocation.value.data_position(),
        ) {
            self.finished = true;
            return Some(Err(error));
        }

        if self.index_allocation.value.stream_position() >= self.index_allocation.value.len() {
            self.finished = true;
            return None;
        }

        // Get the current record.
        let record = iter_try!(NtfsIndexRecord::new(
            fs,
            &mut self.index_allocation.value,
            self.index_record_size
        ));

        Some(Ok(record))
    }
}

/// Iterator over
///   all index records of an [`NtfsIndexAllocation`],
///   returning an [`NtfsIndexRecord`] for each record,
///   implementing [`Iterator`] and [`FusedIterator`].
///
/// This iterator is returned from the [`NtfsIndexRecords::attach`] function.
/// Conceptually the same as [`NtfsIndexRecords`], but mutably borrows the filesystem
/// to implement aforementioned traits.
#[derive(Debug)]
pub struct NtfsIndexRecordsAttached<'n, 'f, 'a, T>
where
    T: Read + Seek,
{
    fs: &'a mut T,
    index_records: NtfsIndexRecords<'n, 'f>,
}

impl<'n, 'f, 'a, T> NtfsIndexRecordsAttached<'n, 'f, 'a, T>
where
    T: Read + Seek,
{
    fn new(fs: &'a mut T, index_records: NtfsIndexRecords<'n, 'f>) -> Self {
        Self { fs, index_records }
    }
    /// Consumes this iterator and returns the inner [`NtfsIndexRecords`].
    pub fn detach(self) -> NtfsIndexRecords<'n, 'f> {
        self.index_records
    }
}

impl<'n, 'f, 'a, T> Iterator for NtfsIndexRecordsAttached<'n, 'f, 'a, T>
where
    T: Read + Seek,
{
    type Item = Result<NtfsIndexRecord>;

    fn next(&mut self) -> Option<Self::Item> {
        self.index_records.next(self.fs)
    }
}

impl<'n, 'f, 'a, T> FusedIterator for NtfsIndexRecordsAttached<'n, 'f, 'a, T> where T: Read + Seek {}