def get_method_randomizers(self, methods, constraints):
"""Create a StructRandomizer for each method"""
state = randomizer.RandomizerState()
method_randomizers = {}
state.push_type_constraints(constraints)
for method_name in methods:
method_constraints = constraints.get(method_name, {})
args_class = methods[method_name]['args_class']
# Create a spec_args tuple for the method args struct type
randomizer_spec_args = (
args_class,
args_class.thrift_spec,
False # isUnion
)
method_randomizer = state.get_randomizer(
Thrift.TType.STRUCT,
randomizer_spec_args,
method_constraints
)
method_randomizers[method_name] = method_randomizer
return method_randomizers
def is_fuzzed(cls, seeds, fuzzed, delta=None):
"""Check whether `fuzzed` could have been
generated by fuzzing any element of `seeds`"""
if delta is None:
# Find the default `fuzz_max_delta` constraint
randomizer_cls = randomizer.RandomizerState().get_randomizer(
cls.ttype, None, {})
delta = randomizer_cls.default_constraints['fuzz_max_delta']
return any(cls._is_fuzzed_single_seed(seed, fuzzed, delta)
for seed in seeds)
def get_randomizer(self, ttypes, spec_args, constraints):
state = randomizer.RandomizerState()
return state.get_randomizer(ttypes, spec_args, constraints)
class TFuzzyHeaderTransport(THeaderTransport):
"""Transport that can optionally fuzz fields in the header or payload.
The initializer accepts `fuzz_fields`, an iterable of names of fields in
the header (or payload) that should be fuzzed (mutated.)
TFuzzyHeaderTransport should ONLY be used for testing purposes.
"""
fuzzable_fields = frozenset({
'flags',
'header_size',
'identity',
'length',
'magic',
'num_transforms',
'padding',
'padding_bytes',
'payload',
'proto_id',
'seq_id',
'transform_id',
'version',
})
# Randomizers for message fields by data type
_randomizer_state = randomizer.RandomizerState()
_i32_randomizer = _randomizer_state.get_randomizer(Thrift.TType.I32, None, {
'fuzz_max_delta': 32
})
_i16_randomizer = _randomizer_state.get_randomizer(Thrift.TType.I16, None, {
'fuzz_max_delta': 32
})
_byte_randomizer = _randomizer_state.get_randomizer(
Thrift.TType.BYTE, None, {})
@classmethod
def _fuzz_char(cls, seed):
"""Randomly modify a (binary) character"""
seed = six.byte2int(seed)
res = cls._byte_randomizer.generate(seed=seed) % (2 ** 8)
return six.int2byte(res)
@classmethod
def _fuzz_i16(cls, seed):
"""Randomly modify a 16-bit positive integer"""
res = cls._i16_randomizer.generate(seed=seed)
return res % (2 ** 16)
@classmethod
def _fuzz_i32(cls, seed):
"""Randomly modify a 32-bit positive integer"""
res = cls._i32_randomizer.generate(seed=seed)
return res % (2 ** 32)
@classmethod
def _fuzz_str(cls, seed):
"""Randomly modify one character of a string"""
fuzz_index = random.randint(0, len(seed) - 1)
fuzzed_char = six.int2byte(random.randint(0, 255))
return seed[:fuzz_index] + fuzzed_char + seed[fuzz_index + 1:]
# Map type names to arguments to struct.pack
_struct_type_map = {
'char': '!c',
'i16': '!H',
'i32': '!I'
}
@classmethod
def _serialize_field(cls, val, type_):
if type_ == 'str':
return six.binary_type(val)
if type_ in cls._struct_type_map:
struct_type = six.binary_type(cls._struct_type_map[type_])
return struct.pack(struct_type, val)
else:
raise NotImplementedError("%s serializer not implemented" % type_)
def __init__(self, trans, client_types=None,
fuzz_fields=None, fuzz_all_if_empty=True, verbose=False):
"""Create a TFuzzyHeaderTransport instance.
@param fuzz_fields(iterable) Collection of string names of fields to
fuzz. Each name must be a member of
TFuzzyHeaderTransport.fuzzable_fields. If the iterable is empty and
fuzz_all_if_empty is True, all fields will be fuzzed. If the iterable
is empty and fuzz_all_if_empty is False, no fields will be fuzzed.
@param verbose(bool) Whether to print fuzzed fields to stdout
@param fuzz_all_if_empty(bool) Whether to fuzz eveyr field if
the fuzz_fields iterable is None or empty.
trans and client_type are forwarded to THeaderProtocol.__init__
"""
cls = self.__class__
THeaderTransport.__init__(self, trans, client_types=client_types)
if fuzz_fields is None:
fuzz_fields = []
self._fuzz_fields = set()
for field_name in fuzz_fields:
if field_name not in cls.fuzzable_fields:
raise NameError("Invalid fuzz field: %s" % field_name)
self._fuzz_fields.add(field_name)
if not self._fuzz_fields and fuzz_all_if_empty:
# No fields were explicitly included
self._fuzz_fields = cls.fuzzable_fields
self._verbose = verbose
def _print(self, s):
if self._verbose:
print(s)
def _get_fuzzy_field(self, name, val, type_):
"""Return a possibly-fuzzed version of val, depending on self._fuzz_set
If `name` is included in the set of fields to fuzz, the field value
will be fuzzed. Otherwise, val will be returned unmodified."""
cls = self.__class__
if name not in self._fuzz_fields:
return val
if type_ == 'char':
fuzzed_val = cls._fuzz_char(val)
elif type_ == 'i16':
fuzzed_val = cls._fuzz_i16(val)
elif type_ == 'i32':
fuzzed_val = cls._fuzz_i32(val)
elif type_ == 'str':
fuzzed_val = cls._fuzz_str(val)
else:
raise NotImplementedError("%s fuzzing not implemented" % type_)
self._print("Fuzzed field %s from %r to %r" % (name, val, fuzzed_val))
return fuzzed_val
def _write_fuzzy_field(self, buf, name, val, type_):
"""Write a message field after (possibly) mutating the value
@param buf(StringIO): Message buffer to write to
@param name(str): Name of the field (for checking whether to fuzz)
@param val(type_): Value to write to field
@param type_(type): Type of field (for fuzzing and serializing)
"""
val = self._get_fuzzy_field(name, val, type_)
serialized = self._serialize_field(val, type_)
buf.write(serialized)
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for CLIENT_TYPE.HEADER
This method writes a message using the same logic as
THeaderTransport._flushHeaderMessage
but mutates fields included in self._fuzz_fields
"""
cls = self.__class__
transform_data = StringIO()
num_transforms = len(self._THeaderTransport__write_transforms)
for trans_id in self._THeaderTransport__write_transforms:
trans_id = self._get_fuzzy_field('transform_id', trans_id, 'i32')
transform_data.write(getVarint(trans_id))
# Add in special flags.
if self._THeaderTransport__identity:
id_version = self._get_fuzzy_field(
'version', self.ID_VERSION, 'str')
self._THeaderTransport__write_headers[self.ID_VERSION_HEADER] = (
id_version)
identity = self._get_fuzzy_field(
'identity', self._THeaderTransport__identity, 'str')
self._THeaderTransport__write_headers[self.IDENTITY_HEADER] = (
identity)
info_data = StringIO()
# Write persistent kv-headers
_flush_info_headers(
info_data,
self._THeaderTransport__write_persistent_headers,
INFO.PERSISTENT)
# Write non-persistent kv-headers
_flush_info_headers(
info_data,
self._THeaderTransport__write_headers,
INFO.NORMAL)
header_data = StringIO()
proto_id = self._get_fuzzy_field(
'proto_id', self._THeaderTransport__proto_id, 'i32')
header_data.write(getVarint(proto_id))
num_transforms = self._get_fuzzy_field(
'num_transforms', num_transforms, 'i32')
header_data.write(getVarint(num_transforms))
header_size = (transform_data.tell() +
header_data.tell() +
info_data.tell())
padding_size = 4 - (header_size % 4)
# Fuzz padding size, but do not let total header size exceed 2**16 - 1
padding_size = min((2 ** 16 - 1) - header_size,
self._get_fuzzy_field('padding', padding_size, 'i16'))
header_size = header_size + padding_size
wsz += header_size + 10
self._write_fuzzy_field(buf, 'length', wsz, 'i32')
self._write_fuzzy_field(buf, 'magic', HEADER_MAGIC >> 16, 'i16')
self._write_fuzzy_field(
buf, 'flags', self.header_flags(), 'i16')
self._write_fuzzy_field(
buf, 'seq_id', self.seq_id, 'i32')
self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16')
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
buf.write(info_data.getvalue())
# Pad out the header with 0x00
if 'padding_bytes' in self._fuzz_fields:
# Print that padding bytes are being fuzzed in favor of printing
# the value of each individual padding byte
self._print("Fuzzing %d padding bytes" % padding_size)
old_verbose, self._verbose = self._verbose, False
for _ in sm.xrange(padding_size):
self._write_fuzzy_field(
buf, 'padding_bytes', six.int2byte(0), 'char')
self._verbose = old_verbose
self._write_fuzzy_field(buf, 'payload', wout, 'str')
def get_randomizer(self, ttypes, spec_args, constraints):
state = randomizer.RandomizerState({"global_constraint": 100})
return state.get_randomizer(ttypes, spec_args, constraints)
