1# Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Helper class for TF Python fuzzing.""" 16 17import atheris_no_libfuzzer as atheris 18import tensorflow as tf 19 20_MIN_INT = -10000 21_MAX_INT = 10000 22 23_MIN_FLOAT = -10000.0 24_MAX_FLOAT = 10000.0 25 26_MIN_LENGTH = 0 27_MAX_LENGTH = 10000 28 29# Max shape can be 8 in length and randomized from 0-8 without running into an 30# OOM error. 31_MIN_SIZE = 0 32_MAX_SIZE = 8 33 34_TF_DTYPES = [ 35 tf.half, tf.float16, tf.float32, tf.float64, tf.bfloat16, tf.complex64, 36 tf.complex128, tf.int8, tf.uint8, tf.uint16, tf.uint32, tf.uint64, tf.int16, 37 tf.int32, tf.int64, tf.bool, tf.string, tf.qint8, tf.quint8, tf.qint16, 38 tf.quint16, tf.qint32, tf.resource, tf.variant 39] 40 41# All types supported by tf.random.uniform 42_TF_RANDOM_DTYPES = [tf.float16, tf.float32, tf.float64, tf.int32, tf.int64] 43 44 45class FuzzingHelper(object): 46 """FuzzingHelper makes handling FuzzedDataProvider easier with TensorFlow Python fuzzing.""" 47 48 def __init__(self, input_bytes): 49 """FuzzingHelper initializer. 50 51 Args: 52 input_bytes: Input randomized bytes used to create a FuzzedDataProvider. 53 """ 54 self.fdp = atheris.FuzzedDataProvider(input_bytes) 55 56 def get_bool(self): 57 """Consume a bool. 58 59 Returns: 60 Consumed a bool based on input bytes and constraints. 61 """ 62 return self.fdp.ConsumeBool() 63 64 def get_int(self, min_int=_MIN_INT, max_int=_MAX_INT): 65 """Consume a signed integer with given constraints. 66 67 Args: 68 min_int: Minimum allowed integer. 69 max_int: Maximum allowed integer. 70 71 Returns: 72 Consumed integer based on input bytes and constraints. 73 """ 74 return self.fdp.ConsumeIntInRange(min_int, max_int) 75 76 def get_float(self, min_float=_MIN_FLOAT, max_float=_MAX_FLOAT): 77 """Consume a float with given constraints. 78 79 Args: 80 min_float: Minimum allowed float. 81 max_float: Maximum allowed float. 82 83 Returns: 84 Consumed float based on input bytes and constraints. 85 """ 86 return self.fdp.ConsumeFloatInRange(min_float, max_float) 87 88 def get_int_list(self, 89 min_length=_MIN_LENGTH, 90 max_length=_MAX_LENGTH, 91 min_int=_MIN_INT, 92 max_int=_MAX_INT): 93 """Consume a signed integer list with given constraints. 94 95 Args: 96 min_length: The minimum length of the list. 97 max_length: The maximum length of the list. 98 min_int: Minimum allowed integer. 99 max_int: Maximum allowed integer. 100 101 Returns: 102 Consumed integer list based on input bytes and constraints. 103 """ 104 length = self.get_int(min_length, max_length) 105 return self.fdp.ConsumeIntListInRange(length, min_int, max_int) 106 107 def get_float_list(self, min_length=_MIN_LENGTH, max_length=_MAX_LENGTH): 108 """Consume a float list with given constraints. 109 110 Args: 111 min_length: The minimum length of the list. 112 max_length: The maximum length of the list. 113 114 Returns: 115 Consumed integer list based on input bytes and constraints. 116 """ 117 length = self.get_int(min_length, max_length) 118 return self.fdp.ConsumeFloatListInRange(length, _MIN_FLOAT, _MAX_FLOAT) 119 120 def get_int_or_float_list(self, 121 min_length=_MIN_LENGTH, 122 max_length=_MAX_LENGTH): 123 """Consume a signed integer or float list with given constraints based on a consumed bool. 124 125 Args: 126 min_length: The minimum length of the list. 127 max_length: The maximum length of the list. 128 129 Returns: 130 Consumed integer or float list based on input bytes and constraints. 131 """ 132 if self.get_bool(): 133 return self.get_int_list(min_length, max_length) 134 else: 135 return self.get_float_list(min_length, max_length) 136 137 def get_tf_dtype(self, allowed_set=None): 138 """Return a random tensorflow dtype. 139 140 Args: 141 allowed_set: An allowlisted set of dtypes to choose from instead of all of 142 them. 143 144 Returns: 145 A random type from the list containing all TensorFlow types. 146 """ 147 if allowed_set: 148 index = self.get_int(0, len(allowed_set) - 1) 149 if allowed_set[index] not in _TF_DTYPES: 150 raise tf.errors.InvalidArgumentError( 151 None, None, 152 'Given dtype {} is not accepted.'.format(allowed_set[index])) 153 return allowed_set[index] 154 else: 155 index = self.get_int(0, len(_TF_DTYPES) - 1) 156 return _TF_DTYPES[index] 157 158 def get_string(self, byte_count=_MAX_INT): 159 """Consume a string with given constraints based on a consumed bool. 160 161 Args: 162 byte_count: Byte count that defaults to _MAX_INT. 163 164 Returns: 165 Consumed string based on input bytes and constraints. 166 """ 167 return self.fdp.ConsumeString(byte_count) 168 169 def get_random_numeric_tensor(self, 170 dtype=None, 171 min_size=_MIN_SIZE, 172 max_size=_MAX_SIZE, 173 min_val=_MIN_INT, 174 max_val=_MAX_INT): 175 """Return a tensor of random shape and values. 176 177 Generated tensors are capped at dimension sizes of 8, as 2^32 bytes of 178 requested memory crashes the fuzzer (see b/34190148). 179 Returns only type that tf.random.uniform can generate. If you need a 180 different type, consider using tf.cast. 181 182 Args: 183 dtype: Type of tensor, must of one of the following types: float16, 184 float32, float64, int32, or int64 185 min_size: Minimum size of returned tensor 186 max_size: Maximum size of returned tensor 187 min_val: Minimum value in returned tensor 188 max_val: Maximum value in returned tensor 189 190 Returns: 191 Tensor of random shape filled with uniformly random numeric values. 192 """ 193 # Max shape can be 8 in length and randomized from 0-8 without running into 194 # an OOM error. 195 if max_size > 8: 196 raise tf.errors.InvalidArgumentError( 197 None, None, 198 'Given size of {} will result in an OOM error'.format(max_size)) 199 200 seed = self.get_int() 201 shape = self.get_int_list( 202 min_length=min_size, 203 max_length=max_size, 204 min_int=min_size, 205 max_int=max_size) 206 207 if dtype is None: 208 dtype = self.get_tf_dtype(allowed_set=_TF_RANDOM_DTYPES) 209 elif dtype not in _TF_RANDOM_DTYPES: 210 raise tf.errors.InvalidArgumentError( 211 None, None, 212 'Given dtype {} is not accepted in get_random_numeric_tensor'.format( 213 dtype)) 214 215 return tf.random.uniform( 216 shape=shape, minval=min_val, maxval=max_val, dtype=dtype, seed=seed) 217