[libFuzzer] add a flag -malloc_limit_mb

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@319590 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 11 insertions, 2 deletions

diff --git a/lib/fuzzer/FuzzerDriver.cpp b/lib/fuzzer/FuzzerDriver.cpp
index e43f581fc..ccb9b0ca7 100644
--- a/lib/fuzzer/FuzzerDriver.cpp
+++ b/lib/fuzzer/FuzzerDriver.cpp
@@ -581,6 +581,9 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
   Options.PurgeAllocatorIntervalSec = Flags.purge_allocator_interval;
   Options.TraceMalloc = Flags.trace_malloc;
   Options.RssLimitMb = Flags.rss_limit_mb;
+  Options.MallocLimitMb = Flags.malloc_limit_mb;
+  if (!Options.MallocLimitMb)
+    Options.MallocLimitMb = Options.RssLimitMb;
   if (Flags.runs >= 0)
     Options.MaxNumberOfRuns = Flags.runs;
   if (!Inputs->empty() && !Flags.minimize_crash_internal_step)
diff --git a/lib/fuzzer/FuzzerFlags.def b/lib/fuzzer/FuzzerFlags.def
index 7326c3dfc..a32102a7d 100644
--- a/lib/fuzzer/FuzzerFlags.def
+++ b/lib/fuzzer/FuzzerFlags.def
@@ -130,6 +130,9 @@ FUZZER_FLAG_INT(trace_malloc, 0, "If >= 1 will print all mallocs/frees. "
     "If >= 2 will also print stack traces.")
 FUZZER_FLAG_INT(rss_limit_mb, 2048, "If non-zero, the fuzzer will exit upon"
     "reaching this limit of RSS memory usage.")
+FUZZER_FLAG_INT(malloc_limit_mb, 0, "If non-zero, the fuzzer will exit "
+    "if the target tries to allocate this number of Mb with one malloc call. "
+    "If zero (default) same limit as rss_limit_mb is applied.")
 FUZZER_FLAG_STRING(exit_on_src_pos, "Exit if a newly found PC originates"
     " from the given source location. Example: -exit_on_src_pos=foo.cc:123. "
     "Used primarily for testing libFuzzer itself.")
diff --git a/lib/fuzzer/FuzzerLoop.cpp b/lib/fuzzer/FuzzerLoop.cpp
index 3cfcfad93..f0de940e0 100644
--- a/lib/fuzzer/FuzzerLoop.cpp
+++ b/lib/fuzzer/FuzzerLoop.cpp
@@ -124,7 +124,7 @@ void FreeHook(const volatile void *ptr) {
 
 // Crash on a single malloc that exceeds the rss limit.
 void Fuzzer::HandleMalloc(size_t Size) {
-  if (!Options.RssLimitMb || (Size >> 20) < (size_t)Options.RssLimitMb)
+  if (!Options.MallocLimitMb || (Size >> 20) < (size_t)Options.MallocLimitMb)
     return;
   Printf("==%d== ERROR: libFuzzer: out-of-memory (malloc(%zd))\n", GetPid(),
          Size);
diff --git a/lib/fuzzer/FuzzerOptions.h b/lib/fuzzer/FuzzerOptions.h
index e0baf69a7..31e1a4de7 100644
--- a/lib/fuzzer/FuzzerOptions.h
+++ b/lib/fuzzer/FuzzerOptions.h
@@ -24,6 +24,7 @@ struct FuzzingOptions {
   int ErrorExitCode = 77;
   int MaxTotalTimeSec = 0;
   int RssLimitMb = 0;
+  int MallocLimitMb = 0;
   bool DoCrossOver = true;
   int MutateDepth = 5;
   bool ReduceDepth = false;
diff --git a/test/fuzzer/fuzzer-oom.test b/test/fuzzer/fuzzer-oom.test
index 9ef7c485d..308c4c5cd 100644
--- a/test/fuzzer/fuzzer-oom.test
+++ b/test/fuzzer/fuzzer-oom.test
@@ -8,7 +8,9 @@ CHECK: ERROR: libFuzzer: out-of-memory (used: {{.*}}; limit: 300Mb)
 CHECK: Test unit written to ./oom-
 SUMMARY: libFuzzer: out-of-memory
 
-RUN: not %t-OutOfMemorySingleLargeMallocTest -rss_limit_mb=300 2>&1 | FileCheck %s --check-prefix=SINGLE_LARGE_MALLOC
+RUN: not %t-OutOfMemorySingleLargeMallocTest -rss_limit_mb=300    2>&1 | FileCheck %s --check-prefix=SINGLE_LARGE_MALLOC
+RUN: not %t-OutOfMemorySingleLargeMallocTest -malloc_limit_mb=300 2>&1 | FileCheck %s --check-prefix=SINGLE_LARGE_MALLOC
+RUN: not %t-OutOfMemorySingleLargeMallocTest -rss_limit_mb=1000 -malloc_limit_mb=300 2>&1 | FileCheck %s --check-prefix=SINGLE_LARGE_MALLOC
 
 We used to check for "out-of-memory (malloc(53{{.*}}))", but that would fail
 sometimes, so now we accept any OOM message.