Merge pull request #480 from 4B5F5F4B/master

enable TinyInst instrumentation mode for winafl-cmin.py

Author: ifratric

CMakeLists.txt

@@ -73,6 +73,10 @@ add_executable(afl-showmap
   afl-showmap.c
   )
 
+if (${TINYINST})
+  target_link_libraries(afl-showmap winafl_tinyinst)
+endif()
+
 project(afl-tmin)
 
 add_executable(afl-tmin

afl-showmap.c

@@ -50,6 +50,14 @@
 #include <sys/stat.h>
 #include <sys/types.h>
 
+#ifdef TINYINST
+int tinyinst_init(int argc, char** argv);
+void tinyinst_set_fuzzer_id(char* fuzzer_id);
+int tinyinst_run(char** argv, uint32_t timeout);
+void tinyinst_killtarget();
+#endif
+
+
 static s32 child_pid;                 /* PID of the tested program         */
 
 static HANDLE child_handle,
@@ -87,6 +95,7 @@ static u8  quiet_mode,                /* Hide non-essential messages?      */
            cmin_mode,                 /* Generate output in afl-cmin mode? */
            binary_mode,               /* Write output as a binary map      */
            drioless = 0;              /* Running without DRIO?             */
+          use_tinyinst = 0;          /* Using TinyInst instrumentation   */
 
 static volatile u8
            stop_soon,                 /* Ctrl-C pressed?                   */
@@ -572,6 +581,13 @@ static void destroy_target_process(int wait_exit) {
   STARTUPINFO si;
   PROCESS_INFORMATION pi;
 
+#ifdef TINYINST
+  if (use_tinyinst) {
+    tinyinst_killtarget();
+    return;
+  }
+#endif
+
   EnterCriticalSection(&critical_section);
 
   if(!child_handle) {
@@ -675,6 +691,13 @@ static int is_child_running() {
 
 static void run_target(char** argv) {
 
+
+#ifdef TINYINST
+  if (use_tinyinst) {
+    return tinyinst_run(argv, exec_tmout);
+  }
+#endif
+
   char command[] = "F";
   DWORD num_read;
   char result = 0;
@@ -963,15 +986,24 @@ int main(int argc, char** argv) {
   optind = 1;
   dynamorio_dir = NULL;
   client_params = NULL;
+  use_tinyinst = 0;
 
 #ifdef USE_COLOR
   enable_ansi_console();
 #endif
 
-  while ((opt = getopt(argc, argv, "+o:m:t:A:D:eqZQbY")) > 0)
-
+  while ((opt = getopt(argc, argv, "+o:m:t:A:D:eqyZQbY")) > 0)
     switch (opt) {
 
+    case 'y':
+        #ifdef TINYINST
+              use_tinyinst = 1;
+        #else
+              FATAL("afl-fuzz was not compiled with TinyInst support");
+        #endif
+        break;
+      
+      
       case 'D': /* dynamorio dir */
 
         if(dynamorio_dir) FATAL("Multiple -D options not supported");
@@ -1094,13 +1126,30 @@ int main(int argc, char** argv) {
 
   if(!out_file) usage(argv[0]);
   if(!drioless) {
-    if(optind == argc || !dynamorio_dir) usage(argv[0]);
+    if(optind == argc || (!dynamorio_dir && !use_tinyinst)) usage(argv[0]);
   }
 
-  extract_client_params(argc, argv);
+ if (use_tinyinst) {
+#ifdef TINYINST
+    int tinyinst_options = tinyinst_init(argc - optind, argv + optind);
+    if (!tinyinst_options) usage(argv[0]);
+    optind += tinyinst_options;
+#endif
+  } else {
+	  extract_client_params(argc, argv);
+  }
   optind++;
 
   setup_shm();
+
+  if (use_tinyinst) {
+#ifdef TINYINST
+    tinyinst_set_fuzzer_id(fuzzer_id);
+#endif
+  }
+
+
+
   setup_watchdog_timer();
   setup_signal_handlers();
 

winafl-cmin.py

@@ -66,23 +66,27 @@ def _to_showmap_options(args, trace_name = '-'):
         if args.static_instr:
             r.append('-Y')
         else:
-            r.extend(['-D', args.dynamorio_dir])
-            r.append('--')
-            r.extend(['-target_module', args.target_module])
+            if args.tinyinst_instr:
+                r.append('-y')
+                r.append('--')
+                for mod in args.instrument_modules:
+                    r.extend(['-instrument_module', mod])
+            else:
+                r.extend(['-D', args.dynamorio_dir])
+                r.append('--')
+                for mod in args.coverage_modules:
+                    r.extend(['-coverage_module', mod])
 
+            r.extend(['-target_module', args.target_module])
             if args.target_method is None:
                 r.extend(['-target_offset', '0x%x' % args.target_offset])
             else:
                 r.extend(['-target_method', args.target_method])
-
             r.extend(['-nargs', '%d' % args.nargs])
             r.extend(['-covtype', args.covtype])
             if args.call_convention is not None:
                 r.extend(['-call_convention', args.call_convention])
 
-            for mod in args.coverage_modules:
-                r.extend(['-coverage_module', mod])
-
         r.append('--')
         r.extend(args.target_cmdline)
         return r
@@ -179,6 +183,9 @@ def setup_argparse():
 
              * Typical use with static instrumentation
               winafl-cmin.py -Y -t 100000 -i in -o minset -- test.instr.exe @@
+            
+             * Typical use with TinyInst mode instrumentation
+             winafl-cmin.py -y -t 100000 -i in -o minset -instrument_module m.dll -target_module test.exe -target_method fuzz -nargs 2 -- test.exe @@
             '''
         ), 100, replace_whitespace = False))
     )
@@ -227,6 +234,26 @@ def setup_argparse():
         metavar = 'dir', help = 'directory containing DynamoRIO binaries (drrun, drconfig)'
     )
 
+    instr_type.add_argument(
+        '-y', '--tinyinst_instr', action = 'store_true',
+        help = 'use the TinyInst instrumentation mode'
+    )
+
+
+    instr_module = group.add_mutually_exclusive_group(required = True)
+    instr_module.add_argument(
+        '-coverage_module', dest = 'coverage_modules', default = None,
+        action = 'append', metavar = 'module', help = 'module for which to record coverage.'
+        ' Multiple module flags are supported'
+    )
+
+    instr_module.add_argument(
+        '-instrument_module', dest = 'instrument_modules', default = None,
+        action = 'append', metavar = 'module', help = 'module for which to record coverage.'
+        ' Multiple module flags are supported'
+    )
+
+
     group.add_argument(
         '-covtype', choices = ('edge', 'bb'), default = 'bb',
         help = 'the type of coverage being recorded (defaults to bb)'
@@ -235,15 +262,12 @@ def setup_argparse():
         '-call_convention', choices = ('stdcall', 'fastcall', 'thiscall', 'ms64'),
         default = 'stdcall', help = 'the calling convention of the target_method'
     )
-    group.add_argument(
-        '-coverage_module', dest = 'coverage_modules', default = None,
-        action = 'append', metavar = 'module', help = 'module for which to record coverage.'
-        ' Multiple module flags are supported'
-    )
+    
     group.add_argument(
         '-target_module', default = None, metavar = 'module',
         help = 'module which contains the target function to be fuzzed'
     )
+
     group.add_argument(
         '-nargs', type = int, default = None, metavar = 'nargs',
         help = 'number of arguments the fuzzed method takes. This is used to save/restore'
@@ -261,6 +285,7 @@ def setup_argparse():
         help = 'offset of the method to fuzz from the start of the module'
     )
 
+
     group = parser.add_argument_group('execution control settings')
     group.add_argument(
         '-t', '--time-limit', type = int, default = 0,
@@ -370,35 +395,44 @@ def validate_args(args):
             return False
 
     if not args.static_instr:
-        # Make sure we have all the arguments we need
-        if len(args.coverage_modules) == 0:
-            logging.error(
-                '[!] -coverage_module is a required option to use'
-                'the dynamic instrumentation'
-            )
-            return False
-
         if None in [args.target_module, args.nargs]:
-            logging.error(
-                '[!] , -target_module and -nargs are required'
-                ' options to use the dynamic instrumentation mode.'
-            )
-            return False
-
+                logging.error(
+                    '[!] , -target_module and -nargs are required'
+                    ' options to use the dynamic instrumentation mode.'
+                )
+                return False
+            
         if args.target_method is None and args.target_offset is None:
             logging.error(
                 '[!] -target_method or -target_offset is required to use the'
                 ' dynamic instrumentation mode'
             )
             return False
+         
+        if not args.tinyinst_instr:
+            # Make sure we have all the arguments we need
+            if len(args.coverage_modules) == 0:
+                logging.error(
+                    '[!] -coverage_module is a required option to use'
+                    'the dynamic instrumentation'
+                )
+                return False
 
-        # If we are using DRIO, one of the thing we need is the DRIO client
-        winafl_path = os.path.join(args.working_dir, 'winafl.dll')
-        if not os.path.isfile(winafl_path):
-            logging.error(
-                '[!] winafl.dll needs to be in %s.', args.working_dir
-            )
-            return False
+            # If we are using DRIO, one of the thing we need is the DRIO client
+            winafl_path = os.path.join(args.working_dir, 'winafl.dll')
+            if not os.path.isfile(winafl_path):
+                logging.error(
+                    '[!] winafl.dll needs to be in %s.', args.working_dir
+                )
+                return False
+        else:
+             # Make sure we have all the arguments we need
+            if len(args.instrument_modules) == 0:
+                logging.error(
+                    '[!] -instrument_module is a required option to use'
+                    'the dynamic instrumentation'
+                )
+                return False
 
     if args.file_read is not None and '@@' not in args.file_read:
         # When a particular input file is specified, first