AFL_CUSTOM_MUTATOR_LATE_SEND added

10 files changed, 82 insertions, 12 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 09ea8cb6..c16214e4 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -11,6 +11,8 @@
       to disable this feature set `AFL_NO_FASTRESUME=1`
       zlib compression is used if zlib is found at compile time
     - improved seed selection algorithm
+    - added `AFL_CUSTOM_MUTATOR_LATE_SEND=1` to call the custom send()
+      function after the target has been restarted.
   - frida_mode:
     - AFL_FRIDA_PERSISTENT_ADDR can now be be any reachable address not just
       a function entry
diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index b7a7032f..3067ceab 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -198,6 +198,11 @@ def deinit():  # optional for Python
     This method can be used if you want to send data to the target yourself,
     e.g. via IPC. This replaces some usage of utils/afl_proxy but requires
     that you start the target with afl-fuzz.
+
+    Setting `AFL_CUSTOM_MUTATOR_LATE_SEND` will call the afl_custom_fuzz_send()
+    function after the target has been restarted. (This is needed for e.g. TCP
+    services.)
+
     Example: [custom_mutators/examples/custom_send.c](../custom_mutators/examples/custom_send.c)
 
 - `queue_new_entry` (optional):
diff --git a/docs/env_variables.md b/docs/env_variables.md
index 6db31df0..eebbcbda 100644
--- a/docs/env_variables.md
+++ b/docs/env_variables.md
@@ -368,6 +368,10 @@ checks or alter some of the more exotic semantics of the tool:
     XML or other highly flexible structured input. For details, see
     [custom_mutators.md](custom_mutators.md).
 
+  - Setting `AFL_CUSTOM_MUTATOR_LATE_SEND` will call the afl_custom_fuzz_send()
+    function after the target has been restarted. (This is needed for e.g. TCP
+    services.)
+
   - Setting `AFL_CYCLE_SCHEDULES` will switch to a different schedule every time
     a cycle is finished.
 
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 45600698..0f0e45d3 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -449,8 +449,9 @@ extern char *power_names[POWER_SCHEDULES_NUM];
 typedef struct afl_env_vars {
 
   u8 afl_skip_cpufreq, afl_exit_when_done, afl_no_affinity, afl_skip_bin_check,
-      afl_dumb_forksrv, afl_import_first, afl_custom_mutator_only, afl_no_ui,
-      afl_force_ui, afl_i_dont_care_about_missing_crashes, afl_bench_just_one,
+      afl_dumb_forksrv, afl_import_first, afl_custom_mutator_only,
+      afl_custom_mutator_late_send, afl_no_ui, afl_force_ui,
+      afl_i_dont_care_about_missing_crashes, afl_bench_just_one,
       afl_bench_until_crash, afl_debug_child, afl_autoresume, afl_cal_fast,
       afl_cycle_schedules, afl_expand_havoc, afl_statsd, afl_cmplog_only_new,
       afl_exit_on_seed_issues, afl_try_affinity, afl_ignore_problems,
diff --git a/include/envs.h b/include/envs.h
index ef522ab4..6a0d329b 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -24,15 +24,16 @@ static char *afl_environment_variables[] = {
     "AFL_DUMP_CYCLOMATIC_COMPLEXITY", "AFL_CMPLOG_MAX_LEN", "AFL_COMPCOV_LEVEL",
     "AFL_CRASH_EXITCODE", "AFL_CRASHING_SEEDS_AS_NEW_CRASH",
     "AFL_CUSTOM_MUTATOR_LIBRARY", "AFL_CUSTOM_MUTATOR_ONLY",
-    "AFL_CUSTOM_INFO_PROGRAM", "AFL_CUSTOM_INFO_PROGRAM_ARGV",
-    "AFL_CUSTOM_INFO_PROGRAM_INPUT", "AFL_CUSTOM_INFO_OUT", "AFL_CXX",
-    "AFL_CYCLE_SCHEDULES", "AFL_DEBUG", "AFL_DEBUG_CHILD", "AFL_DEBUG_GDB",
-    "AFL_DEBUG_UNICORN", "AFL_DISABLE_REDUNDANT", "AFL_NO_REDUNDANT",
-    "AFL_DISABLE_TRIM", "AFL_NO_TRIM", "AFL_DISABLE_LLVM_INSTRUMENTATION",
-    "AFL_DONT_OPTIMIZE", "AFL_DRIVER_STDERR_DUPLICATE_FILENAME",
-    "AFL_DUMB_FORKSRV", "AFL_EARLY_FORKSERVER", "AFL_ENTRYPOINT",
-    "AFL_EXIT_WHEN_DONE", "AFL_EXIT_ON_TIME", "AFL_EXIT_ON_SEED_ISSUES",
-    "AFL_FAST_CAL", "AFL_FINAL_SYNC", "AFL_FORCE_UI", "AFL_FRIDA_DEBUG_MAPS",
+    "AFL_CUSTOM_MUTATOR_LATE_SEND", "AFL_CUSTOM_INFO_PROGRAM",
+    "AFL_CUSTOM_INFO_PROGRAM_ARGV", "AFL_CUSTOM_INFO_PROGRAM_INPUT",
+    "AFL_CUSTOM_INFO_OUT", "AFL_CXX", "AFL_CYCLE_SCHEDULES", "AFL_DEBUG",
+    "AFL_DEBUG_CHILD", "AFL_DEBUG_GDB", "AFL_DEBUG_UNICORN",
+    "AFL_DISABLE_REDUNDANT", "AFL_NO_REDUNDANT", "AFL_DISABLE_TRIM",
+    "AFL_NO_TRIM", "AFL_DISABLE_LLVM_INSTRUMENTATION", "AFL_DONT_OPTIMIZE",
+    "AFL_DRIVER_STDERR_DUPLICATE_FILENAME", "AFL_DUMB_FORKSRV",
+    "AFL_EARLY_FORKSERVER", "AFL_ENTRYPOINT", "AFL_EXIT_WHEN_DONE",
+    "AFL_EXIT_ON_TIME", "AFL_EXIT_ON_SEED_ISSUES", "AFL_FAST_CAL",
+    "AFL_FINAL_SYNC", "AFL_FORCE_UI", "AFL_FRIDA_DEBUG_MAPS",
     "AFL_FRIDA_DRIVER_NO_HOOK", "AFL_FRIDA_EXCLUDE_RANGES",
     "AFL_FRIDA_INST_CACHE_SIZE", "AFL_FRIDA_INST_COVERAGE_ABSOLUTE",
     "AFL_FRIDA_INST_COVERAGE_FILE", "AFL_FRIDA_INST_DEBUG_FILE",
diff --git a/include/forkserver.h b/include/forkserver.h
index 3fd813a4..6c649528 100644
--- a/include/forkserver.h
+++ b/include/forkserver.h
@@ -210,6 +210,11 @@ typedef struct afl_forkserver {
   u8 *persistent_trace_bits;                   /* Persistent copy of bitmap */
 #endif
 
+  void *custom_data_ptr;
+  u8   *custom_input;
+  u32   custom_input_len;
+  void (*late_send)(void *, const u8 *, size_t);
+
 } afl_forkserver_t;
 
 typedef enum fsrv_run_result {
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index a998c10f..cec91f76 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -292,6 +292,9 @@ void afl_fsrv_init_dup(afl_forkserver_t *fsrv_to, afl_forkserver_t *from) {
   fsrv_to->use_fauxsrv = 0;
   fsrv_to->last_run_timed_out = 0;
 
+  fsrv_to->late_send = from->late_send;
+  fsrv_to->custom_data_ptr = from->custom_data_ptr;
+
   fsrv_to->init_child_func = from->init_child_func;
   // Note: do not copy ->add_extra_func or ->persistent_record*
 
@@ -1952,6 +1955,13 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
 
   }
 
+  if (unlikely(fsrv->late_send)) {
+
+    fsrv->late_send(fsrv->custom_data_ptr, fsrv->custom_input,
+                    fsrv->custom_input_len);
+
+  }
+
   exec_ms = read_s32_timed(fsrv->fsrv_st_fd, &fsrv->child_status, timeout,
                            stop_soon_p);
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index c234fc42..2f244a1d 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -194,7 +194,17 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
 
       if (el->afl_custom_fuzz_send) {
 
-        el->afl_custom_fuzz_send(el->data, *mem, new_size);
+        if (!afl->afl_env.afl_custom_mutator_late_send) {
+
+          el->afl_custom_fuzz_send(el->data, *mem, new_size);
+
+        } else {
+
+          afl->fsrv.custom_input = *mem;
+          afl->fsrv.custom_input_len = new_size;
+
+        }
+
         sent = 1;
 
       }
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index dd684a19..eead3e50 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -300,6 +300,13 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
             afl->afl_env.afl_custom_mutator_only =
                 get_afl_env(afl_environment_variables[i]) ? 1 : 0;
 
+          } else if (!strncmp(env, "AFL_CUSTOM_MUTATOR_LATE_SEND",
+
+                              afl_environment_variable_len)) {
+
+            afl->afl_env.afl_custom_mutator_late_send =
+                get_afl_env(afl_environment_variables[i]) ? 1 : 0;
+
           } else if (!strncmp(env, "AFL_CMPLOG_ONLY_NEW",
 
                               afl_environment_variable_len)) {
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index b53a9a2d..8a84d447 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -2142,6 +2142,31 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
+  if (afl->custom_mutators_count && afl->afl_env.afl_custom_mutator_late_send) {
+
+    u32 count_send = 0;
+    LIST_FOREACH(&afl->custom_mutator_list, struct custom_mutator, {
+
+      if (el->afl_custom_fuzz_send) {
+
+        if (count_send) {
+
+          FATAL(
+              "You can only have one custom send() function if you are using "
+              "AFL_CUSTOM_MUTATOR_LATE_SEND!");
+
+        }
+
+        afl->fsrv.late_send = el->afl_custom_fuzz_send;
+        afl->fsrv.custom_data_ptr = el->data;
+        count_send = 1;
+
+      }
+
+    });
+
+  }
+
   if (afl->limit_time_sig > 0 && afl->custom_mutators_count) {
 
     if (afl->custom_only) {