3 files changed, 50 insertions, 18 deletions
diff --git a/compiler-rt/test/fuzzer/AFLDriverTest.cpp b/compiler-rt/test/fuzzer/AFLDriverTest.cpp
index d2937d004d1..84b5f9f6ba2 100644
--- a/compiler-rt/test/fuzzer/AFLDriverTest.cpp
+++ b/compiler-rt/test/fuzzer/AFLDriverTest.cpp
@@ -2,28 +2,33 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-// Contains dummy functions used to avoid dependency on AFL.
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 
+// Dummy functions used to avoid dependency on AFL.
 extern "C" void __afl_manual_init() {}
 
 extern "C" int __afl_persistent_loop(unsigned int N) {
   static int Count = N;
-  fprintf(stderr, "__afl_persistent_loop calle, Count = %d\n", Count);
-  if (Count--) return 1;
-  return 0;
+  fprintf(stderr, "__afl_persistent_loop called, Count = %d\n", Count);
+  return Count--;
 }
 
 // This declaration exists to prevent the Darwin linker
 // from complaining about this being a missing weak symbol.
 extern "C" int LLVMFuzzerInitialize(int *argc, char ***argv) {
-  fprintf(stderr, "LLVMFuzzerInitialize called\n");
   return 0;
 }
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
-  fprintf(stderr, "LLVMFuzzerTestOneInput called; Size = %zd\n", Size);
-  return 0;
+  puts("STDOUT MESSAGE");
+  fflush(stdout);
+  fprintf(stderr, "STDERR MESSAGE\n"
+                  "LLVMFuzzerTestOneInput called; Size = %zd\n",
+          Size);
+  if (Size < 4)
+    return 0;
+
+  return Data[Size];
 }
diff --git a/compiler-rt/test/fuzzer/afl-driver-close-fd-mask.test b/compiler-rt/test/fuzzer/afl-driver-close-fd-mask.test
new file mode 100644
index 00000000000..71f74e27ec4
--- /dev/null
+++ b/compiler-rt/test/fuzzer/afl-driver-close-fd-mask.test
@@ -0,0 +1,31 @@
+REQUIRES: linux
+RUN: %no_fuzzer_cpp_compiler %S/AFLDriverTest.cpp %libfuzzer_src/afl/afl_driver.cpp -o %t-AFLDriverTest
+
+; Test that not specifying AFL_DRIVER_CLOSE_FD_MASK works as intended.
+RUN: echo -n "abc" > %t.file3
+RUN: unset AFL_DRIVER_CLOSE_FD_MASK
+RUN: %run %t-AFLDriverTest < %t.file3 2>&1 | FileCheck %s --check-prefixes=STDERR,STDOUT
+STDOUT: STDOUT MESSAGE
+STDERR: STDERR MESSAGE
+
+; Test that stdout is closed properly.
+RUN: AFL_DRIVER_CLOSE_FD_MASK=1 %run %t-AFLDriverTest < %t.file3 2>&1 | FileCheck %s --check-prefixes=NOT_STDOUT,STDERR
+NOT_STDOUT-NOT: STDOUT MESSAGE
+
+; Test that stderr is closed properly.
+RUN: AFL_DRIVER_CLOSE_FD_MASK=2 %run %t-AFLDriverTest < %t.file3 2>&1 | FileCheck %s --check-prefixes=NOT_STDERR,STDOUT
+NOT_STDERR-NOT: STDERR MESSAGE
+
+; Test that both are closed properly.
+RUN: AFL_DRIVER_CLOSE_FD_MASK=3 %run %t-AFLDriverTest < %t.file3 2>&1 | FileCheck %s --check-prefixes=NOT_STDERR,NOT_STDOUT
+
+; Test that a stack is printed when we close stderr
+RUN: echo -n "abcd" > %t.file4
+RUN: AFL_DRIVER_CLOSE_FD_MASK=2 not %run %t-AFLDriverTest < %t.file4 2>&1  | FileCheck %s --check-prefixes=ASAN_CRASH,STDOUT,NOT_STDERR
+ASAN_CRASH: ERROR: AddressSanitizer
+
+; Test that a stack is written to the stderr duplicate file when we close stderr
+; and specify a duplicate.
+RUN: rm -f %t.stderr
+RUN: AFL_DRIVER_STDERR_DUPLICATE_FILENAME=%t.stderr AFL_DRIVER_CLOSE_FD_MASK=2 not %run %t-AFLDriverTest < %t.file4
+RUN: cat %t.stderr | FileCheck %s --check-prefixes=ASAN_CRASH,NOT_STDERR
diff --git a/compiler-rt/test/fuzzer/afl-driver.test b/compiler-rt/test/fuzzer/afl-driver.test
index 552bafb0bf3..58f422f1ed8 100644
--- a/compiler-rt/test/fuzzer/afl-driver.test
+++ b/compiler-rt/test/fuzzer/afl-driver.test
@@ -3,27 +3,23 @@ REQUIRES: linux
 RUN: %no_fuzzer_cpp_compiler %S/AFLDriverTest.cpp %libfuzzer_src/afl/afl_driver.cpp -o %t-AFLDriverTest
 
 RUN: echo -n "abc" > %t.file3
-RUN: echo -n "abcd" > %t.file4
-
 RUN: %run %t-AFLDriverTest < %t.file3 2>&1 | FileCheck %s --check-prefix=CHECK1
-CHECK1: __afl_persistent_loop calle, Count = 1000
+CHECK1: __afl_persistent_loop called, Count = 1000
 CHECK1: LLVMFuzzerTestOneInput called; Size = 3
 
-
 RUN: %run %t-AFLDriverTest < %t.file3 -42 2>&1 | FileCheck %s --check-prefix=CHECK2
-CHECK2: __afl_persistent_loop calle, Count = 42
+CHECK2: __afl_persistent_loop called, Count = 42
 CHECK2: LLVMFuzzerTestOneInput called; Size = 3
 
-
 RUN: %run %t-AFLDriverTest < %t.file3 666 2>&1 | FileCheck %s --check-prefix=CHECK3
 CHECK3: WARNING: using the deprecated call style
-CHECK3: __afl_persistent_loop calle, Count = 666
+CHECK3: __afl_persistent_loop called, Count = 666
 CHECK3: LLVMFuzzerTestOneInput called; Size = 3
 
-
 RUN: %run %t-AFLDriverTest %t.file3 2>&1 | FileCheck %s --check-prefix=CHECK4
 CHECK4: LLVMFuzzerTestOneInput called; Size = 3
 
-RUN: %run %t-AFLDriverTest %t.file3 %t.file4  2>&1 | FileCheck %s --check-prefix=CHECK5
-CHECK5: LLVMFuzzerTestOneInput called; Size = 3
-CHECK5: LLVMFuzzerTestOneInput called; Size = 4
+RUN: echo -n "ab" > %t.file2
+RUN: %run %t-AFLDriverTest %t.file2 %t.file3 2>&1 | FileCheck %s --check-prefix=CHECK5
+CHECK5: LLVMFuzzerTestOneInput called; Size = 2
+CHECK5: LLVMFuzzerTestOneInput called; Size = 3
+\ No newline at end of file