From 5742988f8e9b7718f425de9745bf758a22d2d0ea Mon Sep 17 00:00:00 2001
From: "Samuel E. Moelius III" <sam@moeli.us>
Date: Tue, 14 Sep 2021 21:11:03 -0400
Subject: [PATCH] Add `execute_with` option

---
 README.md        |  7 +++++++
 macro/src/lib.rs | 30 ++++++++++++++++++++----------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index abaa6841..2aaef8e0 100644
--- a/README.md
+++ b/README.md
@@ -139,6 +139,13 @@ The primary effects of the `test_fuzz` macro are:
 
   **WARNING**: Setting `enable_in_production` could introduce a denial-of-service vector. For example, setting this option for a function that is called many times with different arguments could fill up the disk. The check of [`TEST_FUZZ_WRITE`](#environment-variables) is meant to provide some defense against this possibility. Nonetheless, consider this option carefully before using it.
 
+- **`execute_with = "function"`** - Rather than call the target directly:
+
+  - construct a closure of type `FnOnce() -> R`, where `R` is the target's return type, so that calling the closure calls the target;
+  - call `function` with the closure.
+
+  Calling the target in this way allows `function` to set up the call's environment. This can be useful, e.g., for fuzzing [Substrate externalities](https://substrate.dev/docs/en/knowledgebase/runtime/tests#mock-runtime-storage).
+
 - **`no_auto_generate`** - Do not try to [auto-generate corpus files](#auto-generated-corpus-files) for the target.
 
 - **`only_concretizations`** - Record the target's concretizations when running tests, but do not generate corpus files and do not implement a fuzzing harness. This can be useful when the target is a generic function, but it is unclear what type parameters should be used for fuzzing.
diff --git a/macro/src/lib.rs b/macro/src/lib.rs
index 2b81e8f6..909040e9 100644
--- a/macro/src/lib.rs
+++ b/macro/src/lib.rs
@@ -12,11 +12,11 @@ use quote::{quote, ToTokens};
 use std::{collections::BTreeMap, env::var, io::Write, str::FromStr};
 use subprocess::{Exec, Redirection};
 use syn::{
-    parse::Parser, parse_macro_input, parse_quote, punctuated::Punctuated, token, Attribute,
-    AttributeArgs, Block, Expr, FnArg, GenericArgument, GenericMethodArgument, GenericParam,
-    Generics, Ident, ImplItem, ImplItemMethod, ItemFn, ItemImpl, ItemMod, Pat, Path, PathArguments,
-    PathSegment, ReturnType, Signature, Stmt, Type, TypePath, TypeReference, Visibility,
-    WhereClause, WherePredicate,
+    parse::Parser, parse_macro_input, parse_quote, parse_str, punctuated::Punctuated, token,
+    Attribute, AttributeArgs, Block, Expr, FnArg, GenericArgument, GenericMethodArgument,
+    GenericParam, Generics, Ident, ImplItem, ImplItemMethod, ItemFn, ItemImpl, ItemMod, Pat, Path,
+    PathArguments, PathSegment, ReturnType, Signature, Stmt, Type, TypePath, TypeReference,
+    Visibility, WhereClause, WherePredicate,
 };
 use toolchain_find::find_installed_component;
 use unzip_n::unzip_n;
@@ -166,6 +166,8 @@ struct TestFuzzOpts {
     #[darling(default)]
     enable_in_production: bool,
     #[darling(default)]
+    execute_with: Option<String>,
+    #[darling(default)]
     no_auto_generate: bool,
     #[darling(default)]
     only_concretizations: bool,
@@ -507,20 +509,28 @@ fn map_method_or_fn(
             )
         }
     };
-    let call_with_deserialized_arguments = {
+    let call_in_environment = if let Some(s) = &opts.execute_with {
+        let execute_with: Expr = parse_str(s).expect("Could not parse `execute_with` argument");
+        parse_quote! {
+            #execute_with (|| #call)
+        }
+    } else {
+        call
+    };
+    let call_in_environment_with_deserialized_arguments = {
         #[cfg(feature = "__persistent")]
         quote! {
             test_fuzz::afl::fuzz!(|data: &[u8]| {
                 let mut args = UsingReader::<_>::read_args #combined_concretization (data);
                 let ret: Option<<Args #combined_concretization as HasRetTy>::RetTy> = args.map(|mut args|
-                    #call
+                    #call_in_environment
                 );
             });
         }
         #[cfg(not(feature = "__persistent"))]
         quote! {
             let ret: Option<<Args #combined_concretization as HasRetTy>::RetTy> = args.map(|mut args|
-                #call
+                #call_in_environment
             );
         }
     };
@@ -644,13 +654,13 @@ fn map_method_or_fn(
                                     #output_args
                                 }
                                 if test_fuzz::runtime::replay_enabled() {
-                                    #call_with_deserialized_arguments
+                                    #call_in_environment_with_deserialized_arguments
                                     #output_ret
                                 }
                             } else {
                                 std::panic::set_hook(std::boxed::Box::new(|_| std::process::abort()));
                                 #input_args
-                                #call_with_deserialized_arguments
+                                #call_in_environment_with_deserialized_arguments
                                 let _ = std::panic::take_hook();
                             }
                         }