lancelot 0.3.1

use anyhow::Result;
use log::debug;
use thiserror::Error;

pub mod imports;
pub mod rsrc;

use crate::{
    aspace::RelativeAddressSpace,
    module::{Arch, Module, Permissions, Section},
    util, VA,
};

#[derive(Error, Debug)]
pub enum PEError {
    #[error("format not supported: {0}")]
    FormatNotSupported(String),

    #[error("malformed PE file: {0}")]
    MalformedPEFile(String),
}

/// A parsed and loaded PE file.
/// The `buf` field contains the raw data.
/// The `module` field contains an address space as the PE would be loaded.
pub struct PE {
    pub buf:    Vec<u8>,
    pub module: Module,
    pub header: goblin::pe::header::Header,
}

impl PE {
    pub fn from_bytes(buf: &[u8]) -> Result<PE> {
        load_pe(buf)
    }

    pub fn executable_sections<'b>(&'b self) -> Box<dyn Iterator<Item = &Section> + 'b> {
        Box::new(
            self.module
                .sections
                .iter()
                .filter(|section| section.permissions.intersects(Permissions::X)),
        )
    }

    pub fn pe(&self) -> Result<goblin::pe::PE> {
        get_pe(&self.buf)
    }
}

fn get_pe(buf: &[u8]) -> Result<goblin::pe::PE> {
    match goblin::Object::parse(buf)? {
        goblin::Object::PE(pe) => {
            if let Some(opt) = pe.header.optional_header {
                if opt.data_directories.get_clr_runtime_header().is_some() {
                    return Err(PEError::FormatNotSupported(".NET assembly".to_string()).into());
                }
            }
            Ok(pe)
        }
        goblin::Object::Elf(_) => Err(PEError::FormatNotSupported("elf".to_string()).into()),
        goblin::Object::Archive(_) => Err(PEError::FormatNotSupported("archive".to_string()).into()),
        goblin::Object::Mach(_) => Err(PEError::FormatNotSupported("macho".to_string()).into()),
        goblin::Object::Unknown(_) => Err(PEError::FormatNotSupported("unknown".to_string()).into()),
    }
}

fn load_pe_header(buf: &[u8], pe: &goblin::pe::PE, base_address: VA) -> Result<Section> {
    let hdr_raw_size = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.size_of_headers,
        // assumption: header is at most 0x200 bytes.
        _ => 0x200,
    };

    //   on disk:
    //
    //   +---------------------------------+
    //   |   header        |  sections...  |
    //   +---------------------------------+
    //   .                  \
    //   .  in memory:       \
    //   .                    \
    //   +-----------------+---+        +-------------
    //   |   header        |   |        |  sections...
    //   +-----------------+---+        +-------------
    //                     ^   ^
    //                     |   +--- virtual size
    //                     |        aligned to 0x200
    //                     +-- raw size
    //                         no alignment

    let hdr_raw_size = std::cmp::min(hdr_raw_size as usize, buf.len());
    let hdr_virt_size = util::align(hdr_raw_size as u64, 0x200);

    Ok(Section {
        physical_range: std::ops::Range {
            start: 0x0,
            end:   hdr_raw_size as u64,
        },
        virtual_range:  std::ops::Range {
            start: base_address,
            end:   base_address + hdr_virt_size,
        },
        permissions:    Permissions::R,
        name:           "header".to_string(),
    })
}

/// The section can be executed as code.
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;

/// The section can be read.
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;

/// The section can be written to.
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

fn load_pe_section(
    base_address: VA,
    section_alignment: u64,
    section: &goblin::pe::section_table::SectionTable,
) -> Result<Section> {
    let name = String::from_utf8_lossy(&section.name[..])
        .into_owned()
        .trim_end_matches('\u{0}')
        .trim_end()
        .splitn(2, '\u{0}')
        .next()
        .unwrap()
        .to_string();

    let virtual_size = util::align(section.virtual_size as u64, section_alignment);

    let mut perms = Permissions::empty();
    if section.characteristics & IMAGE_SCN_MEM_READ > 0 {
        perms.insert(Permissions::R);
    }
    if section.characteristics & IMAGE_SCN_MEM_WRITE > 0 {
        perms.insert(Permissions::W);
    }
    if section.characteristics & IMAGE_SCN_MEM_EXECUTE > 0 {
        perms.insert(Permissions::X);
    }

    debug!(
        "pe: section: {} at {:#x}",
        name,
        base_address + section.virtual_address as u64
    );

    Ok(Section {
        physical_range: std::ops::Range {
            start: section.pointer_to_raw_data as u64,
            end:   (section.pointer_to_raw_data + section.size_of_raw_data) as u64,
        },
        virtual_range: std::ops::Range {
            start: base_address + section.virtual_address as u64,
            end:   base_address + section.virtual_address as u64 + virtual_size,
        },
        permissions: perms,
        name,
    })
}

// lots of further detail here: https://github.com/corkami/docs/blob/master/PE/PE.md
fn load_pe(buf: &[u8]) -> Result<PE> {
    let pe = get_pe(buf)?;

    let arch = match pe.is_64 {
        false => Arch::X32,
        true => Arch::X64,
    };
    debug!("pe: arch: {:?}", arch);

    let (base_address, section_alignment) = match pe.header.optional_header {
        Some(opt) => (
            opt.windows_fields.image_base,
            opt.windows_fields.section_alignment as u64,
        ),
        _ => {
            debug!("pe: base address: using default: 0x40:000");
            (0x40_000, 0x1000)
        }
    };
    debug!("pe: base address: {:#x}", base_address);

    let mut sections = vec![load_pe_header(buf, &pe, base_address)?];
    for section in pe.sections.iter() {
        sections.push(load_pe_section(base_address, section_alignment as u64, section)?);
    }

    let max_address = sections.iter().map(|sec| sec.virtual_range.end).max().unwrap();
    let max_page_address = util::align(max_address, 0x1000) - base_address;
    debug!("pe: address space: capacity: {:#x}", max_page_address);

    let mut address_space = RelativeAddressSpace::with_capacity(max_page_address);

    for section in sections.iter() {
        let pstart = section.physical_range.start as usize;
        let pend = section.physical_range.end as usize;
        let psize = pend - pstart;
        let pbuf = &buf[pstart..pend];

        // the section range contains VAs,
        // while we're writing to the RelativeAddressSpace.
        // so shift down by `base_address`.
        let vstart = section.virtual_range.start;
        let rstart = vstart - base_address;
        let vsize = util::align(
            section.virtual_range.end - section.virtual_range.start,
            section_alignment,
        );
        let vend = vstart + vsize;
        let mut vbuf = vec![0u8; vsize as usize];

        if vsize as usize >= psize {
            // vsize > psize, so there will be NULL bytes padding the physical data.
            let dest = &mut vbuf[0..psize as usize];
            dest.copy_from_slice(pbuf);
        } else {
            // psize > vsize, but vsize wins, so we only read a subset of physical data.
            let src = &pbuf[0..vsize as usize];
            vbuf.copy_from_slice(src);
        }

        address_space.map.writezx(rstart, &vbuf)?;

        debug!(
            "pe: address space: mapped {:#x} - {:#x} {:?}",
            vstart, vend, section.permissions
        );
    }

    let module = Module {
        arch,
        sections,
        address_space: address_space.into_absolute(base_address)?,
    };

    debug!("pe: loaded");
    Ok(PE {
        buf: buf.to_vec(),
        module,
        header: pe.header,
    })
}

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use crate::{aspace::AddressSpace, rsrc::*};

    #[test]
    fn base_address() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x1_8000_0000, pe.module.address_space.base_address);

        Ok(())
    }

    #[test]
    fn mz_header() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        // relative read
        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        // absolute read
        assert_eq!(0x4d, pe.module.address_space.read_u8(0x1_8000_0000 + 0x0)?);
        assert_eq!(0x5a, pe.module.address_space.read_u8(0x1_8000_0000 + 0x1)?);

        Ok(())
    }

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    // this demonstrates that the PE will be loaded and sections padded out to their
    // virtual range.
    #[test]
    fn read_each_section() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        for section in pe.module.sections.iter() {
            let start = section.virtual_range.start;
            let size = section.virtual_range.end - section.virtual_range.start;
            pe.module
                .address_space
                .read_bytes(start, size as usize)
                .expect(&format!("read section {} {:#x} {:#x}", section.name, start, size));
        }

        Ok(())
    }
}