Add a Python fuzzer for tf.raw_ops.ImmutableConst.

PiperOrigin-RevId: 354425432
Change-Id: I9d07ef399031a91321446590f6ae86a5340a4e65

tensorflow/security/fuzzing/immutableConst_fuzz.py

@@ -0,0 +1,46 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""This is a Python API fuzzer for tf.raw_ops.ImmutableConst."""
+import sys
+import atheris_no_libfuzzer as atheris
+from python_fuzzing import FuzzingHelper
+import tensorflow as tf
+
+_DEFAULT_FILENAME = '/tmp/test.txt'
+
+
+def TestOneInput(input_bytes):
+  """Test randomized integer fuzzing input for tf.raw_ops.ImmutableConst."""
+  fh = FuzzingHelper(input_bytes)
+
+  dtype = fh.get_tf_dtype()
+  shape = fh.get_int_list()
+  try:
+    with open(_DEFAULT_FILENAME, 'w') as f:
+      f.write(fh.get_string())
+    _ = tf.raw_ops.ImmutableConst(
+        dtype=dtype, shape=shape, memory_region_name=_DEFAULT_FILENAME)
+  except (tf.errors.InvalidArgumentError, tf.errors.InternalError,
+          UnicodeEncodeError, UnicodeDecodeError):
+    pass
+
+
+def main():
+  atheris.Setup(sys.argv, TestOneInput, enable_python_coverage=True)
+  atheris.Fuzz()
+
+
+if __name__ == '__main__':
+  main()

tensorflow/security/fuzzing/python_fuzzing.py

@@ -15,6 +15,7 @@
 """Helper class for TF Python fuzzing."""
 
 import atheris_no_libfuzzer as atheris
+import tensorflow as tf
 
 _MIN_INT = -10000
 _MAX_INT = 10000
@@ -25,6 +26,13 @@ _MAX_FLOAT = 10000.0
 _MIN_LENGTH = 0
 _MAX_LENGTH = 10000
 
+_TF_DTYPES = [
+    tf.float16, tf.float32, tf.float64, tf.bfloat16, tf.complex64,
+    tf.complex128, tf.int8, tf.uint8, tf.uint16, tf.uint32, tf.uint64, tf.int16,
+    tf.int32, tf.int64, tf.bool, tf.string, tf.qint8, tf.quint8, tf.qint16,
+    tf.quint16, tf.qint32, tf.resource, tf.variant
+]
+
 
 class FuzzingHelper(object):
   """FuzzingHelper makes handling FuzzedDataProvider easier with TensorFlow Python fuzzing."""
@@ -117,3 +125,23 @@ class FuzzingHelper(object):
       return self.get_int_list(min_length, max_length)
     else:
       return self.get_float_list(min_length, max_length)
+
+  def get_tf_dtype(self):
+    """Return a random tensorflow type.
+
+    Returns:
+      A random type from the list containing all TensorFlow types.
+    """
+    index = self.get_int(0, len(_TF_DTYPES) - 1)
+    return _TF_DTYPES[index]
+
+  def get_string(self, byte_count=_MAX_INT):
+    """Consume a string with given constraints based on a consumed bool.
+
+    Args:
+      byte_count: Byte count that defaults to _MAX_INT.
+
+    Returns:
+      Consumed string based on input bytes and constraints.
+    """
+    return self.fdp.ConsumeString(byte_count)

tensorflow/security/fuzzing/raggedCountSparseOutput_fuzz.py

@@ -20,7 +20,7 @@ import tensorflow as tf
 
 
 def TestOneInput(input_bytes):
-  """Test randomized integer fuzzing input for tf.raw_ops.RaggedCountSparseOutput."""
+  """Test randomized integer/float fuzzing input for tf.raw_ops.RaggedCountSparseOutput."""
   fh = FuzzingHelper(input_bytes)
 
   splits = fh.get_int_list()