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)))
for str_key, str_value in write_headers.iteritems():
info_data.write(THeaderTransport._serialize_string(str_key))
info_data.write(THeaderTransport._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 x in range(0, padding_size, 1):
buf.write(pack("!c", b'\0'))
# Send data section
buf.write(wout)
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 self.HEADERS_CLIENT_TYPE
@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
THeaderTransport._flush_info_headers(info_data,
self.get_write_persistent_headers(),
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
THeaderTransport._flush_info_headers(info_data,
self.__write_headers,
self.INFO_KEYVALUE)
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
wsz += header_size + 10
buf.write(pack(b"!I", wsz))
buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags))
buf.write(pack(b"!I", self.__seq_id))
buf.write(pack(b"!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 x in range(0, padding_size, 1):
buf.write(pack(b"!c", b'\0'))
# Send data section
buf.write(wout)
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for self.HEADERS_CLIENT_TYPE
@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
THeaderTransport._flush_info_headers(
info_data, self.get_write_persistent_headers(),
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
THeaderTransport._flush_info_headers(info_data, self.__write_headers,
self.INFO_KEYVALUE)
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
wsz += header_size + 10
buf.write(pack(b"!I", wsz))
buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags))
buf.write(pack(b"!I", self.__seq_id))
buf.write(pack(b"!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 x in range(0, padding_size, 1):
buf.write(pack(b"!c", b'\0'))
# Send data section
buf.write(wout)
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)))
if sys.version_info[0] >= 3:
write_headers_iter = write_headers.items()
else:
write_headers_iter = write_headers.iteritems()
for str_key, str_value in write_headers_iter:
info_data.write(THeaderTransport._serialize_string(str_key))
info_data.write(THeaderTransport._serialize_string(str_value))
write_headers.clear()
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)))
if sys.version_info[0] >= 3:
write_headers_iter = write_headers.items()
else:
write_headers_iter = write_headers.iteritems()
for str_key, str_value in write_headers_iter:
info_data.write(THeaderTransport._serialize_string(str_key))
info_data.write(THeaderTransport._serialize_string(str_value))
write_headers.clear()
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 _serialize_string(str_):
if sys.version_info[0] >= 3 and not isinstance(str_, bytes):
str_ = str_.encode()
return getVarint(len(str_)) + str_
def _serialize_string(str_):
if PY3 and not isinstance(str_, bytes):
str_ = str_.encode()
return getVarint(len(str_)) + str_
def _serialize_string(str_):
if PY3 and not isinstance(str_, bytes):
str_ = str_.encode()
return getVarint(len(str_)) + str_
def flush(self):
wout = self.__wbuf.getvalue()
wout = self.transform(wout)
wsz = len(wout)
# reset wbuf before write/flush to preserve state on underlying failure
self.__wbuf.seek(0)
self.__wbuf.truncate()
if self.__proto_id == 1 and self.__client_type != self.HTTP_CLIENT_TYPE:
raise TTransportException(TTransportException.INVALID_CLIENT_TYPE,
"Trying to send JSON encoding over binary")
buf = StringIO()
if self.__client_type == self.HEADERS_CLIENT_TYPE:
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))
if self.__hmac_func:
num_transforms += 1
transform_data.write(getVarint(self.HMAC_TRANSFORM))
transform_data.write(b'\0') # size of hmac, fixup later.
# 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
THeaderTransport._flush_info_headers(info_data,
self.__write_persistent_headers,
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
THeaderTransport._flush_info_headers(info_data,
self.__write_headers,
self.INFO_KEYVALUE)
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
wsz += header_size + 10
buf.write(pack(b"!I", wsz))
buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags))
buf.write(pack(b"!I", self.__seq_id))
buf.write(pack(b"!H", header_size // 4))
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
hmac_loc = buf.tell() - 1 # Fixup hmac size later
buf.write(info_data.getvalue())
# Pad out the header with 0x00
for x in range(0, padding_size, 1):
buf.write(pack(b"!c", b'\0'))
# Send data section
buf.write(wout)
# HMAC calculation should always be last.
if self.__hmac_func:
hmac_data = buf.getvalue()[4:]
hmac = self.__hmac_func(hmac_data)
# Fill in hmac size.
buf.seek(hmac_loc)
buf.write(chr(len(hmac)))
buf.seek(0, os.SEEK_END)
buf.write(hmac)
# Fix packet size since we appended data.
new_sz = buf.tell() - 4
buf.seek(0)
buf.write(pack(b"!I", new_sz))
elif self.__client_type == self.FRAMED_DEPRECATED:
buf.write(pack(b"!I", wsz))
buf.write(wout)
elif self.__client_type == self.UNFRAMED_DEPRECATED:
buf.write(wout)
elif self.__client_type == self.HTTP_CLIENT_TYPE:
# Reset the client type if we sent something -
# oneway calls via HTTP expect a status response otherwise
buf.write(self.header.getvalue())
buf.write(wout)
self.__client_type == self.HEADERS_CLIENT_TYPE
elif self.__client_type == self.UNKNOWN_CLIENT_TYPE:
raise TTransportException(TTransportException.INVALID_CLIENT_TYPE,
"Unknown client type")
# We don't include the framing bytes as part of the frame size check
# (so -4)
if buf.tell() - 4 > self.__max_frame_size:
raise TTransportException(TTransportException.INVALID_FRAME_SIZE,
"Attempting to send frame that is too large")
self.__trans.write(buf.getvalue())
self.__trans.flush()
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')
def _serialize_string(str_): return getVarint(len(str_)) + str_
def flush(self):
wout = self.__wbuf.getvalue()
wout = self.transform(wout)
wsz = len(wout)
# reset wbuf before write/flush to preserve state on underlying failure
self.__wbuf.seek(0)
self.__wbuf.truncate()
if self.__proto_id == 1 and self.__client_type != self.HTTP_CLIENT_TYPE:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Trying to send JSON encoding over binary")
buf = StringIO()
if self.__client_type == self.HEADERS_CLIENT_TYPE:
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))
if self.__hmac_func:
num_transforms += 1
transform_data.write(getVarint(self.HMAC_TRANSFORM))
transform_data.write(b'\0') # size of hmac, fixup later.
# 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
THeaderTransport._flush_info_headers(
info_data, self.__write_persistent_headers,
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
THeaderTransport._flush_info_headers(info_data,
self.__write_headers,
self.INFO_KEYVALUE)
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
wsz += header_size + 10
buf.write(pack(b"!I", wsz))
buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags))
buf.write(pack(b"!I", self.__seq_id))
buf.write(pack(b"!H", header_size // 4))
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
hmac_loc = buf.tell() - 1 # Fixup hmac size later
buf.write(info_data.getvalue())
# Pad out the header with 0x00
for x in range(0, padding_size, 1):
buf.write(pack(b"!c", b'\0'))
# Send data section
buf.write(wout)
# HMAC calculation should always be last.
if self.__hmac_func:
hmac_data = buf.getvalue()[4:]
hmac = self.__hmac_func(hmac_data)
# Fill in hmac size.
buf.seek(hmac_loc)
buf.write(chr(len(hmac)))
buf.seek(0, os.SEEK_END)
buf.write(hmac)
# Fix packet size since we appended data.
new_sz = buf.tell() - 4
buf.seek(0)
buf.write(pack(b"!I", new_sz))
elif self.__client_type == self.FRAMED_DEPRECATED:
buf.write(pack(b"!I", wsz))
buf.write(wout)
elif self.__client_type == self.UNFRAMED_DEPRECATED:
buf.write(wout)
elif self.__client_type == self.HTTP_CLIENT_TYPE:
# Reset the client type if we sent something -
# oneway calls via HTTP expect a status response otherwise
buf.write(self.header.getvalue())
buf.write(wout)
self.__client_type == self.HEADERS_CLIENT_TYPE
elif self.__client_type == self.UNKNOWN_CLIENT_TYPE:
raise TTransportException(TTransportException.INVALID_CLIENT_TYPE,
"Unknown client type")
# We don't include the framing bytes as part of the frame size check
# (so -4)
if buf.tell() - 4 > self.__max_frame_size:
raise TTransportException(
TTransportException.INVALID_FRAME_SIZE,
"Attempting to send frame that is too large")
self.__trans.write(buf.getvalue())
self.__trans.flush()
def _serialize_string(str_):
if sys.version_info[0] >= 3 and not isinstance(str_, bytes):
str_ = str_.encode()
return getVarint(len(str_)) + str_
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for self.HEADERS_CLIENT_TYPE
@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))
if self.__hmac_func:
num_transforms += 1
transform_data.write(getVarint(self.HMAC_TRANSFORM))
transform_data.write(b'\0') # size of hmac, fixup later.
# 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
THeaderTransport._flush_info_headers(info_data,
self.__write_persistent_headers,
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
THeaderTransport._flush_info_headers(info_data, self.__write_headers,
self.INFO_KEYVALUE)
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
wsz += header_size + 10
buf.write(pack(b"!I", wsz))
buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags))
buf.write(pack(b"!I", self.__seq_id))
buf.write(pack(b"!H", header_size // 4))
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
hmac_loc = buf.tell() - 1 # Fixup hmac size later
buf.write(info_data.getvalue())
# Pad out the header with 0x00
for x in range(0, padding_size, 1):
buf.write(pack(b"!c", b'\0'))
# Send data section
buf.write(wout)
# HMAC calculation should always be last.
if self.__hmac_func:
hmac_data = buf.getvalue()[4:]
hmac = self.__hmac_func(hmac_data)
# Fill in hmac size.
buf.seek(hmac_loc)
buf.write(chr(len(hmac)))
buf.seek(0, os.SEEK_END)
buf.write(hmac)
# Fix packet size since we appended data.
new_sz = buf.tell() - 4
buf.seek(0)
buf.write(pack(b"!I", new_sz))
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for self.HEADERS_CLIENT_TYPE
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))
if self._THeaderTransport__hmac_func:
num_transforms += 1
transform_data.write(getVarint(self.HMAC_TRANSFORM))
transform_data.write(b'\0') # size of hmac, fixup later.
# 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
cls._flush_info_headers(
info_data,
self._THeaderTransport__write_persistent_headers,
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
cls._flush_info_headers(
info_data,
self._THeaderTransport__write_headers,
self.INFO_KEYVALUE)
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', self.HEADER_MAGIC, 'i16')
self._write_fuzzy_field(
buf, 'flags', self._THeaderTransport__flags, 'i16')
self._write_fuzzy_field(
buf, 'seq_id', self._THeaderTransport__seq_id, 'i32')
self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16')
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
hmac_loc = buf.tell() - 1 # Fixup hmac size later
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')
# HMAC calculation should always be last.
if self._THeaderTransport__hmac_func:
hmac_data = buf.getvalue()[4:]
hmac = self._THeaderTransport__hmac_func(hmac_data)
# Fill in hmac size.
buf.seek(hmac_loc)
self._write_hmac_size(buf, hmac)
buf.seek(0, os.SEEK_END)
self._write_hmac(buf, hmac)
# Fix packet size since we appended data.
new_sz = buf.tell() - 4
buf.seek(0)
self._write_wsz(buf, new_sz)
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for self.HEADERS_CLIENT_TYPE
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))
if self._THeaderTransport__hmac_func:
num_transforms += 1
transform_data.write(getVarint(self.HMAC_TRANSFORM))
transform_data.write(b'\0') # size of hmac, fixup later.
# 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
cls._flush_info_headers(
info_data, self._THeaderTransport__write_persistent_headers,
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
cls._flush_info_headers(info_data,
self._THeaderTransport__write_headers,
self.INFO_KEYVALUE)
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', self.HEADER_MAGIC, 'i16')
self._write_fuzzy_field(buf, 'flags', self._THeaderTransport__flags,
'i16')
self._write_fuzzy_field(buf, 'seq_id', self._THeaderTransport__seq_id,
'i32')
self._write_fuzzy_field(buf, 'header_size', header_size // 4, 'i16')
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
hmac_loc = buf.tell() - 1 # Fixup hmac size later
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')
# HMAC calculation should always be last.
if self._THeaderTransport__hmac_func:
hmac_data = buf.getvalue()[4:]
hmac = self._THeaderTransport__hmac_func(hmac_data)
# Fill in hmac size.
buf.seek(hmac_loc)
self._write_hmac_size(buf, hmac)
buf.seek(0, os.SEEK_END)
self._write_hmac(buf, hmac)
# Fix packet size since we appended data.
new_sz = buf.tell() - 4
buf.seek(0)
self._write_wsz(buf, new_sz)
def _flushHeaderMessage(self, buf, wout, wsz):
"""Write a message for self.HEADERS_CLIENT_TYPE
@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))
if self.__hmac_func:
num_transforms += 1
transform_data.write(getVarint(self.HMAC_TRANSFORM))
transform_data.write(b'\0') # size of hmac, fixup later.
# 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
THeaderTransport._flush_info_headers(info_data,
self.__write_persistent_headers,
self.INFO_PKEYVALUE)
# Write non-persistent kv-headers
THeaderTransport._flush_info_headers(info_data,
self.__write_headers,
self.INFO_KEYVALUE)
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
wsz += header_size + 10
buf.write(pack(b"!I", wsz))
buf.write(pack(b"!HH", self.HEADER_MAGIC, self.__flags))
buf.write(pack(b"!I", self.__seq_id))
buf.write(pack(b"!H", header_size // 4))
buf.write(header_data.getvalue())
buf.write(transform_data.getvalue())
hmac_loc = buf.tell() - 1 # Fixup hmac size later
buf.write(info_data.getvalue())
# Pad out the header with 0x00
for x in range(0, padding_size, 1):
buf.write(pack(b"!c", b'\0'))
# Send data section
buf.write(wout)
# HMAC calculation should always be last.
if self.__hmac_func:
hmac_data = buf.getvalue()[4:]
hmac = self.__hmac_func(hmac_data)
# Fill in hmac size.
buf.seek(hmac_loc)
buf.write(chr(len(hmac)))
buf.seek(0, os.SEEK_END)
buf.write(hmac)
# Fix packet size since we appended data.
new_sz = buf.tell() - 4
buf.seek(0)
buf.write(pack(b"!I", new_sz))
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)