Skip to content

Commit

Permalink
Merge pull request #120 from bjorn3/inject_refactor
Browse files Browse the repository at this point in the history 
Refactor headcrab_inject
  • Loading branch information
@nbaksalyar
nbaksalyar committed Oct 4, 2020
2 parents 8866e42 + 2bb3e22 commit 444ba54
Show file tree
Hide file tree
Showing 5 changed files with 267 additions and 231 deletions.
10 changes: 5 additions & 5 deletions examples/repl.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ mod example {
};

#[cfg(target_os = "linux")]
use headcrab_inject::{compile_clif_code, DataId, FuncId, InjectionContext};
use headcrab_inject::{inject_clif_code, DataId, FuncId, InjectionModule};

use repl_tools::HighlightAndComplete;
use rustyline::{completion::Pair, CompletionType};
Expand Down Expand Up @@ -808,8 +808,8 @@ mod example {
fn inject_clif(context: &mut Context, file: PathBuf) -> CrabResult<()> {
context.load_debuginfo_if_necessary()?;

let mut inj_ctx = InjectionContext::new(context.remote()?)?;
let run_function = headcrab_inject::inject_clif_code(
let mut inj_ctx = InjectionModule::new(context.remote()?)?;
let run_function = inject_clif_code(
&mut inj_ctx,
&|sym| context.debuginfo().get_symbol_address(sym).unwrap() as u64,
&std::fs::read_to_string(file)?,
Expand Down Expand Up @@ -851,7 +851,7 @@ mod example {
fn inject_lib(context: &mut Context, file: PathBuf) -> CrabResult<()> {
context.load_debuginfo_if_necessary()?;

let mut inj_ctx = InjectionContext::new(context.remote()?)?;
let mut inj_ctx = InjectionModule::new(context.remote()?)?;
inj_ctx.define_function(
FuncId::from_u32(0),
context.debuginfo().get_symbol_address("dlopen").unwrap() as u64,
Expand All @@ -875,7 +875,7 @@ mod example {
for func in functions {
ctx.clear();
ctx.func = func;
compile_clif_code(&mut inj_ctx, &*isa, &mut ctx)?;
inj_ctx.compile_clif_code(&*isa, &mut ctx)?;
}

let run_function = inj_ctx.lookup_function(FuncId::from_u32(2));
Expand Down
268 changes: 55 additions & 213 deletions headcrab_inject/src/lib.rs
View file
Original file line number Diff line number Diff line change
@@ -1,161 +1,26 @@
// FIXME make this work on other systems too.
#![cfg(all(target_arch = "x86_64", target_os = "linux"))]

use std::collections::HashMap;

use cranelift_codegen::{
binemit, ir,
isa::{self, TargetIsa},
settings::{self, Configurable},
};
use headcrab::{target::LinuxTarget, CrabResult};
use headcrab::CrabResult;

use headcrab::target::LinuxTarget;

mod memory;
mod module;

pub use cranelift_codegen::Context;
pub use cranelift_module::{DataId, FuncId, FuncOrDataId};
pub use cranelift_reader::parse_functions;
pub use memory::Memory;
pub use module::InjectionModule;

#[derive(Debug)]
struct RelocEntry {
offset: binemit::CodeOffset,
reloc: binemit::Reloc,
name: ir::ExternalName,
addend: binemit::Addend,
}

#[derive(Default)]
struct VecRelocSink(Vec<RelocEntry>);

impl binemit::RelocSink for VecRelocSink {
fn reloc_block(&mut self, _: binemit::CodeOffset, _: binemit::Reloc, _: binemit::CodeOffset) {
todo!()
}
fn reloc_external(
&mut self,
offset: binemit::CodeOffset,
_: ir::SourceLoc,
reloc: binemit::Reloc,
name: &ir::ExternalName,
addend: binemit::Addend,
) {
self.0.push(RelocEntry {
offset,
reloc,
name: name.clone(),
addend,
});
}
fn reloc_constant(&mut self, _: binemit::CodeOffset, _: binemit::Reloc, _: ir::ConstantOffset) {
todo!()
}
fn reloc_jt(&mut self, _: binemit::CodeOffset, _: binemit::Reloc, _: ir::entities::JumpTable) {
todo!()
}
}

// FIXME unmap memory when done
pub struct InjectionContext<'a> {
target: &'a LinuxTarget,
code: Memory,
readonly: Memory,
readwrite: Memory,
functions: HashMap<FuncId, u64>,
data_objects: HashMap<DataId, u64>,
breakpoint_trap: u64,
}

impl<'a> InjectionContext<'a> {
pub fn new(target: &'a LinuxTarget) -> CrabResult<Self> {
let mut inj_ctx = Self {
target,
code: Memory::new_executable(),
readonly: Memory::new_readonly(),
readwrite: Memory::new_writable(),
functions: HashMap::new(),
data_objects: HashMap::new(),
breakpoint_trap: 0,
};

inj_ctx.breakpoint_trap = inj_ctx.code.allocate(target, 1, 8)?;
inj_ctx
.target
.write()
.write(&0xcc, inj_ctx.breakpoint_trap as usize)
.apply()?;

Ok(inj_ctx)
}

pub fn target(&self) -> &LinuxTarget {
self.target
}

pub fn breakpoint_trap(&self) -> u64 {
self.breakpoint_trap
}

/// Allocate a new stack and return the bottom of the stack.
pub fn new_stack(&mut self, size: u64) -> CrabResult<u64> {
let stack = self.allocate_readwrite(size)?;

// Ensure that we hit a breakpoint trap when returning from the injected function.
self.target()
.write()
.write(
&(self.breakpoint_trap() as usize),
stack as usize + size as usize - std::mem::size_of::<usize>(),
)
.apply()?;

// Stack grows downwards on x86_64
Ok(stack + size - std::mem::size_of::<usize>() as u64)
}

pub fn allocate_code(&mut self, size: u64) -> CrabResult<u64> {
self.code.allocate(self.target, size, 8)
}

pub fn allocate_readonly(&mut self, size: u64) -> CrabResult<u64> {
self.readonly.allocate(self.target, size, 8)
}

pub fn allocate_readwrite(&mut self, size: u64) -> CrabResult<u64> {
self.readwrite.allocate(self.target, size, 8)
}

pub fn define_function(&mut self, func_id: FuncId, addr: u64) {
assert!(self.functions.insert(func_id, addr).is_none());
}

pub fn lookup_function(&self, func_id: FuncId) -> u64 {
self.functions[&func_id]
}

pub fn define_data_object(&mut self, data_id: DataId, addr: u64) {
assert!(self.data_objects.insert(data_id, addr).is_none());
}

pub fn lookup_data_object(&self, data_id: DataId) -> u64 {
self.data_objects[&data_id]
}

pub fn define_data_object_with_bytes(
&mut self,
data_id: DataId,
bytes: &[u8],
) -> CrabResult<()> {
let alloc = self.allocate_readonly(bytes.len() as u64)?;
self.target
.write()
.write_slice(bytes, alloc as usize)
.apply()?;
self.define_data_object(data_id, alloc);

Ok(())
}
}
const EXECUTABLE_DATA_ALIGNMENT: u64 = 0x10;
const WRITABLE_DATA_ALIGNMENT: u64 = 0x8;
const READONLY_DATA_ALIGNMENT: u64 = 0x1;

pub fn target_isa() -> Box<dyn TargetIsa> {
let mut flag_builder = settings::builder();
Expand All @@ -166,70 +31,6 @@ pub fn target_isa() -> Box<dyn TargetIsa> {
.finish(flags)
}

pub fn compile_clif_code(
inj_ctx: &mut InjectionContext,
isa: &dyn TargetIsa,
ctx: &mut Context,
) -> CrabResult<()> {
let mut code_mem = Vec::new();
let mut relocs = VecRelocSink::default();

ctx.compile_and_emit(
isa,
&mut code_mem,
&mut relocs,
&mut binemit::NullTrapSink {},
&mut binemit::NullStackmapSink {},
)
.unwrap();

let code_region = inj_ctx.allocate_code(code_mem.len() as u64)?;

for reloc_entry in relocs.0.drain(..) {
let sym = match reloc_entry.name {
ir::ExternalName::User {
namespace: 0,
index,
} => inj_ctx.lookup_function(FuncId::from_u32(index)),
ir::ExternalName::User {
namespace: 1,
index,
} => inj_ctx.lookup_data_object(DataId::from_u32(index)),
_ => todo!("{:?}", reloc_entry.name),
};
match reloc_entry.reloc {
binemit::Reloc::Abs8 => {
code_mem[reloc_entry.offset as usize..reloc_entry.offset as usize + 8]
.copy_from_slice(&u64::to_ne_bytes((sym as i64 + reloc_entry.addend) as u64));
}
_ => todo!("reloc kind for {:?}", reloc_entry),
}
}

inj_ctx
.target
.write()
.write_slice(&code_mem, code_region as usize)
.apply()?;

inj_ctx.define_function(
match ctx.func.name {
ir::ExternalName::User { namespace, index } => {
assert_eq!(namespace, 0);
FuncId::from_u32(index)
}
ir::ExternalName::TestCase {
length: _,
ascii: _,
} => todo!(),
ir::ExternalName::LibCall(_) => panic!("Can't define libcall"),
},
code_region,
);

Ok(())
}

fn parse_func_or_data(s: &str) -> FuncOrDataId {
let (kind, index) = s.split_at(4);
let index: u32 = index.parse().unwrap();
Expand All @@ -242,7 +43,7 @@ fn parse_func_or_data(s: &str) -> FuncOrDataId {
}

pub fn inject_clif_code(
inj_ctx: &mut InjectionContext,
inj_module: &mut InjectionModule,
lookup_symbol: &dyn Fn(&str) -> u64,
code: &str,
) -> CrabResult<u64> {
Expand All @@ -263,10 +64,10 @@ pub fn inject_clif_code(
let content = content.trim_start();
match parse_func_or_data(id) {
FuncOrDataId::Func(func_id) => {
inj_ctx.define_function(func_id, lookup_symbol(content));
inj_module.define_function(func_id, lookup_symbol(content));
}
FuncOrDataId::Data(data_id) => {
inj_ctx.define_data_object(data_id, lookup_symbol(content));
inj_module.define_data_object(data_id, lookup_symbol(content));
}
}
}
Expand All @@ -280,7 +81,8 @@ pub fn inject_clif_code(
.trim_matches('"')
.replace("\\n", "\n")
.replace("\\0", "\0");
inj_ctx.define_data_object_with_bytes(data_id, content.as_bytes())?;
inj_module
.define_data_object_with_bytes(data_id, content.as_bytes())?;
} else {
todo!();
}
Expand Down Expand Up @@ -313,10 +115,50 @@ pub fn inject_clif_code(
for func in functions {
ctx.clear();
ctx.func = func;
compile_clif_code(inj_ctx, &*isa, &mut ctx)?;
inj_module.compile_clif_code(&*isa, &mut ctx)?;
}

let run_function = inj_ctx.lookup_function(run_function.expect("Missing `run` directive"));
let run_function = inj_module.lookup_function(run_function.expect("Missing `run` directive"));

Ok(run_function)
}

pub struct InjectionContext<'a> {
target: &'a LinuxTarget,
code: Memory,
readonly: Memory,
readwrite: Memory,
}

impl<'a> InjectionContext<'a> {
pub fn new(target: &'a LinuxTarget) -> Self {
Self {
target,
code: Memory::new_executable(),
readonly: Memory::new_readonly(),
readwrite: Memory::new_writable(),
}
}

pub fn target(&self) -> &'a LinuxTarget {
self.target
}

pub fn allocate_code(&mut self, size: u64, align: Option<u64>) -> CrabResult<u64> {
self.code.allocate(
self.target,
size,
align.unwrap_or(EXECUTABLE_DATA_ALIGNMENT),
)
}

pub fn allocate_readonly(&mut self, size: u64, align: Option<u64>) -> CrabResult<u64> {
self.readonly
.allocate(self.target, size, align.unwrap_or(READONLY_DATA_ALIGNMENT))
}

pub fn allocate_readwrite(&mut self, size: u64, align: Option<u64>) -> CrabResult<u64> {
self.readwrite
.allocate(self.target, size, align.unwrap_or(WRITABLE_DATA_ALIGNMENT))
}
}
6 changes: 3 additions & 3 deletions headcrab_inject/src/memory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,9 +88,9 @@ impl Memory {
}
}

pub fn allocate(&mut self, target: &LinuxTarget, size: u64, align: u8) -> CrabResult<u64> {
if self.position % align as u64 != 0 {
self.position += align as u64 - self.position % align as u64;
pub fn allocate(&mut self, target: &LinuxTarget, size: u64, align: u64) -> CrabResult<u64> {
if self.position % align != 0 {
self.position += align - self.position % align;
debug_assert!(self.position % align as u64 == 0);
}

Expand Down
Loading

0 comments on commit 444ba54

Please sign in to comment.