refactor(cli): Add size option to noops injection

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<i16>,
+    ) -> 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<InjectionFn> {
-    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,20 +108,23 @@ 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<i16>) -> 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) * 1024 * 1024;
 
-        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;
     })
 }
 
 /// # 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);

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.
 //!     

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<i16>,
+
         #[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