Example #1
0
    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()
Example #2
0
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
Example #3
0
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()