Improve pygment's fuzzer (#6212)

4 files changed, 47 insertions, 12 deletions

diff --git a/projects/pygments/Dockerfile b/projects/pygments/Dockerfile
index b23fc1b6d..3c750be6c 100644
--- a/projects/pygments/Dockerfile
+++ b/projects/pygments/Dockerfile
@@ -58,4 +58,4 @@ RUN cat fuzzing/dictionaries/aff.dict \
         fuzzing/dictionaries/yara.dict \
       > $OUT/pygments_fuzzer.dict
 
-COPY build.sh pygments_fuzzer.py $SRC/
+COPY build.sh fuzz_guesser.py fuzz_lexers.py $SRC/
diff --git a/projects/pygments/build.sh b/projects/pygments/build.sh
index 4e59eb6ae..b76e29aa9 100644
--- a/projects/pygments/build.sh
+++ b/projects/pygments/build.sh
@@ -19,7 +19,7 @@
 pip3 install .
 
 # Build fuzzers in $OUT.
-for fuzzer in $(find $SRC -name '*_fuzzer.py'); do
+for fuzzer in $(find $SRC -name 'fuzz_*.py'); do
   fuzzer_basename=$(basename -s .py $fuzzer)
   fuzzer_package=${fuzzer_basename}.pkg
   pyinstaller --distpath $OUT --onefile --name $fuzzer_package $fuzzer
diff --git a/projects/pygments/pygments_fuzzer.py b/projects/pygments/fuzz_guesser.py
index feea49934..ecc2411cf 100644
--- a/projects/pygments/pygments_fuzzer.py
+++ b/projects/pygments/fuzz_guesser.py
@@ -14,24 +14,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
 import atheris
 with atheris.instrument_imports():
+  import sys
   import pygments
-  import pygments.formatters
   import pygments.lexers
 
-
 @atheris.instrument_func
-def TestOneInput(input_bytes):
-  fdp = atheris.FuzzedDataProvider(input_bytes)
-  data = fdp.ConsumeUnicode(atheris.ALL_REMAINING)
-
+def TestOneInput(data: bytes) -> int:
   try:
-    lexer = pygments.lexers.guess_lexer(data)
+    lexer = pygments.lexers.guess_lexer(str(data))
   except ValueError:
-    return
-  pygments.highlight(data, lexer, pygments.formatters.HtmlFormatter())
+    return 0
+  return 0
 
 
 def main():
diff --git a/projects/pygments/fuzz_lexers.py b/projects/pygments/fuzz_lexers.py
new file mode 100644
index 000000000..0d610c103
--- /dev/null
+++ b/projects/pygments/fuzz_lexers.py
@@ -0,0 +1,40 @@
+#!/usr/bin/python3
+
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import atheris
+with atheris.instrument_imports():
+  import sys
+  import pygments
+  import pygments.formatters.html
+  import pygments.lexers
+
+formatter = pygments.formatters.html.HtmlFormatter()
+# pygments.LEXERS.values() is a list of tuples like this, with some of then empty:
+# (textual class name, longname, tuple of aliases, tuple of filename patterns, tuple of mimetypes)
+LEXERS = [l[2][0] for l in pygments.lexers.LEXERS.values() if l[2]]
+
+
+def TestOneInput(data: bytes) -> int:
+  fdp = atheris.FuzzedDataProvider(data)
+  random_lexer = pygments.lexers.get_lexer_by_name(fdp.PickValueInList(LEXERS))
+  str_data = fdp.ConsumeUnicode(atheris.ALL_REMAINING)
+
+  pygments.highlight(str_data, random_lexer, formatter)
+  return 0
+
+
+atheris.Setup(sys.argv, TestOneInput)
+atheris.Fuzz()