Skip to content

Commit 57ca2f7

Browse files
dimitrovmaksimdimitrovmaksim
authored
Add function arg to cli (#20)
* refactor: Add function arg. Remove hardcoded validate_block. Move malloc index to heap injection only Signed-off-by: Maksim Dimitrov <dimitrov.maksim@gmail.com> * chore(readme): Update readme Signed-off-by: Maksim Dimitrov <dimitrov.maksim@gmail.com> * chore: Run fmt Signed-off-by: Maksim Dimitrov <dimitrov.maksim@gmail.com> * chore: Update readme Signed-off-by: Maksim Dimitrov <dimitrov.maksim@gmail.com> --------- Signed-off-by: Maksim Dimitrov <dimitrov.maksim@gmail.com>
1 parent 2817b0b commit 57ca2f7

4 files changed

Lines changed: 150 additions & 72 deletions

File tree

‎README.md‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,15 @@ Options:
3838

3939
### Inject:
4040
```sh
41-
Usage: wasm_injector inject [OPTIONS] <injection> <SOURCE> [DESTINATION]
41+
Inject invalid instructions into a wasm module
42+
43+
Usage: wasm_injector inject [OPTIONS] <injection> <function> <source> [destination]
4244

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

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

5457
### Convert:
5558
```sh
56-
Usage: wasm_injector convert [OPTIONS] <SOURCE> [DESTINATION]
59+
Convert from `hexified` and/or `compressed` to `raw` wasm module and vice versa
60+
61+
Usage: wasm_injector convert [OPTIONS] <source> [destination]
5762

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

6267
Options:
63-
--raw Saves the file as raw wasm (default). Can not be used with `--compressed` or `--hexified`.
64-
--compressed Compresses the wasm (zstd compression). Can be used with `--hexified`.
68+
--raw Saves the file as raw wasm (default). Can not be used with `--compressed` or `--hexified`
69+
--compressed Compresses the wasm (zstd compression). Can be used with `--hexified`
6570
--hexified Hexifies the wasm. Can be used with `--compressed`
6671
-h, --help Print help
6772
```
@@ -72,13 +77,13 @@ Options:
7277
To inject code into a wasm file, compress and hexify it, you can run:
7378

7479
```sh
75-
./wasm_injector inject noops my_wasm_file.wasm --compressed --hexified
80+
./wasm_injector inject noops validate_block my_wasm_file.wasm --compressed --hexified
7681
```
7782

7883
To specify a custom destination path, you can run:
7984

8085
```sh
81-
./wasm_injector inject noops my_wasm_file.wasm my_destination_directory/injected_new_file.wasm
86+
./wasm_injector inject noops validate_block my_wasm_file.wasm my_destination_directory/injected_new_file.wasm
8287
```
8388

8489
### Convert:

‎src/injecting/injections.rs‎

Lines changed: 35 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@ pub enum Injection {
1515
}
1616

1717
impl Injection {
18-
pub fn inject(self, module: &mut Module) -> Result<(), String> {
19-
get_injection(self)(module)
18+
pub fn inject(self, module: &mut Module, function: &str) -> Result<(), String> {
19+
get_injection(self)(module, function)
2020
}
2121
}
2222

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

3636
pub fn get_injection(injection: Injection) -> Box<InjectionFn> {
3737
Box::new(match injection {
@@ -43,8 +43,8 @@ pub fn get_injection(injection: Injection) -> Box<InjectionFn> {
4343
})
4444
}
4545

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

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

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

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

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

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

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

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

108-
let mut code_with_allocation =
109-
vec![[Instruction::I32Const(33_554_431), Instruction::Call(index)]; 8]
110-
.into_iter()
111-
.flatten()
112-
.collect::<Vec<Instruction>>();
104+
module.map_function(function_name, |func_body: &mut FuncBody| {
105+
let code = func_body.code_mut();
106+
107+
let mut code_with_allocation = vec![
108+
[
109+
Instruction::I32Const(33_554_431),
110+
Instruction::Call(*malloc_index as u32)
111+
];
112+
8
113+
]
114+
.into_iter()
115+
.flatten()
116+
.collect::<Vec<Instruction>>();
113117

114-
code_with_allocation.append(code.elements_mut());
118+
code_with_allocation.append(code.elements_mut());
115119

116-
*code.elements_mut() = code_with_allocation;
117-
},
118-
)
120+
*code.elements_mut() = code_with_allocation;
121+
})
119122
}
120123

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

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

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

151155
let injection = Injection::InfiniteLoop;
152-
assert!(injection.inject(&mut module).is_ok());
156+
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());
153157

154158
let function_body = get_function_body(&mut module);
155159

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

170-
assert!(injection.inject(&mut module).is_ok());
174+
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());
171175

172176
let function_body = get_function_body(&mut module);
173177

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

182186
let injection = Injection::StackOverflow;
183-
assert!(injection.inject(&mut module).is_ok());
187+
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());
184188

185189
let function_body = get_function_body(&mut module);
186190

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

199203
let injection = Injection::Noops;
200-
assert!(injection.inject(&mut module).is_ok());
204+
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());
201205

202206
let function_body = get_function_body(&mut module);
203207

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

212216
let injection = Injection::HeapOverflow;
213-
assert!(injection.inject(&mut module).is_ok());
217+
assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok());
214218

215219
let index = module.get_malloc_index().unwrap() as u32;
216220
let function_body = get_function_body(&mut module);

‎src/injecting/injector.rs‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@ pub trait FunctionMapper {
44
fn map_function(
55
&mut self,
66
function_name: &str,
7-
body_mapper: impl Fn(&mut FuncBody, usize),
7+
body_mapper: impl Fn(&mut FuncBody),
88
) -> Result<(), String>;
99

1010
fn map_functions(
1111
&mut self,
12-
function_name_body_mapper_pairs: Vec<(&str, impl Fn(&mut FuncBody, usize))>,
12+
function_name_body_mapper_pairs: Vec<(&str, impl Fn(&mut FuncBody))>,
1313
) -> Result<(), String> {
1414
function_name_body_mapper_pairs
1515
.into_iter()
@@ -94,10 +94,11 @@ impl FunctionMapper for Module {
9494

9595
Ok(malloc_index)
9696
}
97+
9798
fn map_function(
9899
&mut self,
99100
function_name: &str,
100-
body_mapper: impl Fn(&mut FuncBody, usize),
101+
body_mapper: impl Fn(&mut FuncBody),
101102
) -> Result<(), String> {
102103
// NOTE:
103104
// Indexing for local functions includes both imported and own (local)
@@ -110,9 +111,6 @@ impl FunctionMapper for Module {
110111
// Count the number of imported functions listed by the module
111112
let import_section_len: usize = self.get_import_section_len()?;
112113

113-
// get the global function index of the dynamic memory allocation method
114-
let malloc_index = self.get_malloc_index()?;
115-
116114
// Compute the local function index (for the code section) by subtracting
117115
// the number of imported functions from the global function index
118116
let local_function_index = global_function_index - import_section_len;
@@ -121,7 +119,7 @@ impl FunctionMapper for Module {
121119
let function_body = self.get_function_body(local_function_index, function_name)?;
122120

123121
// Map over the function_body
124-
body_mapper(function_body, malloc_index);
122+
body_mapper(function_body);
125123

126124
Ok(())
127125
}

0 commit comments

Comments
 (0)