diff --git a/src/arch/x86_64/test/bin.rs b/src/arch/x86_64/test/bin.rs index eddbdf2..e181297 100644 --- a/src/arch/x86_64/test/bin.rs +++ b/src/arch/x86_64/test/bin.rs @@ -1,6 +1,6 @@ use crate::{ arch::x86_64::*, - backend::{Instr as BInstr, Program}, + backend::{Instr as BInstr, Program, pe::Import}, }; use std::{fs::OpenOptions, io::Write, os::unix::fs::OpenOptionsExt, process::Command}; @@ -76,17 +76,25 @@ fn windows() { })]); program.entry = Some(entry); let linked = program.compile().expect("failed to compile"); - let binary = linked.to_pe(); + + let imports = &[Import { + name: "KERNEL32.dll".to_string(), + names: ["GetStdHandle", "WriteFile", "ExitProcess"] + .map(String::from) + .to_vec(), + }]; + + let binary = linked.to_pe(imports); let path = "./x86_64_test.exe"; write(path, &binary); - let mut cmd = Command::new("wine"); - cmd.arg("x86_64_test"); - let mut proc = cmd.spawn().expect("failed to run"); - let status = proc.wait().expect("failed to wait"); - if let Some(code) = status.code() { - std::process::exit(code); - } + // let mut cmd = Command::new("wine"); + // cmd.arg("x86_64_test"); + // let mut proc = cmd.spawn().expect("failed to run"); + // let status = proc.wait().expect("failed to wait"); + // if let Some(code) = status.code() { + // std::process::exit(code); + // } } fn write(path: &str, binary: &[u8]) { diff --git a/src/backend/container/elf.rs b/src/backend/container/elf.rs index 71d2dca..342bcf7 100644 --- a/src/backend/container/elf.rs +++ b/src/backend/container/elf.rs @@ -1,6 +1,7 @@ -use crate::backend::{Addr, LinkedProgram, container::as_u8_slice}; +use crate::backend::{LinkedProgram, container::encode::ByteEncoder}; #[repr(C)] +#[derive(Default)] pub struct ELF64Header { magic: u32, class: u8, @@ -25,6 +26,7 @@ pub struct ELF64Header { } #[repr(C)] +#[derive(Default)] pub struct ProgramHeader { ty: u32, flags: u32, @@ -80,19 +82,17 @@ pub fn create(program: &[u8], start_offset: u64) -> Vec { let page_size = 0x1000; // I don't know if I have to add addr_start here, idk how it maps the memory let program_size = std::mem::size_of_val(program) as u64; - let program_header = ProgramHeader { - ty: 0x1, // LOAD - flags: 0b101, // executable, readable - offset: 0x0, - vaddr: addr_start, - paddr: 0x0, - filesz: program_size, - memsz: program_size, - align: page_size, - }; - let header_len = (size_of::() + size_of::()) as u64; - let program_pos = header_len; - let header = ELF64Header { + + let mut data = ByteEncoder::default(); + let header = data.reserve::(); + + let program_header_offset = data.pos() as u64; + let program_header = data.reserve::(); + + let program_pos = data.pos() as u64; + data.extend(program); + + data[header] = ELF64Header { magic: 0x7f_45_4c_46u32.swap_bytes(), class: 0x2, // 64 bit endianness: 0x1, // little endian @@ -103,8 +103,8 @@ pub fn create(program: &[u8], start_offset: u64) -> Vec { ty: if pie { EType::Dyn } else { EType::Exec } as u16, machine: Arch::X86_64.machine(), e_version: 0x1, - entry: addr_start + program_pos + start_offset as u64, - program_header_offset: size_of::() as u64, + entry: addr_start + program_pos + start_offset, + program_header_offset, section_header_offset: 0x0, // C ABI (16 bit instruction align) + double precision floats flags: 0x1 | 0x4, @@ -115,11 +115,17 @@ pub fn create(program: &[u8], start_offset: u64) -> Vec { section_header_num: 0x0, section_header_str_idx: 0x0, }; - let mut bytes: Vec = Vec::new(); - bytes.extend(as_u8_slice(&header)); - bytes.extend(as_u8_slice(&program_header)); - bytes.extend(program); - bytes + data[program_header] = ProgramHeader { + ty: 0x1, // LOAD + flags: 0b101, // executable, readable + offset: 0x0, + vaddr: addr_start, + paddr: 0x0, + filesz: program_size, + memsz: program_size, + align: page_size, + }; + data.data } impl LinkedProgram { diff --git a/src/backend/container/encode.rs b/src/backend/container/encode.rs new file mode 100644 index 0000000..0f03e02 --- /dev/null +++ b/src/backend/container/encode.rs @@ -0,0 +1,134 @@ +use std::ops::{Index, IndexMut}; + +#[derive(Default)] +pub struct ByteEncoder { + pub data: Vec, +} + +impl ByteEncoder { + pub fn push(&mut self, byte: u8) { + self.data.push(byte); + } + + pub fn val(&mut self, val: &T) -> Reserved { + let pos = self.pos(); + let slice = + unsafe { core::slice::from_raw_parts((val as *const T) as *const u8, size_of::()) }; + self.data.extend(slice); + Reserved::new(pos) + } + + pub fn pos(&self) -> usize { + self.data.len() + } + + pub fn align(&mut self, align: usize) { + self.data.resize(self.data.len().next_multiple_of(align), 0); + } + + #[must_use] + pub fn reserve(&mut self) -> Reserved { + let pos = self.pos(); + self.data.resize(self.data.len() + size_of::(), 0); + Reserved::new(pos) + } + + pub fn pad(&mut self, amt: usize) { + self.data.resize(self.data.len() + amt, 0); + } + + #[must_use] + pub fn reserve_arr(&mut self, len: usize) -> ReservedArr { + let pos = self.pos(); + self.data.resize(self.data.len() + size_of::() * len, 0); + ReservedArr::new(pos, len) + } +} + +pub struct Reserved { + pos: usize, + _pd: std::marker::PhantomData, +} + +impl Clone for Reserved { + fn clone(&self) -> Self { + *self + } +} +impl Copy for Reserved {} + +pub struct ReservedArr { + pos: usize, + len: usize, + _pd: std::marker::PhantomData, +} + +impl Clone for ReservedArr { + fn clone(&self) -> Self { + *self + } +} +impl Copy for ReservedArr {} + +impl Reserved { + fn new(pos: usize) -> Self { + Self { + pos, + _pd: std::marker::PhantomData, + } + } +} + +impl ReservedArr { + fn new(pos: usize, len: usize) -> Self { + Self { + pos, + len, + _pd: std::marker::PhantomData, + } + } +} + +impl Index> for ByteEncoder { + type Output = T; + + fn index(&self, index: Reserved) -> &Self::Output { + let slice = &self.data[index.pos..index.pos + size_of::()]; + unsafe { &core::slice::from_raw_parts((slice as *const [u8]) as *const T, 1)[0] } + } +} + +impl IndexMut> for ByteEncoder { + fn index_mut(&mut self, index: Reserved) -> &mut Self::Output { + let slice = &mut self.data[index.pos..index.pos + size_of::()]; + unsafe { &mut core::slice::from_raw_parts_mut((slice as *mut [u8]) as *mut T, 1)[0] } + } +} + +impl Index> for ByteEncoder { + type Output = [T]; + + fn index(&self, index: ReservedArr) -> &Self::Output { + let slice = &self.data[index.pos..index.pos + size_of::() * index.len]; + unsafe { core::slice::from_raw_parts((slice as *const [u8]) as *const T, index.len) } + } +} + +impl IndexMut> for ByteEncoder { + fn index_mut(&mut self, index: ReservedArr) -> &mut Self::Output { + let slice = &mut self.data[index.pos..index.pos + size_of::() * index.len]; + unsafe { core::slice::from_raw_parts_mut((slice as *mut [u8]) as *mut T, index.len) } + } +} + +impl Extend for ByteEncoder { + fn extend>(&mut self, iter: T) { + self.data.extend(iter); + } +} + +impl<'a> Extend<&'a u8> for ByteEncoder { + fn extend>(&mut self, iter: T) { + self.data.extend(iter); + } +} diff --git a/src/backend/container/mod.rs b/src/backend/container/mod.rs index 8bf524f..1f1708d 100644 --- a/src/backend/container/mod.rs +++ b/src/backend/container/mod.rs @@ -1,38 +1,5 @@ pub mod elf; +mod encode; pub mod pe; -#[derive(Default)] -pub struct ByteEncoder { - pub data: Vec, -} - -trait Pushable { - fn push(self, data: &mut Vec); -} - -impl ByteEncoder { - pub fn push(&mut self, byte: u8) { - self.data.push(byte); - } - pub fn extend(&mut self, iter: impl IntoIterator) { - self.data.extend(iter); - } - pub fn val(&mut self, val: &T) { - self.data.extend(as_u8_slice(val)); - } - pub fn fill(&mut self, pos: usize, val: &T) { - self.data[pos..pos + size_of::()].copy_from_slice(as_u8_slice(val)); - } - pub fn pos(&self) -> usize { - self.data.len() - } - pub fn reserve(&mut self) -> usize { - let pos = self.pos(); - self.val(&T::default()); - pos - } -} - -fn as_u8_slice(p: &T) -> &[u8] { - unsafe { core::slice::from_raw_parts((p as *const T) as *const u8, size_of::()) } -} +use encode::*; diff --git a/src/backend/container/pe/header.rs b/src/backend/container/pe/header.rs index 6a71fa7..10574d2 100644 --- a/src/backend/container/pe/header.rs +++ b/src/backend/container/pe/header.rs @@ -18,7 +18,6 @@ pub struct PeHeader { } #[repr(C)] -#[derive(Default)] pub struct OptHeader64 { pub magic: u16, pub major_linker_ver: u8, @@ -52,7 +51,6 @@ pub struct OptHeader64 { } #[repr(C)] -#[derive(Default)] pub struct Section { pub name: [u8; 8], pub virtual_size: u32, diff --git a/src/backend/container/pe/import.rs b/src/backend/container/pe/import.rs index bdd6013..eda26f6 100644 --- a/src/backend/container/pe/import.rs +++ b/src/backend/container/pe/import.rs @@ -1,33 +1,56 @@ -use crate::backend::ByteEncoder; +use crate::backend::pe::data_dir::DataDir; + +use super::ByteEncoder; pub struct Import { pub name: String, pub names: Vec, } -pub fn encode(data: &mut ByteEncoder, imports: &[Import]) -> usize { - let mut names = 0; - for import in imports { - for name in &import.names { - data.extend(ImportLookupEntry::name().bytes()); - names += 1; +pub fn encode(data: &mut ByteEncoder, imports: &[Import]) -> DataDir { + let start = data.pos() as u32; + let idt = data.reserve_arr::(imports.len()); + // null entry to mark end + data.pad(size_of::()); + let end = data.pos() as u32; + + for (i, import) in imports.iter().enumerate() { + // name + let name_rva = data.pos() as u32; + data.extend(import.name.as_bytes()); + data.push(0); + + // lookup table + data.align(size_of::()); + let lookup_start = data.pos(); + let lookup = data.reserve_arr::(import.names.len()); + data.pad(size_of::()); + let lookup_end = data.pos(); + + for (i, name) in import.names.iter().enumerate() { + let rva = hint_name_entry(data, 0, name); + data[lookup][i] = ImportLookupEntry::name(rva); } - data.extend(ImportLookupEntry::NULL.bytes()); - } - let table_addr = data.pos(); - for import in imports { - let idt = ImportDirTable { - lookup_table_rva: todo!(), + + // address table + data.align(size_of::()); + let addr_start = data.pos(); + let len = lookup_end - lookup_start; + data.pad(len); + data.data.copy_within(lookup_start..lookup_end, addr_start); + + // entry + data[idt][i] = ImportDirTable { + lookup_table_rva: lookup_start as u32, time_date_stamp: 0, forwarder_chain: 0, - name_rva: todo!(), - address_table_rva: todo!(), + name_rva, + address_table_rva: addr_start as u32, }; } - for import in imports { - for name in &import.names { - hint_name_entry(data, 0, &name); - } + DataDir { + virt_addr_rva: start, + size: end - start, } } @@ -40,16 +63,6 @@ pub struct ImportDirTable { pub address_table_rva: u32, } -impl ImportDirTable { - pub const NULL: Self = Self { - lookup_table_rva: 0, - time_date_stamp: 0, - forwarder_chain: 0, - name_rva: 0, - address_table_rva: 0, - }; -} - #[repr(C)] pub struct ImportLookupEntry(u64); @@ -67,11 +80,11 @@ impl ImportLookupEntry { } } -pub fn hint_name_entry(out: &mut Vec, hint: u16, name: &str) { - out.extend(hint.to_le_bytes()); - out.extend(name.as_bytes()); - out.push(0); - if out.len() % 2 == 1 { - out.push(0); - } +pub fn hint_name_entry(data: &mut ByteEncoder, hint: u16, name: &str) -> u32 { + let pos = data.pos() as u32; + data.extend(hint.to_le_bytes()); + data.extend(name.as_bytes()); + data.push(0); + data.align(2); + pos } diff --git a/src/backend/container/pe/mod.rs b/src/backend/container/pe/mod.rs index 40b3523..6a6de9f 100644 --- a/src/backend/container/pe/mod.rs +++ b/src/backend/container/pe/mod.rs @@ -2,62 +2,64 @@ mod data_dir; mod header; mod import; +use super::*; +use crate::backend::LinkedProgram; + use data_dir::*; use header::*; +pub use import::Import; -use crate::backend::{ByteEncoder, LinkedProgram, pe::import::Import}; - -pub fn create(program: &[u8], start_offset: u64) -> Vec { +pub fn create(program: &[u8], start_offset: u64, imports: &[Import]) -> Vec { let mut data = ByteEncoder::default(); let file_align = 1; let section_align = 1; - let code_size = program.len() as u32; + let num_of_rva_and_sizes: u32 = (size_of::() / size_of::()) as u32; - let mzheader = MZHeader { - magic: u16::from_ne_bytes(*b"MZ"), - stuff: [0; _], - lfanew: size_of::() as u32, - }; - data.val(&mzheader); - let header = PeHeader { + let mz_header = data.reserve::(); + + let pe_header_pos = data.pos(); + data.val(&PeHeader { magic: u32::from_ne_bytes(*b"PE\0\0"), machine: 0x8664, num_sections: 1, time_date_stamp: 0, sym_tab_ptr: 0, num_symbols: 0, - opt_header_size: size_of::() as u16, + opt_header_size: (size_of::() + size_of::()) as u16, // https://learn.microsoft.com/en-us/windows/win32/debug/pe-format // executable | can handle >2GB addrs | debug info removed characteristics: 0x2 | 0x20 | 0x0200, + }); + + data[mz_header] = MZHeader { + magic: u16::from_ne_bytes(*b"MZ"), + stuff: [0; _], + lfanew: pe_header_pos as u32, }; - data.val(&header); - let opt_header_pos = data.reserve::(); + let opt_header = data.reserve::(); - let num_of_rva_and_sizes: u32 = (size_of::() / size_of::()) as u32; - data.val(&DataDirs::default()); + let data_dirs = data.val(&DataDirs::default()); - let code_sect_pos = data.reserve::
(); + let code_sect = data.reserve::
(); let hdr_size = data.pos() as u32; - let imports = [Import { - name: "kernel32.lib".to_string(), - names: ["GetStdHandle", "WriteFile", "ExitProcess"] - .map(String::from) - .to_vec(), - }]; - import::encode(&mut data, &imports); + let code_start = data.pos() as u32; + if !imports.is_empty() { + let import_rva = import::encode(&mut data, &imports); + data[data_dirs].import = import_rva; + } - let code_addr = data.pos() as u32; - data.data.extend(program); + let program_start = data.pos() as u32; + data.extend(program); + let code_size = data.pos() as u32 - code_start; - let code_section = Section { + data[code_sect] = Section { name: *b".text\0\0\0", virtual_size: code_size, virtual_addr: hdr_size.next_multiple_of(section_align), raw_data_size: code_size.next_multiple_of(file_align), - raw_data_ptr: code_addr, + raw_data_ptr: code_start, reloc_ptr: 0, line_num_ptr: 0, num_relocs: 0, @@ -67,15 +69,15 @@ pub fn create(program: &[u8], start_offset: u64) -> Vec { let file_size = data.pos() as u32; - let opt_header = OptHeader64 { + data[opt_header] = OptHeader64 { magic: 0x20b, major_linker_ver: 8, minor_linker_ver: 0, code_size: code_size.next_multiple_of(file_align), init_data_size: 0, uninit_data_size: 0, - entry_addr: code_addr + start_offset as u32, - code_base: code_addr, + entry_addr: program_start + start_offset as u32, + code_base: code_start, image_base: 0x400000, section_align, file_align, @@ -99,14 +101,11 @@ pub fn create(program: &[u8], start_offset: u64) -> Vec { num_of_rva_and_sizes, }; - data.fill(opt_header_pos, &opt_header); - data.fill(code_sect_pos, &code_section); - data.data } impl LinkedProgram { - pub fn to_pe(&self) -> Vec { - create(&self.code, self.entry.expect("no start")) + pub fn to_pe(&self, imports: &[Import]) -> Vec { + create(&self.code, self.entry.expect("no start"), imports) } } diff --git a/x86_64_test b/x86_64_test index ae8f616..d8879f7 100755 Binary files a/x86_64_test and b/x86_64_test differ diff --git a/x86_64_test.exe b/x86_64_test.exe index 98cb72a..87d20b4 100755 Binary files a/x86_64_test.exe and b/x86_64_test.exe differ