1 files changed, 24 insertions, 1 deletions
diff --git a/lib/fuzzer/scripts/collect_data_flow.py b/lib/fuzzer/scripts/collect_data_flow.py
index d13f6dcc4..c3faf71c0 100755
--- a/lib/fuzzer/scripts/collect_data_flow.py
+++ b/lib/fuzzer/scripts/collect_data_flow.py
@@ -11,9 +11,15 @@
 # the complete trace for all input bytes (running it on all bytes at once
 # may fail if DFSan runs out of labels).
 # Usage:
-#   collect_data_flow.py BINARY INPUT [RESULT]
+#
+#   # Collect dataflow for one input, store it in OUTPUT (default is stdout)
+#   collect_data_flow.py BINARY INPUT [OUTPUT]
+#
+#   # Collect dataflow for all inputs in CORPUS_DIR, store them in OUTPUT_DIR
+#   collect_data_flow.py BINARY CORPUS_DIR OUTPUT_DIR
 #===------------------------------------------------------------------------===#
 import atexit
+import hashlib
 import sys
 import os
 import subprocess
@@ -26,9 +32,26 @@ def cleanup(d):
   print "removing: ", d
   shutil.rmtree(d)
 
+def collect_dataflow_for_corpus(self, exe, corpus_dir, output_dir):
+  print "Collecting dataflow for corpus:", corpus_dir, \
+    "output_dir:", output_dir
+  assert not os.path.exists(output_dir)
+  os.mkdir(output_dir)
+  for root, dirs, files in os.walk(corpus_dir):
+    for f in files:
+      path = os.path.join(root, f)
+      sha1 = hashlib.sha1(open(path).read()).hexdigest()
+      output = os.path.join(output_dir, sha1)
+      subprocess.call([self, exe, path, output])
+  functions_txt = open(os.path.join(output_dir, "functions.txt"), "w")
+  subprocess.call([exe], stdout=functions_txt)
+
+
 def main(argv):
   exe = argv[1]
   inp = argv[2]
+  if os.path.isdir(inp):
+    return collect_dataflow_for_corpus(argv[0], exe, inp, argv[3])
   size = os.path.getsize(inp)
   q = [[0, size]]
   tmpdir = tempfile.mkdtemp(prefix="libfuzzer-tmp-")