Skip to content

Commit

Permalink
Add function arg to cli (#20)
Browse files Browse the repository at this point in the history 
* refactor: Add function arg. Remove hardcoded validate_block. Move malloc index to heap injection only

Signed-off-by: Maksim Dimitrov <[email protected]>

* chore(readme): Update readme

Signed-off-by: Maksim Dimitrov <[email protected]>

* chore: Run fmt

Signed-off-by: Maksim Dimitrov <[email protected]>

* chore: Update readme

Signed-off-by: Maksim Dimitrov <[email protected]>

---------

Signed-off-by: Maksim Dimitrov <[email protected]>
  • Loading branch information
@dimitrovmaksim
dimitrovmaksim authored Aug 9, 2023
1 parent 2817b0b commit 57ca2f7
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 72 deletions.
25 changes: 15 additions & 10 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,15 @@ Options:

### Inject:
```sh
Usage: wasm_injector inject [OPTIONS] <injection> <SOURCE> [DESTINATION]
Inject invalid instructions into a wasm module

Usage: wasm_injector inject [OPTIONS] <injection> <function> <source> [destination]

Arguments:
<injection> [possible values: infinite-loop, bad-return-value, stack-overflow, noops, heap-overflow]
<SOURCE> Wasm source file path. Can be compressed and/or hexified.
[DESTINATION] Destination file path (optional). If not specified, the output file will be a prefixed source file name.
<function> The name of the exported function to be injected with the instructions
<source> Wasm source file path. Can be compressed and/or hexified
[destination] Destination file path (optional). If not specified, the output file will be a prefixed source file name

Options:
--compressed Compresses the wasm. Can be used with `--hexified`
Expand All @@ -53,15 +56,17 @@ Options:

### Convert:
```sh
Usage: wasm_injector convert [OPTIONS] <SOURCE> [DESTINATION]
Convert from `hexified` and/or `compressed` to `raw` wasm module and vice versa

Usage: wasm_injector convert [OPTIONS] <source> [destination]

Arguments:
<SOURCE> Wasm source file path. Can be compressed and/or hexified.
[DESTINATION] Destination file path (optional). If not specified, the output file will be a prefixed source file name.
<source> Wasm source file path. Can be compressed and/or hexified.
[destination] Destination file path (optional). If not specified, the output file will be a prefixed source file name

Options:
--raw Saves the file as raw wasm (default). Can not be used with `--compressed` or `--hexified`.
--compressed Compresses the wasm (zstd compression). Can be used with `--hexified`.
--raw Saves the file as raw wasm (default). Can not be used with `--compressed` or `--hexified`
--compressed Compresses the wasm (zstd compression). Can be used with `--hexified`
--hexified Hexifies the wasm. Can be used with `--compressed`
-h, --help Print help
```
Expand All @@ -72,13 +77,13 @@ Options:
To inject code into a wasm file, compress and hexify it, you can run:

```sh
./wasm_injector inject noops my_wasm_file.wasm --compressed --hexified
./wasm_injector inject noops validate_block my_wasm_file.wasm --compressed --hexified
```

To specify a custom destination path, you can run:

```sh
./wasm_injector inject noops my_wasm_file.wasm my_destination_directory/injected_new_file.wasm
./wasm_injector inject noops validate_block my_wasm_file.wasm my_destination_directory/injected_new_file.wasm
```

### Convert:
Expand Down
66 changes: 35 additions & 31 deletions src/injecting/injections.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ pub enum Injection {
}

impl Injection {
pub fn inject(self, module: &mut Module) -> Result<(), String> {
get_injection(self)(module)
pub fn inject(self, module: &mut Module, function: &str) -> Result<(), String> {
get_injection(self)(module, function)
}
}

Expand All @@ -31,7 +31,7 @@ impl Display for Injection {
}
}
}
type InjectionFn = dyn FnMut(&mut Module) -> Result<(), String>;
type InjectionFn = dyn FnMut(&mut Module, &str) -> Result<(), String>;

pub fn get_injection(injection: Injection) -> Box<InjectionFn> {
Box::new(match injection {
Expand All @@ -43,8 +43,8 @@ pub fn get_injection(injection: Injection) -> Box<InjectionFn> {
})
}

pub fn inject_infinite_loop(module: &mut Module) -> Result<(), String> {
module.map_function("validate_block", |func_body: &mut FuncBody, _| {
pub fn inject_infinite_loop(module: &mut Module, function_name: &str) -> Result<(), String> {
module.map_function(function_name, |func_body: &mut FuncBody| {
let code = func_body.code_mut();

let mut code_with_loop = vec![
Expand All @@ -60,8 +60,8 @@ pub fn inject_infinite_loop(module: &mut Module) -> Result<(), String> {
})
}

fn inject_bad_return_value(module: &mut Module) -> Result<(), String> {
module.map_function("validate_block", |func_body: &mut FuncBody, _| {
fn inject_bad_return_value(module: &mut Module, function_name: &str) -> Result<(), String> {
module.map_function(function_name, |func_body: &mut FuncBody| {
*func_body.code_mut() = Instructions::new(vec![
// Last value on the stack gets returned
Instruction::I64Const(123456789),
Expand All @@ -70,8 +70,8 @@ fn inject_bad_return_value(module: &mut Module) -> Result<(), String> {
})
}

fn inject_stack_overflow(module: &mut Module) -> Result<(), String> {
module.map_function("validate_block", |func_body: &mut FuncBody, _| {
fn inject_stack_overflow(module: &mut Module, function_name: &str) -> Result<(), String> {
module.map_function(function_name, |func_body: &mut FuncBody| {
let code = func_body.code_mut();

let mut code_with_allocation = vec![
Expand All @@ -86,8 +86,8 @@ fn inject_stack_overflow(module: &mut Module) -> Result<(), String> {
})
}

fn inject_noops(module: &mut Module) -> Result<(), String> {
module.map_function("validate_block", |func_body: &mut FuncBody, _| {
fn inject_noops(module: &mut Module, function_name: &str) -> Result<(), String> {
module.map_function(function_name, |func_body: &mut FuncBody| {
// Add half a billion NoOperations to (hopefully) slow down interpretation-time
let code = func_body.code_mut();

Expand All @@ -98,24 +98,27 @@ fn inject_noops(module: &mut Module) -> Result<(), String> {
})
}

fn inject_heap_overflow(module: &mut Module) -> Result<(), String> {
module.map_function(
"validate_block",
|func_body: &mut FuncBody, malloc_index| {
let code = func_body.code_mut();
let index: u32 = malloc_index as u32;
fn inject_heap_overflow(module: &mut Module, function_name: &str) -> Result<(), String> {
let malloc_index = &module.get_malloc_index().expect("No malloc function");

let mut code_with_allocation =
vec![[Instruction::I32Const(33_554_431), Instruction::Call(index)]; 8]
.into_iter()
.flatten()
.collect::<Vec<Instruction>>();
module.map_function(function_name, |func_body: &mut FuncBody| {
let code = func_body.code_mut();

let mut code_with_allocation = vec![
[
Instruction::I32Const(33_554_431),
Instruction::Call(*malloc_index as u32)
];
8
]
.into_iter()
.flatten()
.collect::<Vec<Instruction>>();

code_with_allocation.append(code.elements_mut());
code_with_allocation.append(code.elements_mut());

*code.elements_mut() = code_with_allocation;
},
)
*code.elements_mut() = code_with_allocation;
})
}

#[cfg(test)]
Expand All @@ -124,6 +127,7 @@ mod injections_tests {
use crate::util::load_module_from_wasm;
use std::path::Path;

const FUNCTION_NAME: &'static str = "validate_block";
const WASM_PATH: &'static str = concat!(env!("CARGO_MANIFEST_DIR"), "/test-wasm/test.wasm");

fn get_function_body(module: &mut Module) -> &mut FuncBody {
Expand All @@ -149,7 +153,7 @@ mod injections_tests {
let mut module = load_module();

let injection = Injection::InfiniteLoop;
assert!(injection.inject(&mut module).is_ok());
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());

let function_body = get_function_body(&mut module);

Expand All @@ -167,7 +171,7 @@ mod injections_tests {
let mut module = load_module();
let injection = Injection::BadReturnValue;

assert!(injection.inject(&mut module).is_ok());
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());

let function_body = get_function_body(&mut module);

Expand All @@ -180,7 +184,7 @@ mod injections_tests {
let mut module = load_module();

let injection = Injection::StackOverflow;
assert!(injection.inject(&mut module).is_ok());
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());

let function_body = get_function_body(&mut module);

Expand All @@ -197,7 +201,7 @@ mod injections_tests {
let mut module = load_module();

let injection = Injection::Noops;
assert!(injection.inject(&mut module).is_ok());
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());

let function_body = get_function_body(&mut module);

Expand All @@ -210,7 +214,7 @@ mod injections_tests {
let mut module = load_module();

let injection = Injection::HeapOverflow;
assert!(injection.inject(&mut module).is_ok());
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());

let index = module.get_malloc_index().unwrap() as u32;
let function_body = get_function_body(&mut module);
Expand Down
12 changes: 5 additions & 7 deletions src/injecting/injector.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,12 @@ pub trait FunctionMapper {
fn map_function(
&mut self,
function_name: &str,
body_mapper: impl Fn(&mut FuncBody, usize),
body_mapper: impl Fn(&mut FuncBody),
) -> Result<(), String>;

fn map_functions(
&mut self,
function_name_body_mapper_pairs: Vec<(&str, impl Fn(&mut FuncBody, usize))>,
function_name_body_mapper_pairs: Vec<(&str, impl Fn(&mut FuncBody))>,
) -> Result<(), String> {
function_name_body_mapper_pairs
.into_iter()
Expand Down Expand Up @@ -94,10 +94,11 @@ impl FunctionMapper for Module {

Ok(malloc_index)
}

fn map_function(
&mut self,
function_name: &str,
body_mapper: impl Fn(&mut FuncBody, usize),
body_mapper: impl Fn(&mut FuncBody),
) -> Result<(), String> {
// NOTE:
// Indexing for local functions includes both imported and own (local)
Expand All @@ -110,9 +111,6 @@ impl FunctionMapper for Module {
// Count the number of imported functions listed by the module
let import_section_len: usize = self.get_import_section_len()?;

// get the global function index of the dynamic memory allocation method
let malloc_index = self.get_malloc_index()?;

// Compute the local function index (for the code section) by subtracting
// the number of imported functions from the global function index
let local_function_index = global_function_index - import_section_len;
Expand All @@ -121,7 +119,7 @@ impl FunctionMapper for Module {
let function_body = self.get_function_body(local_function_index, function_name)?;

// Map over the function_body
body_mapper(function_body, malloc_index);
body_mapper(function_body);

Ok(())
}
Expand Down
Loading

0 comments on commit 57ca2f7

Please sign in to comment.