Pass command line options rather than env vars

Author: GJDuck

afl-rt.c

@@ -209,7 +209,7 @@ void init(int argc, const char **argv, char **envp, void *_unused,
         /*
          * This is a shared library.  For this, we set up a dummy area so the
          * instrumentation does not crash during program initialization.  The
-         * main executable is repsonsible for setting up AFL proper.
+         * main executable is responsible for setting up AFL proper.
          */
         (void)mmap(AREA_BASE, AREA_SIZE,
                 PROT_READ | PROT_WRITE,

e9AFLPlugin.cpp

@@ -30,6 +30,8 @@
 #include <set>
 #include <vector>
 
+#include <getopt.h>
+
 #include "e9plugin.h"
 
 using namespace e9tool;
@@ -47,7 +49,6 @@ enum Option
     OPTION_ALWAYS
 };
 static Option option_debug      = OPTION_DEFAULT;
-static Option option_instrument = OPTION_DEFAULT;
 static Option option_Oselect    = OPTION_DEFAULT;
 static Option option_Oblock     = OPTION_DEFAULT;
 
@@ -113,11 +114,68 @@ typedef std::map<intptr_t, unsigned> Ids;
  */
 static std::set<intptr_t> instrument;
 
+/*
+ * Options.
+ */
+enum
+{
+    OPTION_COUNTER,
+    OPTION_OBLOCK,
+    OPTION_OSELECT,
+    OPTION_DEBUG,
+    OPTION_PATH,
+};
+
 /*
  * Initialization.
  */
 extern void *e9_plugin_init_v1(const Context *cxt)
 {
+    static const struct option long_options[] =
+    {
+        {"counter", required_argument, nullptr, OPTION_COUNTER},
+        {"Oblock",  required_argument, nullptr, OPTION_OBLOCK},
+        {"Oselect", required_argument, nullptr, OPTION_OSELECT},
+        {"debug",   no_argument,       nullptr, OPTION_DEBUG},
+        {"path",    required_argument, nullptr, OPTION_PATH},
+        {nullptr,   no_argument,       nullptr, 0}
+    };
+    std::string option_path(".");
+    Counter option_counter = COUNTER_CLASSIC;
+    optind = 1;
+    char * const *argv = cxt->argv->data();
+    int argc = (int)cxt->argv->size();
+    while (true)
+    {
+        int idx;
+        int opt = getopt_long_only(argc, argv, "Po:v", long_options, &idx);
+        if (opt < 0)
+            break;
+        switch (opt)
+        {
+            case OPTION_COUNTER:
+                option_counter = parseCounter(optarg);
+                break;
+            case OPTION_OBLOCK:
+                option_Oblock = parseOption(optarg);
+                break;
+            case OPTION_OSELECT:
+                option_Oselect = parseOption(optarg);
+                break;
+            case OPTION_DEBUG:
+                option_debug = OPTION_ALWAYS;
+                break;
+            case OPTION_PATH:
+                option_path = optarg;
+                break;
+            default:
+                error("invalid command-line options for %s", argv[0]);
+        }
+    }
+    if (option_Oblock == OPTION_ALWAYS)
+        warning("always removing AFL instrumentation for bad blocks; coverage "
+            "may be incomplete");
+
     // Make seed depend on filename.
     unsigned seed = 0;
     const char *filename = getELFFilename(cxt->elf);
@@ -132,28 +190,6 @@ extern void *e9_plugin_init_v1(const Context *cxt)
     // Reserve memory used by the afl_area_ptr:
     sendReserveMessage(cxt->out, afl_area_ptr, AREA_SIZE, /*absolute=*/true);
 
-    const char *str = nullptr;
-    std::string option_path(".");
-    Counter option_counter = COUNTER_CLASSIC;
-    if ((str = getenv("E9AFL_COUNTER")) != nullptr)
-        option_counter = parseCounter(str);
-    if ((str = getenv("E9AFL_DEBUG")) != nullptr)
-        option_debug = parseOption(str);
-    if ((str = getenv("E9AFL_INSTRUMENT")) != nullptr)
-        option_instrument = parseOption(str);
-    if ((str = getenv("E9AFL_OBLOCK")) != nullptr)
-        option_Oblock = parseOption(str);
-    if ((str = getenv("E9AFL_OSELECT")) != nullptr)
-        option_Oselect = parseOption(str);
-    if ((str = getenv("E9AFL_PATH")) != nullptr)
-        option_path = str;
- 
-    if (option_instrument == OPTION_NEVER)
-        return nullptr;
-    if (option_Oblock == OPTION_ALWAYS)
-        warning("always removing AFL instrumentation for bad blocks; coverage "
-            "may be incomplete");
-
     // Send the AFL runtime (if not shared object):
     std::string path(option_path);
     path += "/afl-rt";
@@ -635,9 +671,6 @@ extern intptr_t e9_plugin_match_v1(const Context *cxt)
  */
 extern void e9_plugin_patch_v1(const Context *cxt, Phase phase)
 {
-    if (option_instrument == OPTION_NEVER)
-        return;
-
     switch (phase)
     {
         case PHASE_CODE:

e9afl.cpp

@@ -244,11 +244,14 @@ int main(int argc, char **argv)
     // Setup environment:
     std::string path;
     getExePath(path);
-    setenv("E9AFL_COUNTER", getCounter(option_counter), true);
-    setenv("E9AFL_OBLOCK",  getValue(option_Oblock), true);
-    setenv("E9AFL_OSELECT", getValue(option_Oselect), true);
-    setenv("E9AFL_DEBUG",   (option_debug? "always": "default"), true);
-    setenv("E9AFL_PATH",    path.c_str(), true);
+    std::string plugin;
+    plugin += '\"';
+    plugin += path;
+    plugin += "/e9AFLPlugin.so\"";
+    std::string plugin_opt;
+    plugin_opt += "--plugin=";
+    plugin_opt += plugin;
+    plugin_opt += ':';
 
     // Construct command:
     std::string command;
@@ -265,14 +268,40 @@ int main(int argc, char **argv)
     command += output;
     command += "\" ";
 
-    command += "-M 'plugin(\"";
-    command += path;
-    command += "/e9AFLPlugin.so\").match()' ";
+    command += "-M 'plugin(";
+    command += plugin;
+    command += ").match()' ";
 
-    command += "-P 'plugin(\"";
-    command += path;
-    command += "/e9AFLPlugin.so\").patch()' ";
+    command += "-P 'plugin(";
+    command += plugin;
+    command += ").patch()' ";
+
+    command += plugin_opt;
+    command += "--counter=";
+    command += getCounter(option_counter);
+    command += ' ';
+
+    command += plugin_opt;
+    command += "-Oblock=";
+    command += getValue(option_Oblock);
+    command += ' ';
+
+    command += plugin_opt;
+    command += "-Oselect=";
+    command += getValue(option_Oselect);
+    command += ' ';
 
+    if (option_debug)
+    {
+        command += plugin_opt;
+        command += "--debug ";
+    }
+
+    command += plugin_opt; 
+    command += "--path='";
+    command += path;
+    command += "' ";
+ 
     for (int i = optind+1; i < argc; i++)
     {
         command += '\'';