def flush(self):
payload = self._write_buffer.getvalue()
self._write_buffer = BufferIO()
buffer = BufferIO()
if self._client_type == THeaderClientType.HEADERS:
for transform_id in self._write_transforms:
transform_fn = WRITE_TRANSFORMS_BY_ID[transform_id]
payload = transform_fn(payload)
headers = BufferIO()
writeVarint(headers, self._protocol_id)
writeVarint(headers, len(self._write_transforms))
for transform_id in self._write_transforms:
writeVarint(headers, transform_id)
if self._write_headers:
writeVarint(headers, TInfoHeaderType.KEY_VALUE)
writeVarint(headers, len(self._write_headers))
for key, value in self._write_headers.items():
_writeString(headers, key)
_writeString(headers, value)
self._write_headers = {}
padding_needed = (4 - (len(headers.getvalue()) % 4)) % 4
headers.write(b"\x00" * padding_needed)
header_bytes = headers.getvalue()
buffer.write(I32.pack(10 + len(header_bytes) + len(payload)))
buffer.write(U16.pack(HEADER_MAGIC))
buffer.write(U16.pack(self.flags))
buffer.write(I32.pack(self.sequence_id))
buffer.write(U16.pack(len(header_bytes) // 4))
buffer.write(header_bytes)
buffer.write(payload)
elif self._client_type in (THeaderClientType.FRAMED_BINARY,
THeaderClientType.FRAMED_COMPACT):
buffer.write(I32.pack(len(payload)))
buffer.write(payload)
elif self._client_type in (THeaderClientType.UNFRAMED_BINARY,
THeaderClientType.UNFRAMED_COMPACT):
buffer.write(payload)
else:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Unknown client type.",
)
# the frame length field doesn't count towards the frame payload size
frame_bytes = buffer.getvalue()
frame_payload_size = len(frame_bytes) - 4
if frame_payload_size > self._max_frame_size:
raise TTransportException(
TTransportException.SIZE_LIMIT,
"Attempting to send frame that is too large.",
)
self._transport.write(frame_bytes)
self._transport.flush()
class THeaderTransport(TTransportBase, CReadableTransport):
def __init__(self, transport, allowed_client_types):
self._transport = transport
self._client_type = THeaderClientType.HEADERS
self._allowed_client_types = allowed_client_types
self._read_buffer = BufferIO(b"")
self._read_headers = {}
self._write_buffer = BufferIO()
self._write_headers = {}
self._write_transforms = []
self.flags = 0
self.sequence_id = 0
self._protocol_id = THeaderSubprotocolID.BINARY
self._max_frame_size = HARD_MAX_FRAME_SIZE
def isOpen(self):
return self._transport.isOpen()
def open(self):
return self._transport.open()
def close(self):
return self._transport.close()
def get_headers(self):
return self._read_headers
def set_header(self, key, value):
if not isinstance(key, bytes):
raise ValueError("header names must be bytes")
if not isinstance(value, bytes):
raise ValueError("header values must be bytes")
self._write_headers[key] = value
def clear_headers(self):
self._write_headers.clear()
def add_transform(self, transform_id):
if transform_id not in WRITE_TRANSFORMS_BY_ID:
raise ValueError("unknown transform")
self._write_transforms.append(transform_id)
def set_max_frame_size(self, size):
if not 0 < size < HARD_MAX_FRAME_SIZE:
raise ValueError("maximum frame size should be < %d and > 0" %
HARD_MAX_FRAME_SIZE)
self._max_frame_size = size
@property
def protocol_id(self):
if self._client_type == THeaderClientType.HEADERS:
return self._protocol_id
elif self._client_type in (THeaderClientType.FRAMED_BINARY,
THeaderClientType.UNFRAMED_BINARY):
return THeaderSubprotocolID.BINARY
elif self._client_type in (THeaderClientType.FRAMED_COMPACT,
THeaderClientType.UNFRAMED_COMPACT):
return THeaderSubprotocolID.COMPACT
else:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Protocol ID not know for client type %d" % self._client_type,
)
def read(self, sz):
# if there are bytes left in the buffer, produce those first.
bytes_read = self._read_buffer.read(sz)
bytes_left_to_read = sz - len(bytes_read)
if bytes_left_to_read == 0:
return bytes_read
# if we've determined this is an unframed client, just pass the read
# through to the underlying transport until we're reset again at the
# beginning of the next message.
if self._client_type in (THeaderClientType.UNFRAMED_BINARY,
THeaderClientType.UNFRAMED_COMPACT):
return bytes_read + self._transport.read(bytes_left_to_read)
# we're empty and (maybe) framed. fill the buffers with the next frame.
self.readFrame(bytes_left_to_read)
return bytes_read + self._read_buffer.read(bytes_left_to_read)
def _set_client_type(self, client_type):
if client_type not in self._allowed_client_types:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Client type %d not allowed by server." % client_type,
)
self._client_type = client_type
def readFrame(self, req_sz):
# the first word could either be the length field of a framed message
# or the first bytes of an unframed message.
first_word = self._transport.readAll(I32.size)
frame_size, = I32.unpack(first_word)
is_unframed = False
if frame_size & TBinaryProtocol.VERSION_MASK == TBinaryProtocol.VERSION_1:
self._set_client_type(THeaderClientType.UNFRAMED_BINARY)
is_unframed = True
elif (byte_index(first_word, 0) == TCompactProtocol.PROTOCOL_ID
and byte_index(first_word, 1) & TCompactProtocol.VERSION_MASK
== TCompactProtocol.VERSION):
self._set_client_type(THeaderClientType.UNFRAMED_COMPACT)
is_unframed = True
if is_unframed:
bytes_left_to_read = req_sz - I32.size
if bytes_left_to_read > 0:
rest = self._transport.read(bytes_left_to_read)
else:
rest = b""
self._read_buffer = BufferIO(first_word + rest)
return
# ok, we're still here so we're framed.
if frame_size > self._max_frame_size:
raise TTransportException(
TTransportException.SIZE_LIMIT,
"Frame was too large.",
)
read_buffer = BufferIO(self._transport.readAll(frame_size))
# the next word is either going to be the version field of a
# binary/compact protocol message or the magic value + flags of a
# header protocol message.
second_word = read_buffer.read(I32.size)
version, = I32.unpack(second_word)
read_buffer.seek(0)
if version >> 16 == HEADER_MAGIC:
self._set_client_type(THeaderClientType.HEADERS)
self._read_buffer = self._parse_header_format(read_buffer)
elif version & TBinaryProtocol.VERSION_MASK == TBinaryProtocol.VERSION_1:
self._set_client_type(THeaderClientType.FRAMED_BINARY)
self._read_buffer = read_buffer
elif (byte_index(second_word, 0) == TCompactProtocol.PROTOCOL_ID
and byte_index(second_word, 1) & TCompactProtocol.VERSION_MASK
== TCompactProtocol.VERSION):
self._set_client_type(THeaderClientType.FRAMED_COMPACT)
self._read_buffer = read_buffer
else:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Could not detect client transport type.",
)
def _parse_header_format(self, buffer):
# make BufferIO look like TTransport for varint helpers
buffer_transport = TMemoryBuffer()
buffer_transport._buffer = buffer
buffer.read(2) # discard the magic bytes
self.flags, = U16.unpack(buffer.read(U16.size))
self.sequence_id, = I32.unpack(buffer.read(I32.size))
header_length = U16.unpack(buffer.read(U16.size))[0] * 4
end_of_headers = buffer.tell() + header_length
if end_of_headers > len(buffer.getvalue()):
raise TTransportException(
TTransportException.SIZE_LIMIT,
"Header size is larger than whole frame.",
)
self._protocol_id = readVarint(buffer_transport)
transforms = []
transform_count = readVarint(buffer_transport)
for _ in range(transform_count):
transform_id = readVarint(buffer_transport)
if transform_id not in READ_TRANSFORMS_BY_ID:
raise TApplicationException(
TApplicationException.INVALID_TRANSFORM,
"Unknown transform: %d" % transform_id,
)
transforms.append(transform_id)
transforms.reverse()
headers = {}
while buffer.tell() < end_of_headers:
header_type = readVarint(buffer_transport)
if header_type == TInfoHeaderType.KEY_VALUE:
count = readVarint(buffer_transport)
for _ in range(count):
key = _readString(buffer_transport)
value = _readString(buffer_transport)
headers[key] = value
else:
break # ignore unknown headers
self._read_headers = headers
# skip padding / anything we didn't understand
buffer.seek(end_of_headers)
payload = buffer.read()
for transform_id in transforms:
transform_fn = READ_TRANSFORMS_BY_ID[transform_id]
payload = transform_fn(payload)
return BufferIO(payload)
def write(self, buf):
self._write_buffer.write(buf)
def flush(self):
payload = self._write_buffer.getvalue()
self._write_buffer = BufferIO()
buffer = BufferIO()
if self._client_type == THeaderClientType.HEADERS:
for transform_id in self._write_transforms:
transform_fn = WRITE_TRANSFORMS_BY_ID[transform_id]
payload = transform_fn(payload)
headers = BufferIO()
writeVarint(headers, self._protocol_id)
writeVarint(headers, len(self._write_transforms))
for transform_id in self._write_transforms:
writeVarint(headers, transform_id)
if self._write_headers:
writeVarint(headers, TInfoHeaderType.KEY_VALUE)
writeVarint(headers, len(self._write_headers))
for key, value in self._write_headers.items():
_writeString(headers, key)
_writeString(headers, value)
self._write_headers = {}
padding_needed = (4 - (len(headers.getvalue()) % 4)) % 4
headers.write(b"\x00" * padding_needed)
header_bytes = headers.getvalue()
buffer.write(I32.pack(10 + len(header_bytes) + len(payload)))
buffer.write(U16.pack(HEADER_MAGIC))
buffer.write(U16.pack(self.flags))
buffer.write(I32.pack(self.sequence_id))
buffer.write(U16.pack(len(header_bytes) // 4))
buffer.write(header_bytes)
buffer.write(payload)
elif self._client_type in (THeaderClientType.FRAMED_BINARY,
THeaderClientType.FRAMED_COMPACT):
buffer.write(I32.pack(len(payload)))
buffer.write(payload)
elif self._client_type in (THeaderClientType.UNFRAMED_BINARY,
THeaderClientType.UNFRAMED_COMPACT):
buffer.write(payload)
else:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Unknown client type.",
)
# the frame length field doesn't count towards the frame payload size
frame_bytes = buffer.getvalue()
frame_payload_size = len(frame_bytes) - 4
if frame_payload_size > self._max_frame_size:
raise TTransportException(
TTransportException.SIZE_LIMIT,
"Attempting to send frame that is too large.",
)
self._transport.write(frame_bytes)
self._transport.flush()
@property
def cstringio_buf(self):
return self._read_buffer
def cstringio_refill(self, partialread, reqlen):
result = bytearray(partialread)
while len(result) < reqlen:
result += self.read(reqlen - len(result))
self._read_buffer = BufferIO(result)
return self._read_buffer
class TTransformTransport(TTransport.TTransportBase):
"""
A transport that transforms the written payload and supplies it as read payload.
The transformation can be any kind of Synchronous operation, either a computation
or a result of a network call.
"""
def __init__(self, value=None, offset=0):
"""
:param value: a value to read from for stringio
If value is set, the read buffer will be initialized with it
otherwise, it is for writing
:param offset: the offset to start reading from.
"""
if value is not None:
self.__read_buffer = BufferIO(value)
else:
self.__read_buffer = BufferIO()
if offset:
self.__read_buffer.seek(offset)
self.__write_buffer = BufferIO()
def isOpen(self):
"""
:return: True if the transport is open, False otherwise
"""
return (not self.__read_buffer.closed) and (
not self.__write_buffer.closed)
def close(self):
"""
closes the transport
"""
self.__read_buffer.close()
self.__write_buffer.close()
def read(self, sz):
"""
reads from the transport
:param sz: number of bytes to read
:return: the data read
"""
return self.__read_buffer.read(sz)
def write(self, buf):
"""
writes to the transport
:param buf: the buffer to write
"""
self.__write_buffer.write(buf)
def _transform(self, buf):
"""
Transforms the data written, and sets it as data to be read
:param buf: The data written to the transport
:return: The data to set as readable from the transport
"""
return buf
def flush(self):
"""
flushes the transport (verify all previous writes actually written)
"""
self.__read_buffer = BufferIO(
self._transform(self.__write_buffer.getvalue()))
self.__write_buffer = BufferIO()
def getvalue(self):
"""
:return: all the current data available for read from the transport
"""
return self.__read_buffer.getvalue()