Adding PFFFT seed corpus generation. (#1790)

3 files changed, 85 insertions, 1 deletions

diff --git a/projects/pffft/Dockerfile b/projects/pffft/Dockerfile
index 4a99a9776..61e69f1ca 100644
--- a/projects/pffft/Dockerfile
+++ b/projects/pffft/Dockerfile
@@ -16,9 +16,10 @@
 
 FROM gcr.io/oss-fuzz-base/base-builder
 MAINTAINER alessiob@webrtc.org
-RUN apt-get update && apt-get install -y mercurial
+RUN apt-get update && apt-get install -y mercurial python-numpy
 RUN hg clone https://bitbucket.org/jpommier/pffft $SRC/pffft
 WORKDIR pffft
 COPY build.sh $SRC
 # TODO(alessiob): Move the fuzzing source code to pffft upstream.
+COPY generate_seed_corpus.py $SRC/pffft
 COPY pffft_fuzzer.cc $SRC/pffft
diff --git a/projects/pffft/build.sh b/projects/pffft/build.sh
index 620bc998b..190f7d75f 100644
--- a/projects/pffft/build.sh
+++ b/projects/pffft/build.sh
@@ -18,11 +18,35 @@
 SRC_DIR=$SRC/pffft
 cd $WORK
 
+# Deploy the seed corpus.
+SEED_CORPUS_ZIP_PATH=$OUT/pffft_fuzzers_seed_corpus.zip
+if [ -d seed_corpus_tmp ]; then 
+  rm -fr seed_corpus_tmp
+fi
+mkdir seed_corpus_tmp
+python $SRC_DIR/generate_seed_corpus.py seed_corpus_tmp
+cd seed_corpus_tmp
+zip -q $SEED_CORPUS_ZIP_PATH ./*
+cd ..
+rm -fr seed_corpus_tmp
+
 build_fuzzer() {
   # Aliases for the arguments.
   fft_transform=$1
 
+  # Fuzzer name.
   fuzz_target_name=pffft_${fft_transform,,}_fuzzer
+
+  # Add a symbolic link for the seed corpus (same corpus for all the
+  # generated fuzzers).
+  ls -la $OUT/*.zip
+  FUZZER_SEED_CORPUS_PATH=$OUT/${fuzz_target_name}_seed_corpus.zip
+  if [ -e $FUZZER_SEED_CORPUS_PATH ]; then
+    rm $FUZZER_SEED_CORPUS_PATH
+  fi
+  ln -s $SEED_CORPUS_ZIP_PATH $FUZZER_SEED_CORPUS_PATH
+
+  # Compile fuzzer.
   $CXX $CXXFLAGS -std=c++11 -msse2 \
        -DTRANSFORM_${fft_transform} \
        $SRC/pffft/pffft.c $SRC/pffft/pffft_fuzzer.cc \
diff --git a/projects/pffft/generate_seed_corpus.py b/projects/pffft/generate_seed_corpus.py
new file mode 100644
index 000000000..5e48cdbac
--- /dev/null
+++ b/projects/pffft/generate_seed_corpus.py
@@ -0,0 +1,59 @@
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+
+import numpy as np
+
+MAX_INPUT_SIZE = int(1e6)
+MAX_FLOAT32 = np.finfo(np.float32).max
+
+def IsValidSize(n):
+  if n == 0:
+    return False
+  # PFFFT only supports transforms for inputs of length N of the form
+  # N = (2^a)*(3^b)*(5^c) where a >= 5, b >=0, c >= 0.
+  FACTORS = [2, 3, 5]
+  factorization = [0, 0, 0]
+  for i, factor in enumerate(FACTORS):
+    while n % factor == 0:
+      n = n // factor
+      factorization[i] += 1
+  return factorization[0] >= 5 and n == 1
+
+
+def main():
+  if len(sys.argv) < 2:
+    print('Usage: %s <path to output directory>' % sys.argv[0])
+    sys.exit(1)
+
+  output_path = sys.argv[1]
+  if not os.path.exists(output_path):
+    print('The output path does not exists.')
+    sys.exit(2)
+
+  # List of valid input sizes.
+  N = [n for n in range(MAX_INPUT_SIZE) if IsValidSize(n)]
+
+  # Set the seed to always generate the same random data.
+  np.random.seed(0)
+
+  # Generate different types of input arrays for each target length.
+  for n in N:
+    # Zeros.
+    z = np.zeros(n, np.float32)
+    z.tofile(os.path.join(output_path, 'zeros_%d' % n))
+    # Max float 32.
+    m = np.ones(n, np.float32) * MAX_FLOAT32
+    m.tofile(os.path.join(output_path, 'max_%d' % n))
+    # Random values in the s16 range.
+    rnd_s16 = 32768.0 * 2.0 * (np.random.rand(n) - 1.0)
+    rnd_s16 = rnd_s16.astype(np.float32)
+    rnd_s16.tofile(os.path.join(output_path, 'rnd_s16_%d' % n))
+  
+  sys.exit(0)
+
+
+if __name__ == '__main__':
+  main()