dnfile 0.4.0

// Several private helpers take `&u32` / `&usize` arguments for historical
// reasons (these types are trivially `Copy`). Changing them ripples through
// many call sites for no behavioural benefit.
#![allow(clippy::trivially_copy_pass_by_ref)]
// `ClrHeaderFlags::new` deliberately returns a `BTreeSet<Self>` (bitfield
// expansion) rather than `Self` — this is part of the public API.
#![allow(clippy::new_ret_no_self)]

use serde::{Deserialize, Serialize};

pub mod error;
pub mod lang;
pub mod stream;
pub mod utils;

use crate::{
    error::Error,
    stream::meta_data_tables::mdtables::{enums::*, *},
};

pub type Result<T> = std::result::Result<T, Error>;

/// A parsed .NET PE file.
///
/// `DnPe` borrows the underlying file buffer (`&'a [u8]`) and exposes the
/// CLR header, metadata streams, tables and CIL method bodies. The caller
/// owns the buffer; back it with `std::fs::read` or `memmap2::Mmap`.
///
/// # Example
///
/// ```no_run
/// let data = std::fs::read("Sample.exe")?;
/// let pe = dnfile::DnPe::parse(&data)?;
/// let clr = pe.net()?;
/// println!("{} functions", clr.functions().len());
/// # Ok::<(), dnfile::error::Error>(())
/// ```
#[derive(Debug, Serialize)]
pub struct DnPe<'a> {
    #[serde(skip_serializing)]
    data: &'a [u8],
    /// PE section table cached at construction.
    /// Avoids re-parsing the whole PE on every `offset()` / `get_data()` call,
    /// which is in the hot path of every metadata-table row and every CIL instruction.
    #[serde(skip_serializing)]
    sections: Vec<goblin::pe::section_table::SectionTable>,
    /// PE optional-header `file_alignment` cached at construction (same reason).
    #[serde(skip_serializing)]
    file_alignment: u32,
    net: Option<ClrData<'a>>,
}

impl<'a> DnPe<'a> {
    pub fn net(&self) -> Result<&ClrData<'a>> {
        self.net.as_ref().ok_or(Error::NotImplementedError)
    }

    /// Re-parse and return a fresh `goblin::pe::PE` view over the file bytes.
    ///
    /// Internal hot-path code uses the cached `sections` + `file_alignment`
    /// fields directly; this method is kept on the public surface for
    /// consumers that need the full `goblin::pe::PE` value. Calling it is
    /// linear in file size, so cache the result yourself if you need it
    /// repeatedly.
    pub fn pe(&self) -> Result<goblin::pe::PE<'_>> {
        match goblin::Object::parse(self.data)? {
            goblin::Object::PE(pe) => Ok(pe),
            _ => Err(Error::UnsupportedBinaryFormat("main")),
        }
    }

    /// Parse a .NET PE from a borrowed byte buffer.
    ///
    /// Zero-copy: heaps, streams and method-body readers all hold slices
    /// into `data` rather than copies.
    pub fn parse(data: &'a [u8]) -> Result<Self> {
        // Parse the PE exactly once and extract the bits we'll need on the
        // hot path. Previously every `offset()` / `get_data()` call re-parsed
        // the entire binary via `self.pe()`.
        let (sections, file_alignment, clr_directory) = {
            let pe = match goblin::Object::parse(data)? {
                goblin::Object::PE(pe) => pe,
                _ => return Err(Error::UnsupportedBinaryFormat("main")),
            };
            let opt_header = pe
                .header
                .optional_header
                .ok_or(Error::UnsupportedBinaryFormat("optional header absence"))?;
            let file_alignment = opt_header.windows_fields.file_alignment;
            let clr_directory = opt_header
                .data_directories
                .get_clr_runtime_header()
                .copied()
                .ok_or(Error::UnsupportedBinaryFormat("ClR runtime header absence"))?;
            (pe.sections.clone(), file_alignment, clr_directory)
        };

        let mut res = DnPe {
            data,
            sections,
            file_alignment,
            net: None,
        };
        let clr_struct: ClrStruct = res.get_data(
            &clr_directory.virtual_address,
            &(clr_directory.size as usize),
        )?;
        res.net = Some(res.new_clrdata(clr_struct)?);
        Ok(res)
    }

    fn offset(&self, rva: u32) -> Result<usize> {
        goblin::pe::utils::find_offset(
            rva as usize,
            &self.sections,
            self.file_alignment,
            &goblin::pe::options::ParseOptions::default(),
        )
        .ok_or(Error::UnresolvedRvaError(rva))
    }

    fn get_data<'b, T>(&'b self, rva: &'b u32, size: &'b usize) -> Result<T>
    where
        T: scroll::ctx::TryFromCtx<'b, goblin::container::Endian, Error = scroll::Error>,
    {
        Ok(goblin::pe::utils::get_data(
            self.data,
            &self.sections,
            goblin::pe::data_directories::DataDirectory {
                virtual_address: *rva,
                size: *size as u32,
            },
            self.file_alignment,
        )?)
    }

    fn get_nullterminated_string(&self, rva: &u32) -> Result<String> {
        let mut res_buf = vec![];
        let mut rrva = *rva;
        let mut c = self.data[self.offset(rrva)?];
        while c != 0 {
            res_buf.push(c);
            rrva += 1;
            c = self.data[self.offset(rrva)?];
        }
        Ok(String::from_utf8(res_buf)?)
    }

    /// Borrowed slice into the file buffer. Replaces the prior `Vec<u8>`-returning
    /// helper to keep stream/heap construction zero-copy.
    fn get_slice(&self, rva: &u32, size: usize) -> Result<&'a [u8]> {
        let offset = self.offset(*rva)?;
        self.data
            .get(offset..offset + size)
            .ok_or(Error::UnresolvedRvaError(*rva))
    }

    /// Owned variant kept for callers (mdtables) that intentionally take ownership.
    fn get_vec(&self, rva: &u32, size: &usize) -> Result<Vec<u8>> {
        Ok(self.get_slice(rva, *size)?.to_vec())
    }

    fn get_dword_at_rva(&self, rva: &u32) -> Result<u32> {
        self.get_data(rva, &4)
    }

    fn new_clrdata(&self, clr_struct: ClrStruct) -> Result<ClrData<'a>> {
        let metadata_struct: MetaDataStruct = self.get_data(
            &clr_struct.meta_data_rva,
            &(clr_struct.meta_data_size as usize),
        )?;
        let metadata = self.new_metadata(&clr_struct.meta_data_rva, metadata_struct)?;
        let flags = ClrHeaderFlags::new(clr_struct.flags as usize);
        let functions = self.parse_functions(&metadata)?;
        Ok(ClrData {
            metadata,
            flags,
            functions,
        })
    }

    fn parse_functions(
        &self,
        metadata: &MetaData<'a>,
    ) -> Result<Vec<lang::cil::function::Function>> {
        let mut res = vec![];
        let method_def_table = metadata.md_table("MethodDef")?;
        for i in 0..method_def_table.row_count() {
            let row = method_def_table.row::<MethodDef>(i)?;
            if !row
                .impl_flags
                .contains(&ClrMethodImpl::MethodCodeType(CorMethodCodeType::IL))
                || row
                    .flags
                    .contains(&ClrMethodAttr::AttrFlag(CorMethodAttrFlag::Abstract))
                || row
                    .flags
                    .contains(&ClrMethodAttr::AttrFlag(CorMethodAttrFlag::PinvokeImpl))
            {
                continue;
            }
            res.push(self.parse_function(row)?);
        }
        Ok(res)
    }

    fn parse_function(&self, row: &MethodDef) -> Result<lang::cil::function::Function> {
        let mut reader = lang::cil::function::reader::Reader::new(self.data);
        reader.seek(self.offset(row.rva)?)?;
        lang::cil::function::Function::new(&mut reader)
    }

    fn new_metadata(
        &self,
        metadata_rva: &u32,
        metadata_struct: MetaDataStruct,
    ) -> Result<MetaData<'a>> {
        let version_offset = self.offset(metadata_rva + 16)?;
        let version = self.data
            [version_offset..version_offset + metadata_struct.version_length as usize]
            .to_vec();
        let flags: u16 =
            self.get_data(&(metadata_rva + 16 + metadata_struct.version_length), &2)?;
        let number_of_streams: u16 = self.get_data(
            &(metadata_rva + 16 + metadata_struct.version_length + 2),
            &2,
        )?;
        let struct_size = 16 + metadata_struct.version_length + 2 + 2;
        let streams = if number_of_streams > 0 {
            let streams_table_rva = metadata_rva + struct_size;
            self.new_streams(
                metadata_rva,
                &streams_table_rva,
                &(number_of_streams as usize),
            )?
        } else {
            std::collections::HashMap::new()
        };
        Ok(MetaData {
            _version: String::from_utf8(version)?,
            flags,
            streams,
        })
    }

    fn new_streams(
        &self,
        metadata_rva: &u32,
        streams_table_rva: &u32,
        number_of_streams: &usize,
    ) -> Result<std::collections::HashMap<String, stream::ClrStream<'a>>> {
        let mut res = std::collections::HashMap::new();
        let mut stream_entry_rva = *streams_table_rva;
        for _i in 0..*number_of_streams {
            let stream = self.new_clr_stream(&stream_entry_rva, metadata_rva)?;
            stream_entry_rva += stream.stream_table_entry_size as u32;
            res.insert(stream.name().to_string(), stream);
        }
        let mut rres = std::collections::HashMap::new();
        for (n, s) in &res {
            rres.insert(n.to_string(), self.parse_clr_stream(s, &res)?);
        }
        Ok(rres)
    }

    fn new_clr_stream(
        &self,
        stream_table_entry_rva: &u32,
        metadata_rva: &u32,
    ) -> Result<stream::ClrStream<'a>> {
        let stream_offset: u32 = self.get_data(stream_table_entry_rva, &4)?;
        let stream_size: u32 = self.get_data(&(stream_table_entry_rva + 4), &4)?;
        let stream_name = self.get_nullterminated_string(&(stream_table_entry_rva + 8))?;
        let stream_data = self.get_slice(&(metadata_rva + stream_offset), stream_size as usize)?;
        self.nnew_clr_stream(
            metadata_rva,
            &stream_offset,
            &(stream_size as usize),
            &stream_name,
            stream_data,
        )
    }
}

#[repr(C)]
#[derive(Deserialize, Debug, scroll::Pread)]
pub struct ClrStruct {
    cb: u32,
    majorr_runtime_version: u16,
    minor_runtime_version: u16,
    meta_data_rva: u32,
    meta_data_size: u32,
    flags: u32,
    entry_point_token_or_rva: u32,
    resources_rva: u32,
    resources_size: u32,
    strong_name_signature_rva: u32,
    strong_name_signature_size: u32,
    code_manager_table_rva: u32,
    code_manager_table_size: u32,
    v_table_fixups_rva: u32,
    v_table_fixups_size: u32,
    export_address_table_jumps_rva: u32,
    export_address_table_jumps_size: u32,
    managed_native_header_rva: u32,
    managed_native_header_size: u32,
}

#[derive(Debug, Serialize, PartialOrd, Ord, PartialEq, Eq)]
pub enum ClrHeaderFlags {
    IlOnly,
    BitRequired32,
    IlLibrary,
    StrongNamesSigned,
    NativeEntryPiont,
    TrackDebugData,
    Prefer32Bit,
}

impl ClrHeaderFlags {
    pub fn new(value: usize) -> std::collections::BTreeSet<Self> {
        let mut res = std::collections::BTreeSet::new();
        if value & 1 != 0 {
            res.insert(Self::IlOnly);
        }
        if value & 2 != 0 {
            res.insert(Self::BitRequired32);
        }
        if value & 4 != 0 {
            res.insert(Self::IlLibrary);
        }
        if value & 8 != 0 {
            res.insert(Self::StrongNamesSigned);
        }
        if value & 0x10 != 0 {
            res.insert(Self::NativeEntryPiont);
        }
        if value & 0x10000 != 0 {
            res.insert(Self::TrackDebugData);
        }
        if value & 0x20000 != 0 {
            res.insert(Self::Prefer32Bit);
        }
        res
    }
}

#[derive(Debug, Serialize)]
pub struct ClrData<'a> {
    pub metadata: MetaData<'a>,
    pub flags: std::collections::BTreeSet<ClrHeaderFlags>,
    pub functions: Vec<lang::cil::function::Function>,
}

impl<'a> ClrData<'a> {
    pub fn md_table(
        &self,
        name: &'static str,
    ) -> Result<&stream::meta_data_tables::mdtables::MetaDataTable> {
        self.metadata.md_table(name)
    }

    pub fn md_table_by_index(
        &self,
        index: &usize,
    ) -> Result<&stream::meta_data_tables::mdtables::MetaDataTable> {
        self.metadata.md_table_by_index(index)
    }

    pub fn resolve_coded_index<T>(
        &self,
        index: &dyn stream::meta_data_tables::mdtables::codedindex::CodedIndex,
    ) -> Result<&T>
    where
        T: stream::meta_data_tables::mdtables::MDTableRowTrait + 'static,
    {
        let table = self.md_table(index.table())?;
        table.row(index.row_index() - 1)
    }

    pub fn functions(&self) -> &Vec<lang::cil::function::Function> {
        &self.functions
    }

    pub fn get_us(&self, rid: usize) -> Result<String> {
        self.metadata.get_us(rid)
    }
}

#[repr(C)]
#[derive(Deserialize, Debug, Clone, scroll::Pread)]
pub struct MetaDataStruct {
    signature: u32,
    major_version: u16,
    minor_version: u16,
    reserved: u32,
    version_length: u32,
}

#[derive(Debug, Serialize)]
pub struct MetaData<'a> {
    #[serde(skip_serializing)]
    _version: String,
    flags: u16,
    pub streams: std::collections::HashMap<String, stream::ClrStream<'a>>,
}

impl<'a> MetaData<'a> {
    pub fn md_table(
        &self,
        name: &'static str,
    ) -> Result<&stream::meta_data_tables::mdtables::MetaDataTable> {
        for s in self.streams.values() {
            if let stream::Stream::MetaDataTables(mt) = &s.stream {
                match mt
                    .tables
                    .get(&stream::meta_data_tables::mdtables::table_name_2_index(
                        name,
                    )?) {
                    Some(s) => return Ok(s),
                    None => return Err(Error::UndefinedMetaDataTableName(name)),
                }
            }
        }
        Err(Error::UndefinedMetaDataTableName(name))
    }

    pub fn md_table_by_index(
        &self,
        index: &usize,
    ) -> Result<&stream::meta_data_tables::mdtables::MetaDataTable> {
        for s in self.streams.values() {
            if let stream::Stream::MetaDataTables(mt) = &s.stream {
                match mt.tables.get(index) {
                    Some(s) => return Ok(s),
                    None => return Err(Error::UndefinedMetaDataTableIndex(*index as u32)),
                }
            }
        }
        Err(Error::UndefinedMetaDataTableIndex(*index as u32))
    }

    pub fn get_us(&self, rid: usize) -> Result<String> {
        for s in self.streams.values() {
            if let stream::Stream::UserStringHeap(us) = &s.stream {
                return us.get_us(rid);
            }
        }
        Err(Error::UndefinedMetaDataTableName("US"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn clr_header_flags_decodes_bits() {
        let flags = ClrHeaderFlags::new(0b0000_0001);
        assert!(flags.contains(&ClrHeaderFlags::IlOnly));
        assert!(!flags.contains(&ClrHeaderFlags::BitRequired32));
    }

    #[test]
    fn clr_header_flags_decodes_all_low_bits() {
        let flags = ClrHeaderFlags::new(0x1F); // bits 0-4
        for v in [
            ClrHeaderFlags::IlOnly,
            ClrHeaderFlags::BitRequired32,
            ClrHeaderFlags::IlLibrary,
            ClrHeaderFlags::StrongNamesSigned,
            ClrHeaderFlags::NativeEntryPiont,
        ] {
            assert!(flags.contains(&v), "missing flag: {v:?}");
        }
    }

    #[test]
    fn clr_header_flags_decodes_high_bits() {
        let flags = ClrHeaderFlags::new(0x30000); // bits 16 + 17
        assert!(flags.contains(&ClrHeaderFlags::TrackDebugData));
        assert!(flags.contains(&ClrHeaderFlags::Prefer32Bit));
    }

    #[test]
    fn clr_header_flags_zero_value_is_empty() {
        let flags = ClrHeaderFlags::new(0);
        assert!(flags.is_empty());
    }
}