From f770567861dc0c92078b23e58086f399ad56fc4a Mon Sep 17 00:00:00 2001
From: DavidKorczynski <david@adalogics.com>
Date: Tue, 19 Jul 2022 19:30:57 +0100
Subject: [PATCH] g-api-auth-library-python: refine set up (#8029)

Split the existing fuzzer. This adds significant speedup in one of the
code areas being tested. Also refines the setup to add more code paths
in the encoding logic.
---
 projects/g-api-auth-library-python/build.sh   |  1 +
 .../g-api-auth-library-python/fuzz_jwt.py     | 26 ++-----
 .../fuzz_jwt_roundtrip.py                     | 70 +++++++++++++++++++
 3 files changed, 77 insertions(+), 20 deletions(-)
 create mode 100644 projects/g-api-auth-library-python/fuzz_jwt_roundtrip.py

diff --git a/projects/g-api-auth-library-python/build.sh b/projects/g-api-auth-library-python/build.sh
index 841b85019..1f44fe0eb 100644
--- a/projects/g-api-auth-library-python/build.sh
+++ b/projects/g-api-auth-library-python/build.sh
@@ -19,6 +19,7 @@
 cp tests/data/*.pem $OUT/
 cp tests/data/*.pub $OUT/
 
+pip3 install -r docs/requirements-docs.txt
 pip3 install .
 
 for fuzzer in $(find $SRC -name 'fuzz_*.py'); do
diff --git a/projects/g-api-auth-library-python/fuzz_jwt.py b/projects/g-api-auth-library-python/fuzz_jwt.py
index cbc94be1e..e167e2cb3 100644
--- a/projects/g-api-auth-library-python/fuzz_jwt.py
+++ b/projects/g-api-auth-library-python/fuzz_jwt.py
@@ -16,15 +16,9 @@
 import atheris
 import sys
 
-with atheris.instrument_imports():
-  from google.auth import jwt
-  from google.auth import crypt
+from google.auth import jwt
+from google.auth import crypt
 
-if os.path.isfile("privatekey.pem"):
-  with open("privatekey.pem", "rb") as fh:
-    PRIVATE_KEY_BYTES = fh.read()
-else:
-  raise Exception("Could not find private key")
 
 if os.path.isfile("public_cert.pem"):
   with open("public_cert.pem", "rb") as fh:
@@ -32,16 +26,7 @@ if os.path.isfile("public_cert.pem"):
 else:
   raise Exception("Could not find public cert")
 
-
-def test_roundtrip_unverified(data):
-  fdp = atheris.FuzzedDataProvider(data)
-  raw_data = fdp.ConsumeString(200)
-  signer = crypt.RSASigner.from_string(PRIVATE_KEY_BYTES, "1")
-  encoded = jwt.encode(signer, raw_data)
-
-  _, decoded_data, _, _ = jwt._unverified_decode(encoded)
-  assert(raw_data == decoded_data)
-
+@atheris.instrument_func
 def test_token_decode(data):
   fdp = atheris.FuzzedDataProvider(data)
   try:
@@ -49,13 +34,14 @@ def test_token_decode(data):
   except ValueError:  # ValueError is thrown if any failed verification checks
     pass
 
+@atheris.instrument_func
 def TestOneInput(data):
-  test_roundtrip_unverified(data)
   test_token_decode(data)
 
+
 def main():
-  atheris.instrument_all()
   atheris.Setup(sys.argv, TestOneInput, enable_python_coverage=True)
+  atheris.instrument_all()
   atheris.Fuzz()
 
 if __name__ == "__main__":
diff --git a/projects/g-api-auth-library-python/fuzz_jwt_roundtrip.py b/projects/g-api-auth-library-python/fuzz_jwt_roundtrip.py
new file mode 100644
index 000000000..50d878eb3
--- /dev/null
+++ b/projects/g-api-auth-library-python/fuzz_jwt_roundtrip.py
@@ -0,0 +1,70 @@
+#!/usr/bin/python3
+# Copyright 2022 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
+import sys
+
+# We instrument all imports below
+from google.auth import jwt
+from google.auth import crypt
+
+if os.path.isfile("privatekey.pem"):
+  with open("privatekey.pem", "rb") as fh:
+    PRIVATE_KEY_BYTES = fh.read()
+else:
+  raise Exception("Could not find private key")
+
+if os.path.isfile("public_cert.pem"):
+  with open("public_cert.pem", "rb") as fh:
+    PUBLIC_CERT_BYTES = fh.read()
+else:
+  raise Exception("Could not find public cert")
+
+@atheris.instrument_func
+def test_roundtrip_unverified(data):
+  fdp = atheris.FuzzedDataProvider(data)
+  signer = crypt.RSASigner.from_string(PRIVATE_KEY_BYTES, "1")
+
+  to_header = fdp.ConsumeIntInRange(1, 100)
+  if to_header < 50:
+    header = None
+  else:
+    header = {
+      "alg" : fdp.ConsumeString(100),
+    }
+  to_keyid = fdp.ConsumeIntInRange(1, 100)
+  raw_data = fdp.ConsumeString(200)
+
+  key_id = fdp.ConsumeString(50) if to_keyid < 50 else None
+  encoded = jwt.encode(signer, raw_data, header = header, key_id = key_id)
+  try:
+    _, decoded_data, _, _ = jwt.decode(encoded, PUBLIC_CERT_BYTES)
+  except ValueError as e:
+    return
+
+
+@atheris.instrument_func
+def TestOneInput(data):
+  test_roundtrip_unverified(data)
+
+
+def main():
+  atheris.Setup(sys.argv, TestOneInput, enable_python_coverage=True)
+  atheris.instrument_all()
+  atheris.Fuzz()
+
+
+if __name__ == "__main__":
+  main()