From 501d3210587d93f8a9408b0fa0f7c2c2b0ad9204 Mon Sep 17 00:00:00 2001 From: Maksim Dimitrov Date: Tue, 15 Aug 2023 14:47:58 +0300 Subject: [PATCH 1/2] refactor(cli): Add size option to noops injection Signed-off-by: Maksim Dimitrov --- src/injecting/injections.rs | 52 +++++++++++++++++-------------- src/lib.rs | 2 +- src/main.rs | 62 ++++++++++++++++++++++++++++++++++--- 3 files changed, 87 insertions(+), 29 deletions(-) diff --git a/src/injecting/injections.rs b/src/injecting/injections.rs index 028a9be..bbca3f7 100644 --- a/src/injecting/injections.rs +++ b/src/injecting/injections.rs @@ -18,7 +18,7 @@ use super::injector::FunctionMapper; /// let source = Path::new("samples/example.wasm"); /// let mut module = load_module_from_wasm(source)?; /// let injection = Injection::InfiniteLoop; -/// injection.inject(&mut module, "validate_block")?; +/// injection.inject(&mut module, "validate_block", None)?; /// # Ok(()) /// # } /// ``` @@ -34,8 +34,19 @@ pub enum Injection { impl Injection { /// # Takes a module and injects the selected injection into the module. - pub fn inject(self, module: &mut Module, function: &str) -> Result<(), String> { - get_injection(self)(module, function) + pub fn inject( + self, + module: &mut Module, + function: &str, + size: Option, + ) -> Result<(), String> { + match self { + Injection::InfiniteLoop => inject_infinite_loop(module, function), + Injection::BadReturnValue => inject_bad_return_value(module, function), + Injection::StackOverflow => inject_stack_overflow(module, function), + Injection::Noops => inject_noops(module, function, size), + Injection::HeapOverflow => inject_heap_overflow(module, function), + } } } @@ -50,18 +61,6 @@ impl Display for Injection { } } } -type InjectionFn = dyn FnMut(&mut Module, &str) -> Result<(), String>; - -/// # Takes an injection and returns a function that performs the injection. -fn get_injection(injection: Injection) -> Box { - Box::new(match injection { - Injection::InfiniteLoop => inject_infinite_loop, - Injection::BadReturnValue => inject_bad_return_value, - Injection::StackOverflow => inject_stack_overflow, - Injection::Noops => inject_noops, - Injection::HeapOverflow => inject_heap_overflow, - }) -} /// # Takes a module and injects an infinite loop in the beginning of the module. fn inject_infinite_loop(module: &mut Module, function_name: &str) -> Result<(), String> { @@ -109,13 +108,14 @@ fn inject_stack_overflow(module: &mut Module, function_name: &str) -> Result<(), }) } -/// # Takes a module and injects 500 million NoOperations in the beginning of the module. -fn inject_noops(module: &mut Module, function_name: &str) -> Result<(), String> { +/// # Takes a module and injects specified size in MB of NoOperations in the beginning of the module. +fn inject_noops(module: &mut Module, function_name: &str, size: Option) -> 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(); + let size = size.expect("No size given"); + let nops_count = (size as usize) * 1000 * 1000; - let mut nops = vec![Instruction::Nop; 500_000_000]; + let mut nops = vec![Instruction::Nop; nops_count]; nops.append(code.elements_mut()); *code.elements_mut() = nops; @@ -123,6 +123,8 @@ fn inject_noops(module: &mut Module, function_name: &str) -> Result<(), String> } /// # Takes a module and tries to allocate a lot of memory multiple times in the beginning of the module. +/// # This injection relies on the existnce of the `ext_allocator_malloc` function as per +/// # Polkadot's specification https://spec.polkadot.network/chap-host-api#id-ext_allocator_malloc fn inject_heap_overflow(module: &mut Module, function_name: &str) -> Result<(), String> { let malloc_index = &module.get_malloc_index().expect("No malloc function"); @@ -178,7 +180,7 @@ mod injections_tests { let mut module = load_module(); let injection = Injection::InfiniteLoop; - assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok()); + assert!(injection.inject(&mut module, FUNCTION_NAME, None).is_ok()); let function_body = get_function_body(&mut module); @@ -196,7 +198,7 @@ mod injections_tests { let mut module = load_module(); let injection = Injection::BadReturnValue; - assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok()); + assert!(injection.inject(&mut module, FUNCTION_NAME, None).is_ok()); let function_body = get_function_body(&mut module); @@ -209,7 +211,7 @@ mod injections_tests { let mut module = load_module(); let injection = Injection::StackOverflow; - assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok()); + assert!(injection.inject(&mut module, FUNCTION_NAME, None).is_ok()); let function_body = get_function_body(&mut module); @@ -226,7 +228,9 @@ mod injections_tests { let mut module = load_module(); let injection = Injection::Noops; - assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok()); + assert!(injection + .inject(&mut module, FUNCTION_NAME, Some(10)) + .is_ok()); let function_body = get_function_body(&mut module); @@ -239,7 +243,7 @@ mod injections_tests { let mut module = load_module(); let injection = Injection::HeapOverflow; - assert!(injection.inject(&mut module, FUNCTION_NAME).is_ok()); + assert!(injection.inject(&mut module, FUNCTION_NAME, None).is_ok()); let index = module.get_malloc_index().unwrap() as u32; let function_body = get_function_body(&mut module); diff --git a/src/lib.rs b/src/lib.rs index c80b2c7..178ef9b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,7 +15,7 @@ //! //! let mut module = load_module_from_wasm(source)?; // supply your own path here //! let injection = Injection::StackOverflow; // choose your injection -//! injection.inject(&mut module, "validate_block")?; // inject the instruction into the specified wasm export function +//! injection.inject(&mut module, "validate_block", None)?; // inject the instruction into the specified wasm export function //! //! save_module_to_wasm(module, destination, compressed, hexified)?; // save the module in your destination. You can choose to compress and/or hexify the module. //! diff --git a/src/main.rs b/src/main.rs index 8fc3fe3..bb58bb5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,6 @@ -use clap::{builder::ArgPredicate, Parser, Subcommand, ValueHint}; +use clap::{ + builder::ArgPredicate, error::ErrorKind, CommandFactory, Parser, Subcommand, ValueHint, +}; use std::path::PathBuf; use wasm_injector::injecting::injections::Injection; use wasm_injector::util::{load_module_from_wasm, modify_file_name, save_module_to_wasm}; @@ -14,12 +16,20 @@ struct Cli { enum Action { #[command(about = "Inject invalid instructions into a wasm module")] Inject { - #[arg(required = true, value_name = "injection",value_hint = ValueHint::Other)] + #[arg(value_enum, required = true, requires_if("noops", "size"), value_name = "injection",value_hint = ValueHint::Other)] injection: Injection, #[arg(required = true, value_name = "function", help = "The name of the exported function to be injected with the instructions", value_hint = ValueHint::Other)] function: String, + #[arg( + long, + value_name = "size", + help = "The number of noops to be injected in MB (1 NOP = 1 byte)", + value_hint = ValueHint::Other + )] + size: Option, + #[command(flatten)] global_opts: GlobalOpts, @@ -89,6 +99,25 @@ struct GlobalOpts { fn main() -> Result<(), String> { let Cli { action } = Cli::parse(); + if let Action::Inject { + injection, size, .. + } = &action + { + match injection { + Injection::Noops => {} + _ => { + if size.is_some() { + let mut cmd = Cli::command(); + cmd.error( + ErrorKind::ArgumentConflict, + "The `size` argument is only valid for the `noops` injection".to_string(), + ) + .exit(); + } + } + } + } + let calculate_default_destination_file_name = |file_name: &str| { let mut file_name = String::from(file_name); @@ -161,11 +190,12 @@ fn main() -> Result<(), String> { if let Action::Inject { injection, function, + size, .. } = action { // Inject the module - injection.inject(&mut module, &function)?; + injection.inject(&mut module, &function, size)?; } save_module_to_wasm(module, destination.as_path(), compressed, hexified)?; @@ -197,10 +227,20 @@ mod cli_tests { #[test] fn test_inject_noops() { assert_eq!( - Cli::try_parse_from(&["test", "inject", "noops", FUNCTION_NAME, "test.wasm"]).unwrap(), + Cli::try_parse_from(&[ + "test", + "inject", + "noops", + "--size", + "20", + FUNCTION_NAME, + "test.wasm" + ]) + .unwrap(), Cli { action: Action::Inject { injection: Injection::Noops, + size: Some(20), function: FUNCTION_NAME.to_string(), global_opts: GlobalOpts { source: PathBuf::from("test.wasm"), @@ -213,6 +253,16 @@ mod cli_tests { ) } + #[test] + fn test_inject_noops_requires_size_arg() { + let result = Cli::try_parse_from(&["test", "inject", "noops", FUNCTION_NAME, "test.wasm"]); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().kind(), + clap::error::ErrorKind::MissingRequiredArgument + ) + } + #[test] fn test_inject_heap_overflow() { assert_eq!( @@ -228,6 +278,7 @@ mod cli_tests { action: Action::Inject { injection: Injection::HeapOverflow, function: FUNCTION_NAME.to_string(), + size: None, global_opts: GlobalOpts { source: PathBuf::from("test.wasm"), destination: None @@ -254,6 +305,7 @@ mod cli_tests { action: Action::Inject { injection: Injection::StackOverflow, function: FUNCTION_NAME.to_string(), + size: None, global_opts: GlobalOpts { source: PathBuf::from("test.wasm"), destination: None @@ -280,6 +332,7 @@ mod cli_tests { action: Action::Inject { injection: Injection::BadReturnValue, function: FUNCTION_NAME.to_string(), + size: None, global_opts: GlobalOpts { source: PathBuf::from("test.wasm"), destination: None @@ -306,6 +359,7 @@ mod cli_tests { action: Action::Inject { injection: Injection::InfiniteLoop, function: FUNCTION_NAME.to_string(), + size: None, global_opts: GlobalOpts { source: PathBuf::from("test.wasm"), destination: None From 6dcf3f0d0805f6cc02d1c294efa0abbf6ee3f2cb Mon Sep 17 00:00:00 2001 From: Mihail Kirov Date: Thu, 17 Aug 2023 13:05:52 +0300 Subject: [PATCH 2/2] fix: 2**10 in kb and 2**10 in mb --- src/injecting/injections.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/injecting/injections.rs b/src/injecting/injections.rs index bbca3f7..bb06983 100644 --- a/src/injecting/injections.rs +++ b/src/injecting/injections.rs @@ -113,7 +113,7 @@ fn inject_noops(module: &mut Module, function_name: &str, size: Option) -> module.map_function(function_name, |func_body: &mut FuncBody| { let code = func_body.code_mut(); let size = size.expect("No size given"); - let nops_count = (size as usize) * 1000 * 1000; + let nops_count = (size as usize) * 1024 * 1024; let mut nops = vec![Instruction::Nop; nops_count]; nops.append(code.elements_mut());