added -I option

3 files changed, 33 insertions, 13 deletions

diff --git a/CHANGES b/CHANGES
index a5f8e47..149b3d7 100644
--- a/CHANGES
+++ b/CHANGES
@@ -2,6 +2,7 @@ Changelog
 =========
 
 https://github.com/vanhauser-thc/afl-dyninst
+ - added -I option (only instrument specific functions)
  - updated the README for guidance to build against dyninst version 10
  - added support for dyninst version 10
  - added -x performance optimization options, before this afl-dyninst was meh,
diff --git a/README.md b/README.md
index c80282e..0c36f77 100644
--- a/README.md
+++ b/README.md
@@ -63,21 +63,22 @@ Depending on the age of your Linux OS you can try to use packages from your dist
 
 ## Commandline options
 ```
-Usage: ./afl-dyninst -dfvD -i INPUT_BINARY -o OUTPUT_BINARY -l INPUT_LIBRARY -e ADDRESS -E ADDRESS -s NUMBER -S FUNCNAME -m SIZE
-   -i: input binary 
-   -o: output binary
+Usage: afl-dyninst -dfvxD -i binary -o  binary -l library -e address -E address -s number -S funcname -I funcname -m size
+   -i: input binary program
+   -o: output binary program
    -d: do not instrument the binary, only supplied libraries
    -l: linked library to instrument (repeat for more than one)
    -r: runtime library to instrument (path to, repeat for more than one)
    -e: entry point address to patch (required for stripped binaries)
    -E: exit point - force exit(0) at this address (repeat for more than one)
    -s: number of initial basic blocks to skip in binary
-   -m: minimum size of a basic bock to instrument (default: 1)
-   -f: try to fix a dyninst bug that leads to crashes
+   -m: minimum size of a basic bock to instrument (default: 10)
+   -f: try to fix a dyninst bug that leads to crashes (loss of 20%% performance)
+   -I: only instrument this function and nothing else (repeat for more than one)
    -S: do not instrument this function (repeat for more than one)
-   -D: instrument fork server and forced exit functions but no basic blocks
-   -x: experimental performance modes (can be set up to three times)
-         -x (level 1) : ~40-50%% improvement
+   -D: instrument only a simple fork server and also forced exit functions
+   -x: experimental performance modes (can be set up to two times)
+         -x (level 1):  ~40-50%% improvement
          -xx (level 2): ~100%% vs normal, ~40%% vs level 1
    -v: verbose output
 ```
@@ -104,7 +105,7 @@ optimization option, as skipping the basic blocks of the initialization
 routines makes things run faster.  If the instrumented binary is crashing by
 itself, try skiping a number of blocks.
 
-Switch -r allows you to specify a path to the library that is loaded
+Switch -r allows you to specify a path to a library that is loaded
 via dlopen() at runtime. Instrumented runtime libraries will be 
 written to the same location with a ".ins" suffix as not to overwrite
 the original ones. Make sure to backup the originals and then rename the
diff --git a/afl-dyninst.cpp b/afl-dyninst.cpp
index 35e55f3..157de21 100644
--- a/afl-dyninst.cpp
+++ b/afl-dyninst.cpp
@@ -39,6 +39,7 @@ set < string > todo;
 set < string > instrumentLibraries;
 set < string > runtimeLibraries;
 set < string > skipAddresses;
+set < string > onlyAddresses;
 set < unsigned long > exitAddresses;
 unsigned int bbMinSize = 10;
 int bbSkip = 0, performance = 0;
@@ -53,8 +54,8 @@ const char *functions[] = { "main", "_main", "_initproc", "_init", "start", "_st
 
 const char *instLibrary = "libAflDyninst.so";
 
-static const char *OPT_STR = "fi:o:l:e:E:vs:dr:m:S:Dx";
-static const char *USAGE = " -dfvxD -i <binary> -o <binary> -l <library> -e <address> -E <address> -s <number> -S <funcname> -m <size>\n \
+static const char *OPT_STR = "fi:o:l:e:E:vs:dr:m:S:I:Dx";
+static const char *USAGE = " -dfvxD -i <binary> -o <binary> -l <library> -e <address> -E <address> -s <number> -S <funcname> -I <funcname> -m <size>\n \
   -i: input binary \n \
   -o: output binary\n \
   -d: do not instrument the binary, only supplied libraries\n \
@@ -65,6 +66,7 @@ static const char *USAGE = " -dfvxD -i <binary> -o <binary> -l <library> -e <add
   -s: number of initial basic blocks to skip in binary\n \
   -m: minimum size of a basic bock to instrument (default: 10)\n \
   -f: try to fix a dyninst bug that leads to crashes (loss of 20%% performance)\n \
+  -I: only instrument this function and nothing else (repeat for more than one)\n \
   -S: do not instrument this function (repeat for more than one)\n \
   -D: instrument only a simple fork server and also forced exit functions\n \
   -x: experimental performance modes (can be set up to two times)\n \
@@ -92,6 +94,9 @@ bool parseOptions(int argc, char **argv) {
         performance = 2;
       }
       break;
+    case 'I':
+      onlyAddresses.insert(optarg);
+      break;
     case 'S':
       skipAddresses.insert(optarg);
       break;
@@ -562,9 +567,11 @@ int main(int argc, char **argv) {
         int do_patch = 1;
 
         curFunc->getName(funcName, 1024);
-        if (string(funcName) == string("_init") || string(funcName) == string("__libc_csu_init") || string(funcName) == string("_start")
-          )
+        if (string(funcName) == string("_init") || string(funcName) == string("__libc_csu_init") || string(funcName) == string("_start")) {
+          if (verbose)
+            cout << "Skipping instrumenting function " << funcName << endl;
           continue;             // here's a bug on hlt // XXX: check what happens if removed
+        }
         if (!skipAddresses.empty()) {
           set < string >::iterator saiter;
           for (saiter = skipAddresses.begin(); saiter != skipAddresses.end() && do_patch == 1; saiter++)
@@ -575,6 +582,17 @@ int main(int argc, char **argv) {
             continue;
           }
         }
+        if (!onlyAddresses.empty()) {
+          do_patch = 0;
+          set < string >::iterator saiter;
+          for (saiter = onlyAddresses.begin(); saiter != onlyAddresses.end() && do_patch == 1; saiter++)
+            if (*saiter == string(funcName))
+              do_patch = 1;
+          if (do_patch == 0) {
+            cout << "Skipping instrumenting function " << funcName << endl;
+            continue;
+          }
+        }
         insertBBCallback(appBin, curFunc, funcName, bbCallback, &bbIndex);
       }
     }