create processes in a job object to enforce memory limitation

gnbon · Aug 22, 2017 · c95ea5c · c95ea5c
1 parent 47dc7e9
commit c95ea5c
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 23 deletions.
diff --git a/afl-fuzz.c b/afl-fuzz.c
@@ -2039,6 +2039,8 @@ static void create_target_process(char** argv) {
   size_t pidsize;
   BOOL inherit_handles = TRUE;
 
+  HANDLE hJob = NULL;
+  JOBOBJECT_EXTENDED_LIMIT_INFORMATION job_limit;
   STARTUPINFO si;
   PROCESS_INFORMATION pi;
 
@@ -2054,16 +2056,15 @@ static void create_target_process(char** argv) {
     20000,                    // client time-out
     NULL);                    // default security attribute
 
-  if (pipe_handle == INVALID_HANDLE_VALUE)
-  {
-      FATAL("CreateNamedPipe failed, GLE=%d.\n", GetLastError());
+  if (pipe_handle == INVALID_HANDLE_VALUE) {
+    FATAL("CreateNamedPipe failed, GLE=%d.\n", GetLastError());
   }
 
   target_cmd = argv_to_cmd(argv);
 
-  ZeroMemory( &si, sizeof(si) );
+  ZeroMemory(&si, sizeof(si));
   si.cb = sizeof(si);
-  ZeroMemory( &pi, sizeof(pi) );
+  ZeroMemory(&pi, sizeof(pi));
 
   if(sinkhole_stds) {
     si.hStdOutput = si.hStdError = devnul_handle;
@@ -2090,26 +2091,54 @@ static void create_target_process(char** argv) {
     );
   }
 
-  if(!CreateProcess(NULL, cmd, NULL, NULL, inherit_handles, /*CREATE_NO_WINDOW*/0, NULL, NULL, &si, &pi)) {
+  if(mem_limit != 0) {
+    hJob = CreateJobObject(NULL, NULL);
+    if(hJob == NULL) {
+      FATAL("CreateJobObject failed, GLE=%d.\n", GetLastError());
+    }
+
+    ZeroMemory(&job_limit, sizeof(job_limit));
+    job_limit.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
+    job_limit.ProcessMemoryLimit = mem_limit * 1024 * 1024;
+
+    if(!SetInformationJobObject(
+      hJob,
+      JobObjectExtendedLimitInformation,
+      &job_limit,
+      sizeof(job_limit)
+    )) {
+      FATAL("SetInformationJobObject failed, GLE=%d.\n", GetLastError());
+    }
+  }
+
+  if(!CreateProcess(NULL, cmd, NULL, NULL, inherit_handles, CREATE_SUSPENDED, NULL, NULL, &si, &pi)) {
     FATAL("CreateProcess failed, GLE=%d.\n", GetLastError());
   }
 
   child_handle = pi.hProcess;
   child_thread_handle = pi.hThread;
 
+  if(mem_limit != 0) {
+    if(!AssignProcessToJobObject(hJob, child_handle)) {
+      FATAL("AssignProcessToJobObject failed, GLE=%d.\n", GetLastError());
+    }
+  }
+
+  ResumeThread(child_thread_handle);
+
   watchdog_timeout_time = get_cur_time() + exec_tmout;
   watchdog_enabled = 1;
 
   if(!ConnectNamedPipe(pipe_handle, NULL)) {
     if(GetLastError() != ERROR_PIPE_CONNECTED) {
-        FATAL("ConnectNamedPipe failed, GLE=%d.\n", GetLastError());
+      FATAL("ConnectNamedPipe failed, GLE=%d.\n", GetLastError());
     }
   }
 
   watchdog_enabled = 0;
 
-  //by the time pipe has connected the pidfile must have been created
   if(drioless == 0) {
+    //by the time pipe has connected the pidfile must have been created
     fp = fopen(pidfile, "rb");
     if(!fp) {
       FATAL("Error opening pidfile.txt");
@@ -2121,6 +2150,7 @@ static void create_target_process(char** argv) {
     fread(buf, pidsize, 1, fp);
     buf[pidsize] = 0;
     fclose(fp);
+    remove(pidfile);
     child_pid = atoi(buf);
     free(buf);
     ck_free(pidfile);

diff --git a/afl-showmap.c b/afl-showmap.c
@@ -438,6 +438,8 @@ static void create_target_process(char** argv) {
   size_t pidsize;
   BOOL inherit_handles = TRUE;
 
+  HANDLE hJob = NULL;
+  JOBOBJECT_EXTENDED_LIMIT_INFORMATION job_limit;
   STARTUPINFO si;
   PROCESS_INFORMATION pi;
 
@@ -459,9 +461,9 @@ static void create_target_process(char** argv) {
 
   target_cmd = argv_to_cmd(argv);
 
-  ZeroMemory( &si, sizeof(si) );
+  ZeroMemory(&si, sizeof(si));
   si.cb = sizeof(si);
-  ZeroMemory( &pi, sizeof(pi) );
+  ZeroMemory(&pi, sizeof(pi));
 
   if(quiet_mode) {
     si.hStdOutput = si.hStdError = devnul_handle;
@@ -470,7 +472,6 @@ static void create_target_process(char** argv) {
     inherit_handles = FALSE;
   }
 
-
   if(drioless) {
     char *static_config = alloc_printf("%s:1", fuzzer_id);
 
@@ -489,26 +490,54 @@ static void create_target_process(char** argv) {
     );
   }
 
-  if(!CreateProcess(NULL, cmd, NULL, NULL, inherit_handles, /*CREATE_NO_WINDOW*/0, NULL, NULL, &si, &pi)) {
+  if(mem_limit != 0) {
+    hJob = CreateJobObject(NULL, NULL);
+    if(hJob == NULL) {
+      FATAL("CreateJobObject failed, GLE=%d.\n", GetLastError());
+    }
+
+    ZeroMemory(&job_limit, sizeof(job_limit));
+    job_limit.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
+    job_limit.ProcessMemoryLimit = mem_limit * 1024 * 1024;
+
+    if(!SetInformationJobObject(
+      hJob,
+      JobObjectExtendedLimitInformation,
+      &job_limit,
+      sizeof(job_limit)
+    )) {
+      FATAL("SetInformationJobObject failed, GLE=%d.\n", GetLastError());
+    }
+  }
+
+  if(!CreateProcess(NULL, cmd, NULL, NULL, inherit_handles, CREATE_SUSPENDED, NULL, NULL, &si, &pi)) {
     FATAL("CreateProcess failed, GLE=%d.\n", GetLastError());
   }
 
   child_handle = pi.hProcess;
   child_thread_handle = pi.hThread;
 
+  if(mem_limit != 0) {
+    if(!AssignProcessToJobObject(hJob, child_handle)) {
+      FATAL("AssignProcessToJobObject failed, GLE=%d.\n", GetLastError());
+    }
+  }
+
+  ResumeThread(child_thread_handle);
+
   watchdog_timeout_time = get_cur_time() + exec_tmout;
   watchdog_enabled = 1;
 
   if(!ConnectNamedPipe(pipe_handle, NULL)) {
     if(GetLastError() != ERROR_PIPE_CONNECTED) {
-        FATAL("ConnectNamedPipe failed, GLE=%d.\n", GetLastError());
+      FATAL("ConnectNamedPipe failed, GLE=%d.\n", GetLastError());
     }
   }
 
   watchdog_enabled = 0;
 
-  //by the time pipe has connected the pidfile must have been created
   if(drioless == 0) {
+    //by the time pipe has connected the pidfile must have been created
     fp = fopen(pidfile, "rb");
     if(!fp) {
       FATAL("Error opening pidfile.txt");
@@ -521,7 +550,6 @@ static void create_target_process(char** argv) {
     buf[pidsize] = 0;
     fclose(fp);
     remove(pidfile);
-
     child_pid = atoi(buf);
     free(buf);
     ck_free(pidfile);

diff --git a/config.h b/config.h
@@ -55,9 +55,9 @@
 /* Default memory limit for child process (MB): */
 
 #ifndef __x86_64__ 
-#  define MEM_LIMIT         25
+#  define MEM_LIMIT         100
 #else
-#  define MEM_LIMIT         50
+#  define MEM_LIMIT         100
 #endif /* ^!__x86_64__ */
 
 /* Default memory limit when running in QEMU mode (MB): */

diff --git a/test_static.cpp b/test_static.cpp
@@ -1,8 +1,13 @@
 /*
-   WinAFL - A simple test binary that crashes on certain inputs:
-     - 'test1' with a normal write access violation at NULL
-     - 'test2' with a /GS stack cookie violation
-     - 'test3' with a hang
+   WinAFL - A simple test binary that exercises various behaviors
+   depending on inputs:
+     - 'test1' crashes with a normal write access violation at NULL
+     - 'test2' crashes with a /GS stack cookie violation
+     - 'test3' triggers a hang
+     - 'test4' triggers an exception that is caught and handled
+     - 'test5' triggers an OutputDebugString
+     - 'test6' triggers an allocation of 120MB (and a crash if the
+               allocation fails)
    -------------------------------------------------------------
 
    Written by Axel "0vercl0k" Souchet <[email protected]>
@@ -96,6 +101,12 @@ int test(int argc, char **argv) {
     else if (c == '5') {
         OutputDebugString(TEXT("hello!"));
     }
+    else if (c == '6') {
+        printf("allocating 120MB\n");
+        char *buffer = (char*)malloc((1024 * 1024) * 120);
+        *buffer = 0;
+        free(buffer);
+    }
     else {
         printf("Error 5\n");
     }

diff --git a/testcases/tests/bigalloc.txt b/testcases/tests/bigalloc.txt
@@ -0,0 +1 @@
+test6
diff --git a/winafl-cmin.py b/winafl-cmin.py
@@ -24,6 +24,7 @@
 import subprocess
 import sys
 import time
+import re
 from textwrap import dedent, wrap
 
 nul = open(os.devnull, 'wb')
@@ -46,7 +47,7 @@ def _to_showmap_options(args, trace_name = '-'):
         '''Takes the argparse namespace, and convert it to the list of options used
         to invoke afl-showmap.exe'''
         r = [
-            'afl-showmap.exe', '-o', trace_name
+            'afl-showmap.exe', '-o', trace_name, '-m', args.memory_limit
         ]
 
         if os.getenv('AFL_NO_SINKHOLE') is None:
@@ -145,6 +146,13 @@ def target_offset(opt):
     except:
         raise argparse.ArgumentTypeError('must be an integer')
 
+def memory_limit(opt):
+    '''Validates that the -m parameter is properly formated, else
+    raises an ArgumentTypeError exception back to the argparse parser.'''
+    if re.match('^\d+[TGkM]{0,1}$', opt) or opt == 'none':
+        return opt
+    raise argparse.ArgumentTypeError('must be an integer followed by either: '
+                                     'T, G, M, k or nothing; or none')
 
 def setup_argparse():
     '''Sets up the argparse configuration.'''
@@ -174,7 +182,7 @@ def setup_argparse():
     group = parser.add_argument_group('basic parameters')
     group.add_argument(
         '-i', '--input', action = 'append', required = True,
-        metavar = 'dir', help = 'input directory with the starting corpus.' \
+        metavar = 'dir', help = 'input directory with the starting corpus.'
         ' Multiple input directories are supported'
     )
     group.add_argument(
@@ -246,6 +254,10 @@ def setup_argparse():
         '-t', '--time-limit', type = int, default = 0,
         metavar = 'msec', help = 'timeout for each run (none)'
     )
+    group.add_argument(
+        '-m', '--memory-limit', default = 'none', type = memory_limit,
+        metavar = 'megs', help = 'memory limit for child process'
+    )
     # Note(0vercl0k): If you use -f, which means you want the input file at
     # a specific location (and a specific name), we have to force the pool
     # process to contain only a single worker as there is a unique location