sys-rs 0.1.1

use goblin::elf;
use nix::{errno::Errno, sys::wait::wait, unistd::Pid};
use procfs::process::{MMPermissions, MMapPath, Process};
use std::{collections::HashMap, fs::File, io::Read, path::Path};

use crate::diag::{Error, Result};

const AT_ENTRY: u64 = 9;
const AT_PHDR: u64 = 3;
const EI_DATA: usize = 5;
const MAX_OPCODE_SIZE: u64 = 16;

/// Process metadata and parsed ELF image used by the tracer.
///
/// `Info` holds the auxiliary vector, an in-memory copy of the ELF file
/// contents, parsed headers/sections, and the offsets needed to translate
/// runtime addresses to file offsets. Construct using `Info::build`.
pub struct Info {
    pid: Pid,
    auxv: HashMap<u64, u64>,
    buffer: Vec<u8>,
    header: elf::Header,
    sections: HashMap<String, elf::SectionHeader>,
    load_vaddr: u64,
    load_offset: u64,
    mem_offset: u64,
}

impl Info {
    /// Builds an `Info` struct by collecting process data.
    ///
    /// # Arguments
    ///
    /// * `path` - The path to the file containing process data.
    /// * `pid` - The process ID.
    ///
    /// # Errors
    ///
    /// Returns an `Err` upon any failure while collecting process data.
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the `Info` struct upon success, or an `Err` upon failure.
    pub fn build(path: &str, pid: Pid) -> Result<Self> {
        // First, wait for the process to start so we can collect its data.
        wait()?;

        let auxv = Self::get_auxv(pid)?;

        let mut file = File::open(Path::new(path))?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer)?;

        let elf = elf::Elf::parse(&buffer)?;
        let header = elf.header;

        let sections: HashMap<String, elf::SectionHeader> = elf
            .section_headers
            .iter()
            .filter_map(|header| {
                elf.shdr_strtab
                    .get_at(header.sh_name)
                    .map(|name| (name.to_string(), header.clone()))
            })
            .collect();
        sections
            .get(".text")
            .ok_or_else(|| Error::from(Errno::ENODATA))?;

        let dynamic = match elf.header.e_type {
            elf::header::ET_DYN => Ok(true),
            elf::header::ET_EXEC => Ok(false),
            _ => Err(Error::from(Errno::ENOEXEC)),
        }?;

        let (load_vaddr, load_offset) = if dynamic {
            elf.program_headers
                .iter()
                .find(|ph| {
                    ph.p_type == elf::program_header::PT_LOAD && ph.is_executable()
                })
                .map(|ph| (ph.p_vaddr, ph.p_offset))
                .ok_or_else(|| Error::from(Errno::ENODATA))?
        } else {
            (0, 0)
        };

        let mem_offset = if dynamic {
            Self::get_mem_offset(path, pid)?
        } else {
            0
        };

        Ok(Self {
            pid,
            auxv,
            buffer,
            header,
            sections,
            load_vaddr,
            load_offset,
            mem_offset,
        })
    }

    fn get_auxv(pid: Pid) -> Result<HashMap<u64, u64>> {
        let process = Process::new(pid.into())?;
        let auxv = process.auxv()?;
        Ok(auxv.into_iter().collect())
    }

    fn get_mem_offset(path: &str, pid: Pid) -> Result<u64> {
        let absolute_path = std::fs::canonicalize(path)?;

        let process = Process::new(pid.into())?;
        let maps = process.maps()?;

        maps.into_iter()
            .find_map(|map| match &map.pathname {
                MMapPath::Path(buf)
                    if buf == &absolute_path
                        && map.perms.contains(MMPermissions::READ)
                        && map.perms.contains(MMPermissions::EXECUTE) =>
                {
                    Some(map.address.0)
                }
                _ => None,
            })
            .ok_or_else(|| Error::from(Errno::ENODATA))
    }

    fn get_buffer_data(&self, offset: u64, len: u64) -> Result<Option<&[u8]>> {
        let offset = usize::try_from(offset)?;
        let len = usize::try_from(len)?;
        Ok(self.buffer.get(offset..offset + len))
    }

    #[must_use]
    /// Return the PID associated with this `Info`.
    ///
    /// # Returns
    ///
    /// The `Pid` belonging to the traced process.
    pub fn pid(&self) -> Pid {
        self.pid
    }

    #[must_use]
    /// Return the ELF data encoding (endianness) byte (`EI_DATA`).
    ///
    /// This value matches the ELF header `e_ident[EI_DATA]` and can be
    /// compared against `ELFDATA2LSB`/`ELFDATA2MSB` constants.
    ///
    /// # Returns
    ///
    /// The ELF `e_ident[EI_DATA]` byte as a `u8`.
    pub fn endianness(&self) -> u8 {
        self.header.e_ident[EI_DATA]
    }

    #[must_use]
    /// Return the offset used to translate file addresses to runtime
    /// addresses (i.e., `mem_offset - load_vaddr`).
    ///
    /// # Returns
    ///
    /// The computed offset which should be added to file addresses to
    /// obtain runtime addresses (equal to `mem_offset - load_vaddr`).
    pub fn offset(&self) -> u64 {
        self.mem_offset - self.load_vaddr
    }

    #[must_use]
    /// Return true if `addr` lies within the named ELF section (e.g.
    /// ".text"). Addresses are compared against the section's runtime
    /// address (`section.sh_addr` + `mem_offset`).
    ///
    /// # Arguments
    ///
    /// * `addr` - Runtime address to test.
    /// * `name` - The section name to check (for example `".text"`).
    ///
    /// # Returns
    ///
    /// `true` if `addr` falls within the runtime range of the named section
    /// (computed as `section.sh_addr + mem_offset` .. `+ sh_size`), otherwise
    /// `false`.
    pub fn is_addr_in_section(&self, addr: u64, name: &str) -> bool {
        self.sections.get(name).is_some_and(|section| {
            let start = section.sh_addr + self.mem_offset;
            let end = start + section.sh_size;
            (start..end).contains(&addr)
        })
    }

    /// Retrieves the runtime entry point address of the process from the auxiliary vector.
    ///
    /// # Errors
    ///
    /// Returns an `Err` if the entry point is not available in the auxiliary vector.
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing a reference to the entry point address (`u64`) from the auxiliary vector (`AT_ENTRY`).
    pub fn entry(&self) -> Result<&u64> {
        self.auxv
            .get(&AT_ENTRY)
            .ok_or_else(|| Error::from(Errno::ENODATA))
    }

    /// Retrieves the data from the specified section.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the section.
    ///
    /// # Errors
    ///
    /// Returns an `Err` if the conversion from `u64` to `usize` fails when getting buffer data.
    ///
    /// # Returns
    ///
    /// Returns an `Ok` containing the data from the specified section as a slice of bytes, or `Err` if the section does not exist or if the conversion from `u64` to `usize` fails.
    pub fn get_section_data(&self, name: &str) -> Result<Option<&[u8]>> {
        self.sections.get(name).map_or(Ok(None), |section| {
            self.get_buffer_data(section.sh_offset, section.sh_size)
        })
    }

    /// Retrieves opcode data from the loaded binary at a given runtime address.
    ///
    /// # Arguments
    ///
    /// * `addr` - The runtime address for which to fetch opcode bytes.
    ///
    /// # Errors
    ///
    /// Returns an `Err` if the conversion from `u64` to `usize` fails when accessing the buffer, or if required auxiliary vector data is missing.
    ///
    /// # Returns
    ///
    /// Returns `Ok(Some(&[u8]))` containing the opcode bytes at the given address, `Ok(None)` if the address is invalid, or `Err` on conversion failure or missing data.
    pub fn get_opcode_from_addr(&self, addr: u64) -> Result<Option<&[u8]>> {
        let phdr = self
            .auxv
            .get(&AT_PHDR)
            .ok_or_else(|| Error::from(Errno::ENODATA))?;
        let offset =
            (addr - phdr + self.header.e_phoff) + self.load_offset - self.load_vaddr;
        self.get_buffer_data(offset, MAX_OPCODE_SIZE)
    }
}

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

    use std::io::Write;

    fn create_temp_elf_file() -> std::io::Result<String> {
        let path = "/tmp/test.elf";
        let mut file = File::create(path)?;
        file.write_all(b"\x7fELF")?;
        Ok(path.to_string())
    }

    #[test]
    fn test_get_section_data_invalid() {
        let pid = Pid::from_raw(1234);
        let path =
            create_temp_elf_file().expect("Failed to create temporary ELF file");
        let info = Info::build(&path, pid);
        if let Ok(info) = info {
            let data = info.get_section_data(".invalid");
            assert!(matches!(data, Ok(None)));
        }
    }

    #[test]
    fn test_get_opcode_from_addr_invalid_auxv() {
        let pid = Pid::from_raw(1234);
        let path =
            create_temp_elf_file().expect("Failed to create temporary ELF file");
        let info = Info::build(&path, pid);
        if let Ok(info) = info {
            let result = info.get_opcode_from_addr(0xdeadbeef);
            assert!(result.is_err());
        }
    }

    #[test]
    fn test_is_addr_in_section_false() {
        let pid = Pid::from_raw(1234);
        let path =
            create_temp_elf_file().expect("Failed to create temporary ELF file");
        let info = Info::build(&path, pid);
        if let Ok(info) = info {
            let found = info.is_addr_in_section(0xdeadbeef, ".text");
            assert!(!found);
        }
    }

    #[test]
    fn test_entry_missing() {
        let pid = Pid::from_raw(1234);
        let path =
            create_temp_elf_file().expect("Failed to create temporary ELF file");
        let info = Info::build(&path, pid);
        if let Ok(info) = info {
            let entry = info.entry();
            assert!(entry.is_err());
        }
    }

    #[test]
    fn test_build_invalid_path() {
        let pid = Pid::from_raw(1234);
        let result = Info::build("/invalid/path", pid);
        assert!(result.is_err());
    }

    #[test]
    fn test_build_invalid_elf() {
        let pid = Pid::from_raw(1234);
        let path =
            create_temp_elf_file().expect("Failed to create temporary ELF file");
        let result = Info::build(&path, pid);
        assert!(result.is_err());
    }

    #[test]
    fn test_get_mem_offset_invalid_path() {
        let pid = Pid::from_raw(1234);
        let result = Info::get_mem_offset("/invalid/path", pid);
        assert!(result.is_err());
    }

    #[test]
    fn test_get_buffer_and_section_helpers() {
        let pid = Pid::from_raw(1);
        let buffer = vec![0u8; 64];
        let header = elf::Header {
            e_ident: [0; 16],
            e_type: 0,
            e_machine: 0,
            e_version: 0,
            e_entry: 0,
            e_phoff: 0,
            e_shoff: 0,
            e_flags: 0,
            e_ehsize: 0,
            e_phentsize: 0,
            e_phnum: 0,
            e_shentsize: 0,
            e_shnum: 0,
            e_shstrndx: 0,
        };
        let sections: HashMap<String, elf::SectionHeader> = HashMap::new();

        let info = Info {
            pid,
            auxv: HashMap::new(),
            buffer,
            header,
            sections,
            load_vaddr: 0,
            load_offset: 0,
            mem_offset: 0,
        };

        let data = info.get_buffer_data(0, 16).expect("get_buffer_data failed");
        assert!(data.is_some());

        let sec = info.get_section_data(".text");
        assert!(sec.expect("get_section_data returned Err").is_none());
    }

    #[test]
    fn test_info_accessors_and_opcode() {
        let pid = Pid::from_raw(42);

        let mut buffer = vec![0u8; 256];
        for i in 100..116 {
            buffer[i] = (i - 100) as u8;
        }

        let header = elf::Header {
            e_ident: [0; 16],
            e_type: 0,
            e_machine: 0,
            e_version: 0,
            e_entry: 0,
            e_phoff: 20,
            e_shoff: 0,
            e_flags: 0,
            e_ehsize: 0,
            e_phentsize: 0,
            e_phnum: 0,
            e_shentsize: 0,
            e_shnum: 0,
            e_shstrndx: 0,
        };

        let mut sections: HashMap<String, elf::SectionHeader> = HashMap::new();
        let sh = elf::SectionHeader {
            sh_name: 0,
            sh_type: 0,
            sh_flags: 0,
            sh_addr: 0x200,
            sh_offset: 50,
            sh_size: 10,
            sh_link: 0,
            sh_info: 0,
            sh_addralign: 0,
            sh_entsize: 0,
        };
        sections.insert(".text".to_string(), sh);

        let mut auxv = HashMap::new();
        auxv.insert(AT_PHDR, 150u64);

        let info = Info {
            pid,
            auxv,
            buffer,
            header,
            sections,
            load_vaddr: 0,
            load_offset: 0,
            mem_offset: 0,
        };

        assert_eq!(info.pid(), pid);
        assert_eq!(info.endianness(), 0);
        assert_eq!(info.offset(), 0);

        let addr = 0x200 + 5;
        assert!(info.is_addr_in_section(addr, ".text"));

        let sec = info
            .get_section_data(".text")
            .expect("get_section_data failed");
        assert!(sec.is_some());
        let sec = sec.unwrap();
        assert_eq!(sec.len(), 10);

        let addr_for_opcode = 150u64 - 20u64 + 100u64;
        let opc = info
            .get_opcode_from_addr(addr_for_opcode)
            .expect("get_opcode_from_addr failed")
            .expect("opcode not found");

        assert_eq!(opc.len(), usize::try_from(MAX_OPCODE_SIZE).unwrap());
        assert_eq!(opc[0], 0u8);
        assert_eq!(opc[15], 15u8);
    }
}