tapher-mach.rs

use std::io::{Cursor, Write};

use goblin::elf::{self, Elf};
use goblin::mach::*;
use scroll::{IOwrite, Pwrite, SizeWith};

use crate::PAGE_SIZE;

#[derive(Debug, Default)]
struct Segment {
	name: [u8; 16],
	address: u64,
	size: u64,
	offset: u64,
	file_size: u64,
	protections: u32,
}

impl Segment {
	pub fn name(mut self, name: &[u8]) -> Self {
		self.name[..name.len()].copy_from_slice(name);
		self
	}

	pub fn address(mut self, address: u64) -> Self {
		self.address = address;
		self
	}

	pub fn size(mut self, size: u64) -> Self {
		self.size = size;
		self
	}

	pub fn offset(mut self, offset: u64) -> Self {
		self.offset = offset;
		self
	}

	pub fn file_size(mut self, file_size: u64) -> Self {
		self.file_size = file_size;
		self
	}

	pub fn protections(mut self, protection: u32) -> Self {
		self.protections |= protection;
		self
	}

	pub fn build(self) -> load_command::SegmentCommand64 {
		load_command::SegmentCommand64 {
			cmd: load_command::LC_SEGMENT_64,
			cmdsize: load_command::SIZEOF_SEGMENT_COMMAND_64 as u32,
			segname: self.name,
			vmaddr: self.address,
			vmsize: self.size,
			fileoff: self.offset,
			filesize: self.file_size,
			maxprot: self.protections,
			initprot: self.protections,
			nsects: 0,
			flags: 0,
		}
	}
}

#[repr(C)]
#[derive(Debug, Pwrite, IOwrite, SizeWith)]
struct UnixThreadCommand {
	command: u32,
	command_size: u32,
	flavor: u32,
	count: u32,
	thread_state: [u64; 21],
}

impl UnixThreadCommand {
	const STATE_COUNT: usize = 21;

	pub fn new(instruction_address: u64) -> Self {
		use std::mem::size_of;
		const X86_THREAD_STATE64: u32 = 4;
		const STATE_SIZE: usize = size_of::<[u64; UnixThreadCommand::STATE_COUNT]>();
		const STATE_COUNT: usize = STATE_SIZE / size_of::<u32>();

		const INSTRUCTION_REGISTER: usize = 16;
		let mut thread_state = [0; Self::STATE_COUNT];
		thread_state[INSTRUCTION_REGISTER] = instruction_address;

		UnixThreadCommand {
			command: load_command::LC_UNIXTHREAD,
			command_size: size_of::<Self>() as u32,
			flavor: X86_THREAD_STATE64,
			count: STATE_COUNT as u32,
			thread_state,
		}
	}
}

#[derive(Debug)]
struct LoadCommands {
	zero: load_command::SegmentCommand64,
	header: load_command::SegmentCommand64,
	segments: Vec<load_command::SegmentCommand64>,
	thread: UnixThreadCommand,
}

impl LoadCommands {
	pub fn new(segments: Vec<load_command::SegmentCommand64>, thread: UnixThreadCommand) -> Self {
		let zero = Segment::default().name(b"__PAGEZERO").size(super::PAGE_SIZE as u64).build();
		let header = Segment::default().name(b"__TEXT").address(super::PAGE_SIZE as u64)
			.protections(constants::VM_PROT_EXECUTE | constants::VM_PROT_READ).build();
		LoadCommands { zero, header, segments, thread }
	}

	pub fn size_count(&self) -> (u32, usize) {
		std::iter::once(self.zero.cmdsize)
			.chain(std::iter::once(self.header.cmdsize))
			.chain(std::iter::once(self.thread.command_size))
			.chain(self.segments.iter().map(|command| command.cmdsize))
			.fold((0, 0), |(size, count), command_size| (size + command_size, count + 1))
	}

	pub fn write(self, target: &mut Cursor<Vec<u8>>) -> crate::Result<()> {
		self.segments.into_iter().try_for_each(|segment| target.iowrite(segment))?;
		target.iowrite(self.header)?;
		target.iowrite(self.thread)?;
		target.iowrite(self.zero)?;
		Ok(())
	}
}

fn segments(executable: &Elf) -> Vec<load_command::SegmentCommand64> {
	executable.program_headers.iter()
		.filter(|header| header.p_type == elf::program_header::PT_LOAD)
		.map(|header| Segment::default().name(b"__TEXT").address(header.p_vaddr)
			.size(header.p_memsz).offset(header.p_offset).file_size(header.p_filesz)
			.protections(protections(header)).build()).collect()
}

fn protections(header: &elf::ProgramHeader) -> u32 {
	let mut protections = 0;
	if header.is_executable() { protections |= constants::VM_PROT_EXECUTE };
	if header.is_write() { protections |= constants::VM_PROT_WRITE };
	if header.is_read() { protections |= constants::VM_PROT_READ };
	protections
}

fn fill_page(offset: &mut usize, bytes: &mut Vec<u8>) {
	if *offset % super::PAGE_SIZE == 0 { return; }
	let alignment = PAGE_SIZE - (*offset % super::PAGE_SIZE);
	let length = bytes.len() + alignment;
	bytes.resize(length, 0);
	*offset += alignment;
}

pub fn translate(executable: &Elf, data: &[u8]) -> crate::Result<Vec<u8>> {
	let thread = UnixThreadCommand::new(executable.entry);
	let mut commands = LoadCommands::new(segments(executable), thread);
	let (size, command_count) = commands.size_count();

	let mut bytes = Vec::new();
	let mut offset = header::SIZEOF_HEADER_64 + size as usize;
	commands.header.filesize = offset as u64;

	fill_page(&mut offset, &mut bytes);
	for segment in &mut commands.segments {
		let file_end = (segment.fileoff + segment.filesize) as usize;
		let slice = &data[segment.fileoff as usize..file_end];
		segment.fileoff = offset as u64;
		offset += bytes.write(slice)?;
		fill_page(&mut offset, &mut bytes);
	}

	let mut target = Cursor::new(Vec::new());
	target.iowrite(header::Header {
		magic: header::MH_MAGIC_64,
		cputype: constants::cputype::CPU_TYPE_X86_64,
		cpusubtype: constants::cputype::CPU_SUBTYPE_X86_64_ALL,
		filetype: header::MH_EXECUTE,
		ncmds: command_count,
		sizeofcmds: size,
		flags: header::MH_NOUNDEFS,
		reserved: 0,
	})?;

	commands.write(&mut target)?;
	target.write_all(&bytes)?;
	Ok(target.into_inner())
}