picori 0.1.0

//! Parse Dolphin executables (`.dol`).
//!
//! # Parse
//!
//! Parse from binary stream by calling [`Dol::from_binary`]. On error a
//! [Error][`crate::Error`] is return. Otherwise, the parsing succeeded and you
//! get back a [`Dol`] struct.
//!
//! ## Example
//!
//! This is an example of how to parse a `.dol` file.
//!
//! ```no_run
//! # use std::fs::File;
//! # use picori::Result;
//! fn main() -> Result<()> {
//!     let mut file = File::open("main.dol")?;
//!     let dol = picori::Dol::from_binary(&mut file)?;
//!     println!("entry point: {:#08x}", dol.entry_point());
//!     Ok(())
//! }
//! ```

use std::io::Cursor;

use crate::helper::alignment::AlignPowerOfTwo;
use crate::helper::{ensure, ParseProblem, Parser, ProblemLocation, Seeker};
use crate::Result;

/// Dolphin executable header.
#[derive(Debug, Clone)]
pub struct Header {
    /// Offset of the text sections.
    pub text_offset: [u32; 7],

    /// Offset of the data sections.
    pub data_offset: [u32; 11],

    /// Address of the text sections.
    pub text_address: [u32; 7],

    /// Address of the data sections.
    pub data_address: [u32; 11],

    /// Size of the text sections.
    pub text_size: [u32; 7],

    /// Size of the data sections.
    pub data_size: [u32; 11],

    /// BSS address.
    pub bss_address: u32,

    /// BSS size.
    pub bss_size: u32,

    /// Entry point.
    pub entry_point: u32,
}

/// Dolphin executable section kind.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum SectionKind {
    /// Text section, e.g. `.init`, `.text`, etc.
    Text,

    /// Data section, e.g. `extab_`, `extabindex_`, `.ctors`, `.dtors`,
    /// `.rodata`, `.data`, `.sdata`, `.sdata2`, etc.
    Data,

    /// BSS section, e.g., `.bss`, `.sbss`, `.sbss2`, etc.
    Bss,
}

/// Dolphin executable section.
#[derive(Debug, Clone)]
pub struct Section {
    /// The kind of section this is (text, data, or bss).
    pub kind: SectionKind,

    /// The section name (e.g. `.text`, `.data`, `.rodata`, etc.), this was
    /// guessed from the type of section and order in which they appear in
    /// the Dolphin executable. This is not guaranteed to be correct.
    /// [`Section::guess_name`] is responsible for guessing the name of the
    /// section.
    pub name: &'static str,

    /// The section address to load the section data at startup.
    pub address: u32,

    /// The section size in bytes.
    pub size: u32,

    /// The section size in bytes, rounded up to the nearest multiple of 32.
    pub aligned_size: u32,

    /// The section data. For `.bss` sections ([`SectionKind::Bss`]), this
    /// will be an empty vector.
    pub data: Vec<u8>,

    /// The offset of the section in the `Dolphin Executable`.
    pub offset: Option<u32>,
}

/// RomCopyInfo represents one entry in the [`__rom_copy_info`][`RomCopyInfo`]
/// symbol generated by the linker at the end of the `.init` section. It has
/// information otherwise lost in the process of converting `.elf` to `.dol`,
/// such as, the original unaligned section size. At startup the
/// [`__rom_copy_info`][`RomCopyInfo`] is used to copy each entry from the ROM
/// to the RAM.
#[derive(Debug, Copy, Clone)]
pub struct RomCopyInfo {
    /// Read Only Memory (ROM) address of the section.
    pub rom_address: u32,

    /// Random Access Memory (RAM) address of the section.
    pub ram_address: u32,

    /// The size of the section in bytes.
    pub size: u32,
}

/// List of [`RomCopyInfo`][`RomCopyInfo`] entries.
#[derive(Debug, Clone)]
pub struct RomCopyInfoList {
    /// The offset of `__rom_copy_info` in the `.init` section.
    pub offset: u32,

    /// Entries in the [`__rom_copy_info`][`RomCopyInfo`] symbol.
    pub entries: Vec<RomCopyInfo>,
}

/// BssInitInfo represents one entry in the [`__bss_init_info`][`BssInitInfo`]
/// symbol generated by the linker at the end of the `.init` section. It has
/// information otherwise lost in the process of converting `.elf` to `.dol`,
/// such as, the original unaligned section size and how many `.bss` (`.sbss`,
/// `.bss2`, etc.) sections exists. The final `.dol` file will have a single
/// `.bss` section with the size of the sum of all the `.bss` sections. At
/// startup the [`__bss_init_info`][`BssInitInfo`] is used to zero out the
/// `.bss` section in RAM.
#[derive(Debug, Copy, Clone)]
pub struct BssInitInfo {
    /// Random Access Memory (RAM) address of the section.
    pub ram_address: u32,

    /// The size of the section in bytes.
    pub size: u32,
}

/// List of [`BssInitInfo`][`BssInitInfo`] entries.
#[derive(Debug, Clone)]
pub struct BssInitInfoList {
    /// The offset of `__bss_init_info` in the `.init` section.
    pub offset: u32,

    /// Entries in the [`__bss_init_info`][`BssInitInfo`] symbol.
    pub entries: Vec<BssInitInfo>,
}

/// `.dol` file object.
#[derive(Debug, Clone)]
pub struct Dol {
    /// Header.
    pub header: Header,

    /// Data of `__rom_copy_info` symbol if it exists.
    pub rom_copy_info: Option<RomCopyInfoList>,

    /// Data of `__bss_init_info` symbol if it exists.
    pub bss_init_info: Option<BssInitInfoList>,

    /// Sections.
    pub sections: Vec<Section>,
}

impl RomCopyInfo {
    /// Parse [`RomCopyInfo`] from binary stream.
    fn from_binary<D: Parser + Seeker>(reader: &mut D) -> Result<Self> {
        let rom_copy_info: Result<_> = {
            let rom_address = reader.bu32()?;
            let ram_address = reader.bu32()?;
            let size = reader.bu32()?;
            Ok(RomCopyInfo {
                rom_address,
                ram_address,
                size,
            })
        };

        rom_copy_info.map_err(|_| {
            ParseProblem::InvalidData("invalid RomCopyInfo", std::panic::Location::current()).into()
        })
    }
}

impl BssInitInfo {
    /// Parse [`BssInitInfo`] from binary stream.
    fn from_binary<D: Parser + Seeker>(reader: &mut D) -> Result<Self> {
        let bss_init_info: Result<_> = {
            let ram_address = reader.bu32()?;
            let size = reader.bu32()?;
            Ok(BssInitInfo { ram_address, size })
        };

        bss_init_info.map_err(|_| {
            ParseProblem::InvalidData("invalid BssInitInfo", std::panic::Location::current()).into()
        })
    }
}

/// Search 0x200 bytes from the end of `data` (from the `.init` section)
/// until we find all `__rom_copy_info` entries.
fn rom_copy_info_search(data: &[u8], address: u32) -> Option<RomCopyInfoList> {
    let offset = if data.len() > 0x200 {
        data.len() - 0x200
    } else {
        0
    };

    for offset in offset..data.len() {
        if let Ok(rom_copy_info) = RomCopyInfo::from_binary(&mut Cursor::new(&data[offset..])) {
            if rom_copy_info.ram_address == address
                && rom_copy_info.rom_address == address
                && rom_copy_info.size < 0x2000000
            {
                let rom_copy_info = data[offset..]
                    .chunks(12)
                    .map(|x| RomCopyInfo::from_binary(&mut Cursor::new(x)))
                    .filter_map(|x| x.ok())
                    .take_while(|x| x.ram_address != 0)
                    .collect::<Vec<_>>();

                return Some(RomCopyInfoList {
                    offset:  offset as u32,
                    entries: rom_copy_info,
                });
            }
        }
    }

    None
}

/// Search 0x200 bytes from the end of `data` (from the `.init` section)
/// until we find all `__bss_init_info` entries.
fn bss_init_info_search(data: &[u8], address: u32) -> Option<BssInitInfoList> {
    let offset = if data.len() > 0x200 {
        data.len() - 0x200
    } else {
        0
    };

    for offset in offset..data.len() {
        if let Ok(bss_init_info) = BssInitInfo::from_binary(&mut Cursor::new(&data[offset..])) {
            if bss_init_info.ram_address == address && bss_init_info.size < 0x2000000 {
                let bss_init_info = data[offset..]
                    .chunks(8)
                    .map(|x| BssInitInfo::from_binary(&mut Cursor::new(x)))
                    .filter_map(|x| x.ok())
                    .take_while(|x| x.ram_address != 0)
                    .collect::<Vec<_>>();

                return Some(BssInitInfoList {
                    offset:  offset as u32,
                    entries: bss_init_info,
                });
            }
        }
    }

    None
}

impl Section {
    fn new(
        kind: SectionKind,
        index: usize,
        offset: u32,
        address: u32,
        size: u32,
        aligned_size: u32,
    ) -> Result<Self> {
        Ok(Self {
            kind,
            name: Section::guess_name(kind, index),
            address,
            size,
            aligned_size,
            data: Vec::new(),
            offset: Some(offset),
        })
    }

    fn read_data<D: Parser + Seeker>(&mut self, reader: &mut D, base: u64) -> Result<()> {
        if self.size > 0 && self.offset.is_some() {
            ensure!(
                self.size <= 0x2000000,
                ParseProblem::InvalidRange(
                    "section size (too large)",
                    std::panic::Location::current()
                )
            );

            reader.goto(base + self.offset.unwrap() as u64)?;
            self.data = reader.read_as_vec(self.size as usize)?;
        }

        Ok(())
    }

    /// Guess section name using kind and index.
    pub fn guess_name(kind: SectionKind, index: usize) -> &'static str {
        match kind {
            SectionKind::Text => match index {
                0 => ".init",
                1 => ".text",
                2 => ".text.2",
                3 => ".text.3",
                4 => ".text.4",
                5 => ".text.5",
                6 => ".text.6",
                _ => unreachable!(),
            },
            SectionKind::Data => match index {
                0 => "extab_",
                1 => "extabindex_",
                2 => ".ctors",
                3 => ".dtors",
                4 => ".rodata",
                5 => ".data",
                6 => ".sdata",
                7 => ".sdata2",
                8 => ".data8",
                9 => ".data9",
                10 => ".data10",
                _ => unreachable!(),
            },
            SectionKind::Bss => match index {
                0 => ".bss",
                1 => ".sbss",
                2 => ".sbss2",
                _ => unreachable!(),
            },
        }
    }
}

impl Dol {
    /// Parse [`Dol`] from binary stream.
    ///
    /// This function _should_ not panic and if any error occurs, it will return
    /// [`Err`] of type [`Error`][`crate::Error`]/[`ParseProblem`].
    pub fn from_binary<D: Parser + Seeker>(reader: &mut D) -> Result<Dol> {
        let base = reader.position()?;

        let text_offset = reader.bu32_array::<7>()?;
        let data_offset = reader.bu32_array::<11>()?;
        let text_address = reader.bu32_array::<7>()?;
        let data_address = reader.bu32_array::<11>()?;
        let text_size = reader.bu32_array::<7>()?;
        let data_size = reader.bu32_array::<11>()?;
        let bss_address = reader.bu32()?;
        let bss_size = reader.bu32()?;
        let entry_point = reader.bu32()?;
        let _ = reader.bu32_array::<7>()?;

        let text_sections = text_offset
            .iter()
            .zip(text_address.iter().zip(text_size.iter()))
            .map(|(offset, (address, size))| (*offset, *address, *size));
        let text_sections = text_sections
            .enumerate()
            .map(|(i, x)| (SectionKind::Text, i, x));

        let data_sections = data_offset
            .iter()
            .zip(data_address.iter().zip(data_size.iter()))
            .map(|(offset, (address, size))| (*offset, *address, *size));
        let data_sections = data_sections
            .enumerate()
            .map(|(i, x)| (SectionKind::Data, i, x));

        let sections = text_sections.chain(data_sections);
        let mut sections: Vec<Section> = sections
            .filter(|(_, _, x)| x.1 > 0)
            .map(|(kind, index, (offset, address, size))| {
                ensure!(
                    offset.checked_add(size).is_some(),
                    ParseProblem::InvalidData(
                        "offset + size overflow",
                        std::panic::Location::current()
                    )
                );
                let section = Section::new(kind, index, offset, address, size, size);
                match section {
                    Ok(mut section) => {
                        section.read_data(reader, base)?;
                        Ok(section)
                    },
                    Err(e) => Err(e),
                }
            })
            .collect::<Result<Vec<_>>>()?;

        // Seek to the end of the dol
        let text_sections = text_offset.iter().zip(text_size.iter());
        let data_sections = data_offset.iter().zip(data_size.iter());
        let end_of_dol = text_sections
            .chain(data_sections)
            .map(|(offset, size)| offset.checked_add(*size).unwrap_or(0))
            .max()
            .unwrap_or(0);
        reader.goto(base + end_of_dol as u64)?;

        let init = sections.iter().find(|x| x.name == ".init");
        let rom_copy_info =
            init.and_then(|init| rom_copy_info_search(init.data.as_slice(), init.address));
        let bss_init_info =
            init.and_then(|init| bss_init_info_search(init.data.as_slice(), bss_address));

        for section in sections.iter_mut() {
            section.size = rom_copy_info
                .as_ref()
                .and_then(|v| v.entries.iter().find(|x| x.rom_address == section.address))
                .map_or(section.size, |x| x.size);
        }

        // If `__bss_init_info` is available we can use it to determine the size and
        // count of the `.bss` sections. Otherwise we assume that there is only one
        // `.bss` section and use the size from the header (which is probably
        // not correct).
        if let Some(bss_init_info) = &bss_init_info {
            ensure!(
                bss_init_info.entries.len() <= 3,
                ParseProblem::InvalidRange(
                    "too many _bss_init_info entries",
                    std::panic::Location::current()
                )
            );

            let bss_sections = bss_init_info
                .entries
                .iter()
                .enumerate()
                .map(|(index, entry)| Section {
                    kind:         SectionKind::Bss,
                    name:         Section::guess_name(SectionKind::Bss, index),
                    address:      entry.ram_address,
                    size:         entry.size,
                    aligned_size: entry.size.align_next(32),
                    data:         vec![],
                    offset:       None,
                });
            sections.extend(bss_sections)
        } else {
            // TODO: We can probably use the data section to determine the .bss sections.
            sections.push(Section {
                kind:         SectionKind::Bss,
                name:         Section::guess_name(SectionKind::Bss, 0),
                address:      bss_address,
                size:         bss_size,
                aligned_size: bss_size,
                data:         vec![],
                offset:       None,
            });
        }

        Ok(Dol {
            header: Header {
                text_offset,
                data_offset,
                text_address,
                data_address,
                text_size,
                data_size,
                bss_address,
                bss_size,
                entry_point,
            },
            rom_copy_info,
            bss_init_info,
            sections,
        })
    }

    /// Returns the entry point of the [DOL][`crate::dol`] file. This is the
    /// address of the first instruction that will be executed. The section
    /// containing the entry point can be found using
    /// [`Dol::section_by_address`]. Entry point is also available via
    /// direct access to the [`Dol::header`].
    #[inline]
    pub fn entry_point(&self) -> u32 { self.header.entry_point }

    /// Returns an [`Some(&Section)`] if the [DOL][`crate::dol`] file contains a
    /// section with the given name `name` or [`None`] otherwise. Section
    /// names are not information provided by the `.dol` format, instead we
    /// assign names based on the section kind [`SectionKind`] and the index
    /// of the section.
    #[inline]
    pub fn section_by_name(&self, name: &str) -> Option<&Section> {
        self.sections.iter().find(|x| x.name == name)
    }

    /// Returns an [`Some(&Section)`] if the [DOL][`crate::dol`] file contains a
    /// section with that contains the given address `address` or [`None`]
    /// otherwise.
    #[inline]
    pub fn section_by_address(&self, address: u32) -> Option<&Section> {
        self.sections
            .iter()
            .find(|x| address >= x.address && address < x.address + x.size)
    }
}