diff --git a/src/arch/x86_64/test/bin.rs b/src/arch/x86_64/test/bin.rs index e181297..238065f 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, pe::Import}, + backend::{Instr as BInstr, Program, pe::LibImport}, }; use std::{fs::OpenOptions, io::Write, os::unix::fs::OpenOptionsExt, process::Command}; @@ -72,14 +72,18 @@ fn linux() { fn windows() { let mut program = Program::::default(); let entry = program.func([BInstr::Asm(Asm { - instrs: vec![push(39), pop(rax), Instr::Ret], + instrs: vec![ + push(39), + pop(rax), + Instr::Ret + ], })]); program.entry = Some(entry); let linked = program.compile().expect("failed to compile"); - let imports = &[Import { + let imports = &[LibImport { name: "KERNEL32.dll".to_string(), - names: ["GetStdHandle", "WriteFile", "ExitProcess"] + syms: ["GetStdHandle", "WriteFile", "ExitProcess"] .map(String::from) .to_vec(), }]; diff --git a/src/backend/container/pe/import.rs b/src/backend/container/pe/import.rs index eda26f6..e8487a0 100644 --- a/src/backend/container/pe/import.rs +++ b/src/backend/container/pe/import.rs @@ -2,12 +2,17 @@ use crate::backend::pe::data_dir::DataDir; use super::ByteEncoder; -pub struct Import { +pub struct LibImport { pub name: String, - pub names: Vec, + pub syms: Vec, } -pub fn encode(data: &mut ByteEncoder, imports: &[Import]) -> DataDir { +pub struct SymImport { + name: String, + usages: Vec, +} + +pub fn encode(data: &mut ByteEncoder, imports: &[LibImport]) -> DataDir { let start = data.pos() as u32; let idt = data.reserve_arr::(imports.len()); // null entry to mark end @@ -23,12 +28,12 @@ pub fn encode(data: &mut ByteEncoder, imports: &[Import]) -> DataDir { // lookup table data.align(size_of::()); let lookup_start = data.pos(); - let lookup = data.reserve_arr::(import.names.len()); + let lookup = data.reserve_arr::(import.syms.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); + for (i, sym) in import.syms.iter().enumerate() { + let rva = hint_name_entry(data, 0, &sym.name); data[lookup][i] = ImportLookupEntry::name(rva); } diff --git a/src/backend/container/pe/mod.rs b/src/backend/container/pe/mod.rs index 6a6de9f..88665c7 100644 --- a/src/backend/container/pe/mod.rs +++ b/src/backend/container/pe/mod.rs @@ -7,9 +7,9 @@ use crate::backend::LinkedProgram; use data_dir::*; use header::*; -pub use import::Import; +pub use import::LibImport; -pub fn create(program: &[u8], start_offset: u64, imports: &[Import]) -> Vec { +pub fn create(program: &[u8], start_offset: u64, imports: &[LibImport]) -> Vec { let mut data = ByteEncoder::default(); let file_align = 1; let section_align = 1; @@ -105,7 +105,7 @@ pub fn create(program: &[u8], start_offset: u64, imports: &[Import]) -> Vec } impl LinkedProgram { - pub fn to_pe(&self, imports: &[Import]) -> Vec { + pub fn to_pe(&self, imports: &[LibImport]) -> Vec { create(&self.code, self.entry.expect("no start"), imports) } } diff --git a/src/backend/ir/mod.rs b/src/backend/ir/mod.rs index 7daad48..f18e211 100644 --- a/src/backend/ir/mod.rs +++ b/src/backend/ir/mod.rs @@ -9,6 +9,7 @@ pub struct Program { pub ro_data: Vec, pub funcs: Vec>, pub entry: Option, + pub external: Vec, sym_count: usize, } @@ -23,6 +24,11 @@ pub struct Func { pub sym: Symbol, } +pub struct External { + pub file: String, + pub syms: Vec, +} + pub enum Instr { Set { dst: VarId, src: Vec }, Call { dst: FnId, args: Vec }, @@ -56,6 +62,22 @@ impl Program { sym } + pub fn external( + &mut self, + file: impl Into, + names: [impl Into; LEN], + ) -> [Symbol; LEN] { + let syms = names.map(|s| { + let sym = self.reserve(); + sym + }); + self.external.push(External { + file: file.into(), + syms: syms.to_vec(), + }); + syms + } + fn reserve(&mut self) -> Symbol { let res = Symbol(self.sym_count); self.sym_count += 1; @@ -78,6 +100,7 @@ impl Default for Program { funcs: Default::default(), entry: Default::default(), sym_count: Default::default(), + external: Default::default(), } } }