def _flush_info_headers(info_data, write_headers, type):
if (len(write_headers) > 0):
info_data.write(getVarint(type))
info_data.write(getVarint(len(write_headers)))
write_headers_iter = write_headers.items()
for str_key, str_value in write_headers_iter:
info_data.write(_serialize_string(str_key))
info_data.write(_serialize_string(str_value))
write_headers.clear()
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for CLIENT_TYPE.HEADER
@param buf(StringIO): Buffer to write message to
@param wout(str): Payload
@param wsz(int): Payload length
"""
transform_data = StringIO()
# For now, all transforms don't require data.
num_transforms = len(self.__write_transforms)
for trans_id in self.__write_transforms:
transform_data.write(getVarint(trans_id))
# Add in special flags.
if self.__identity:
self.__write_headers[self.ID_VERSION_HEADER] = self.ID_VERSION
self.__write_headers[self.IDENTITY_HEADER] = self.__identity
info_data = StringIO()
# Write persistent kv-headers
_flush_info_headers(info_data,
self.get_write_persistent_headers(),
INFO.PERSISTENT)
# Write non-persistent kv-headers
_flush_info_headers(info_data,
self.__write_headers,
INFO.NORMAL)
header_data = StringIO()
header_data.write(getVarint(self.__proto_id))
header_data.write(getVarint(num_transforms))
header_size = transform_data.tell() + header_data.tell() + \
info_data.tell()
padding_size = 4 - (header_size % 4)
header_size = header_size + padding_size
# MAGIC(2) | FLAGS(2) + SEQ_ID(4) + HEADER_SIZE(2)
wsz += header_size + 10
if wsz > MAX_FRAME_SIZE:
buf.write(pack("!I", BIG_FRAME_MAGIC))
buf.write(pack("!Q", wsz))
else:
buf.write(pack("!I", wsz))
buf.write(pack("!HH", HEADER_MAGIC >> 16, self.__flags))
buf.write(pack("!I", self.seq_id))
buf.write(pack("!H", header_size // 4))
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
buf.write(info_data.getvalue())
# Pad out the header with 0x00
for _ in range(0, padding_size, 1):
buf.write(pack("!c", b'\0'))
# Send data section
buf.write(wout)
def _serialize_string(str_):
if PY3 and not isinstance(str_, bytes):
str_ = str_.encode()
return getVarint(len(str_)) + str_
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
"""
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')