def read_header_format(self, sz, header_size, data):
# clear out any previous transforms
self.__read_transforms = []
header_size = header_size * 4
if header_size > sz:
raise TTransportException(TTransportException.INVALID_FRAME_SIZE,
"Header size is larger than frame")
end_header = header_size + data.tell()
self.__proto_id = readVarint(data)
num_headers = readVarint(data)
if self.__proto_id == 1 and self.__client_type != \
CLIENT_TYPE.HTTP_SERVER:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Trying to recv JSON encoding over binary")
# Read the headers. Data for each header varies.
for _ in range(0, num_headers):
trans_id = readVarint(data)
if trans_id in (TRANSFORM.ZLIB, TRANSFORM.SNAPPY, TRANSFORM.ZSTD):
self.__read_transforms.insert(0, trans_id)
elif trans_id == TRANSFORM.HMAC:
raise TApplicationException(
TApplicationException.INVALID_TRANSFORM,
"Hmac transform is no longer supported: %i" % trans_id)
else:
# TApplicationException will be sent back to client
raise TApplicationException(
TApplicationException.INVALID_TRANSFORM,
"Unknown transform in client request: %i" % trans_id)
# Clear out previous info headers.
self.__read_headers.clear()
# Read the info headers.
while data.tell() < end_header:
info_id = readVarint(data)
if info_id == INFO.NORMAL:
_read_info_headers(data, end_header, self.__read_headers)
elif info_id == INFO.PERSISTENT:
_read_info_headers(data, end_header,
self.__read_persistent_headers)
else:
break # Unknown header. Stop info processing.
if self.__read_persistent_headers:
self.__read_headers.update(self.__read_persistent_headers)
# Skip the rest of the header
data.seek(end_header)
payload = data.read(sz - header_size)
# Read the data section.
self.__rbuf = StringIO(self.untransform(payload))
def read(self, sz):
try:
buff = self.handle.recv(sz)
if len(buff) == 0:
raise TTransportException(type=TTransportException.END_OF_FILE,
message='TSocket read 0 bytes')
except socket.error as e:
raise TTransportException(type=TTransportException.END_OF_FILE,
message='Socket read failed: {}'.format(
str(e)))
return buff
def set_max_frame_size(self, size):
if size > MAX_BIG_FRAME_SIZE:
raise TTransportException(
TTransportException.INVALID_FRAME_SIZE,
"Cannot set max frame size > %s" % MAX_BIG_FRAME_SIZE)
if size > MAX_FRAME_SIZE and self.__client_type != CLIENT_TYPE.HEADER:
raise TTransportException(
TTransportException.INVALID_FRAME_SIZE,
"Cannot set max frame size > %s for clients other than HEADER"
% MAX_FRAME_SIZE)
self.__max_frame_size = size
def open(self):
address = None
try:
res0 = self._resolveAddr(self.family)
for res in res0:
address = res[4]
handle = socket.socket(res[0], res[1])
self.setHandle(handle)
handle.settimeout(self._timeout)
self.setCloseOnExec(self.close_on_exec)
try:
handle.connect(address)
except socket.error:
self.close()
if res is not res0[-1]:
continue
else:
raise
break
except socket.error as e:
if self._unix_socket:
msg = 'socket error connecting to path %s: %s' % (
self._unix_socket, repr(e))
else:
msg = 'socket error connecting to host %s, port %s (%s): %s' % (
self.host, self.port, repr(address), repr(e))
raise TTransportException(TTransportException.NOT_OPEN, msg)
def write(self, buff):
if not self.handle:
raise TTransportException(TTransportException.NOT_OPEN,
'Transport not open')
sent = 0
have = len(buff)
while sent < have:
try:
plus = self.handle.send(buff)
except socket.error as e:
raise TTransportException(
type=TTransportException.END_OF_FILE,
message='Socket write failed: {}'.format(str(e)))
assert plus > 0
sent += plus
buff = buff[plus:]
def open(self):
TSocket.open(self)
self.write("CONNECT %s:%d HTTP/1.1\r\nHost: %s:%d\r\n\r\n" % (
self.remote_host, self.remote_port,
self.remote_host, self.remote_port))
res = self.read(4096)
try:
status = res.split()[1]
if status != '200':
self.close()
raise TTransportException(TTransportException.NOT_OPEN,
"Error response from proxy server: %s" % res)
except IndexError:
self.close()
raise TTransportException(TTransportException.NOT_OPEN,
"Error response from proxy server: %s" % res)
def _handle_message(self, frame, clear_timeout):
try:
tmi = TReadOnlyBuffer(frame)
iprot = self.THEADER_PROTOCOL_FACTORY(
client_type=self.client_type, ).getProtocol(tmi)
(fname, mtype, seqid) = iprot.readMessageBegin()
except TTransportException as ex:
self.fail_all_futures(ex)
self.transport.close()
return
except Exception as ex:
te = TTransportException(type=TTransportException.END_OF_FILE,
message=str(ex))
self.fail_all_futures(te)
self.transport.close()
return
if clear_timeout:
try:
timeout_task = self.pending_tasks.pop(seqid)
except KeyError:
# Task doesn't have a timeout or has already been cancelled
# and pruned from `pending_tasks`.
pass
else:
timeout_task.cancel()
self._handle_message_received(iprot, fname, mtype, seqid)
def fileno(self):
if not self.handles:
raise TTransportException(TTransportException.NOT_OPEN,
'Transport not open')
if sys.version_info[0] >= 3:
return list(self.handles.values())[0].fileno()
else:
return self.handles.values()[0].fileno()
def flushImpl(self, oneway):
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 != CLIENT_TYPE.HTTP_SERVER:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Trying to send JSON encoding over binary")
buf = StringIO()
if self.__client_type == CLIENT_TYPE.HEADER:
self._flushHeaderMessage(buf, wout, wsz)
elif self.__client_type in (CLIENT_TYPE.FRAMED_DEPRECATED,
CLIENT_TYPE.FRAMED_COMPACT):
buf.write(pack("!i", wsz))
buf.write(wout)
elif self.__client_type in (CLIENT_TYPE.UNFRAMED_DEPRECATED,
CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED):
buf.write(wout)
elif self.__client_type == CLIENT_TYPE.HTTP_SERVER:
# 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 == CLIENT_TYPE.HEADER
elif self.__client_type == CLIENT_TYPE.UNKNOWN:
raise TTransportException(TTransportException.INVALID_CLIENT_TYPE,
"Unknown client type")
# We don't include the framing bytes as part of the frame size check
frame_size = buf.tell() - (4 if wsz < MAX_FRAME_SIZE else 12)
_frame_size_check(frame_size,
self.__max_frame_size,
header=self.__client_type == CLIENT_TYPE.HEADER)
self.getTransport().write(buf.getvalue())
if oneway:
self.getTransport().onewayFlush()
else:
self.getTransport().flush()
def open(self):
TSocketOverHttpTunnel.open(self)
try:
sslh = ssl.SSLSocket(self.handle,
ssl_version=self.ssl_version,
cert_reqs=self.cert_reqs,
keyfile=self.keyfile,
certfile=self.certfile,
ca_certs=self.ca_certs)
self.handle = sslh
except ssl.SSLError as e:
self.close()
raise TTransportException(TTransportException.NOT_OPEN,
"SSL error during handshake: " + str(e))
except socket.error as e:
self.close()
raise TTransportException(
TTransportException.NOT_OPEN,
"socket error during SSL handshake: " + str(e))
def transform(self, buf):
for trans_id in self.__write_transforms:
if trans_id == TRANSFORM.ZLIB:
buf = zlib.compress(buf)
elif trans_id == TRANSFORM.SNAPPY:
buf = snappy.compress(buf)
elif trans_id == TRANSFORM.ZSTD:
buf = zstd.ZstdCompressor(
write_content_size=True).compress(buf)
else:
raise TTransportException(
TTransportException.INVALID_TRANSFORM,
"Unknown transform during send")
return buf
def __init__(self, host, port, proxy_host, proxy_port):
TSocket.__init__(self, proxy_host, proxy_port)
try:
# Use IP address since sometimes proxy_host cannot resolve
# external hostnames using unbound
info = socket.getaddrinfo(
host,
None,
socket.AF_INET | socket.AF_INET6,
socket.SOCK_STREAM,
socket.IPPROTO_TCP)
self.remote_host = info[0][4][0]
except socket.error as e:
raise TTransportException(TTransportException.NOT_OPEN, str(e))
self.remote_port = port
def _sock_accept(self):
if self._queue:
return self._queue.pop()
if hasattr(select, "epoll"):
poller = ConnectionEpoll()
else:
poller = ConnectionSelect()
for filenos in self.handles.keys():
poller.read(filenos)
r, _, x = poller.process(0)
for fd in r:
self._queue.append(self.handles[fd].accept())
if not self._queue:
raise TTransportException("Accept interrupt without client?")
return self._queue.pop()
def listen(self):
res0 = self._resolveAddr(self.family)
for res in res0:
if res[0] == socket.AF_INET6 and res[4][0] == socket.AF_INET6:
# This happens if your version of python was built without IPv6
# support. getaddrinfo() will return IPv6 addresses, but the
# contents of the address field are bogus.
# (For example, see http://bugs.python.org/issue8858)
#
# Ignore IPv6 addresses if python doesn't have IPv6 support.
continue
# We need remove the old unix socket if the file exists and
# nobody is listening on it.
if self._unix_socket:
self._cleanup_unix_socket(res)
# Don't complain if we can't create a socket
# since this is handled below.
try:
handle = socket.socket(res[0], res[1])
except Exception:
continue
handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self._setHandleCloseOnExec(handle)
# Always set IPV6_V6ONLY for IPv6 sockets when not on Windows
if res[0] == socket.AF_INET6 and sys.platform != 'win32':
handle.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY,
True)
handle.settimeout(None)
handle.bind(res[4])
handle.listen(self.tcp_backlog)
self.handles[handle.fileno()] = handle
if not self.handles:
raise TTransportException("No valid interfaces to listen on!")
def ZstdCompressor(self, write_content_size):
raise TTransportException(TTransportException.INVALID_TRANSFORM,
'zstd module not available')
def readFrame(self, req_sz):
self.__rbuf_frame = True
word1 = self.getTransport().readAll(4)
sz = unpack('!I', word1)[0]
proto_id = word1[0] if PY3 else ord(word1[0])
if proto_id == TBinaryProtocol.PROTOCOL_ID:
# unframed
self.__client_type = CLIENT_TYPE.UNFRAMED_DEPRECATED
self.__proto_id = T_BINARY_PROTOCOL
if req_sz <= 4: # check for reads < 0.
self.__rbuf = StringIO(word1)
else:
self.__rbuf = StringIO(word1 +
self.getTransport().read(req_sz - 4))
elif proto_id == TCompactProtocol.PROTOCOL_ID:
self.__client_type = CLIENT_TYPE.UNFRAMED_COMPACT_DEPRECATED
self.__proto_id = T_COMPACT_PROTOCOL
if req_sz <= 4: # check for reads < 0.
self.__rbuf = StringIO(word1)
else:
self.__rbuf = StringIO(word1 +
self.getTransport().read(req_sz - 4))
elif sz == HTTP_SERVER_MAGIC:
self.__client_type = CLIENT_TYPE.HTTP_SERVER
mf = self.getTransport().handle.makefile('rb', -1)
self.handler = RequestHandler(mf, 'client_address:port', '')
self.header = self.handler.wfile
self.__rbuf = StringIO(self.handler.data)
else:
if sz == BIG_FRAME_MAGIC:
sz = unpack('!Q', self.getTransport().readAll(8))[0]
# could be header format or framed. Check next two bytes.
magic = self.getTransport().readAll(2)
proto_id = magic[0] if PY3 else ord(magic[0])
if proto_id == TCompactProtocol.PROTOCOL_ID:
self.__client_type = CLIENT_TYPE.FRAMED_COMPACT
self.__proto_id = T_COMPACT_PROTOCOL
_frame_size_check(sz, self.__max_frame_size, header=False)
self.__rbuf = StringIO(magic +
self.getTransport().readAll(sz - 2))
elif proto_id == TBinaryProtocol.PROTOCOL_ID:
self.__client_type = CLIENT_TYPE.FRAMED_DEPRECATED
self.__proto_id = T_BINARY_PROTOCOL
_frame_size_check(sz, self.__max_frame_size, header=False)
self.__rbuf = StringIO(magic +
self.getTransport().readAll(sz - 2))
elif magic == PACKED_HEADER_MAGIC:
self.__client_type = CLIENT_TYPE.HEADER
_frame_size_check(sz, self.__max_frame_size)
# flags(2), seq_id(4), header_size(2)
n_header_meta = self.getTransport().readAll(8)
self.__flags, self.seq_id, header_size = unpack(
'!HIH', n_header_meta)
data = StringIO()
data.write(magic)
data.write(n_header_meta)
data.write(self.getTransport().readAll(sz - 10))
data.seek(10)
self.read_header_format(sz - 10, header_size, data)
else:
self.__client_type = CLIENT_TYPE.UNKNOWN
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Could not detect client transport type")
if self.__client_type not in self.__supported_client_types:
raise TTransportException(
TTransportException.INVALID_CLIENT_TYPE,
"Client type {} not supported on server".format(
self.__client_type))
def _read_string(bufio, buflimit):
str_sz = readVarint(bufio)
if str_sz + bufio.tell() > buflimit:
raise TTransportException(TTransportException.INVALID_FRAME_SIZE,
"String read too big")
return bufio.read(str_sz)
def decompress(self, buf):
raise TTransportException(TTransportException.INVALID_TRANSFORM,
'snappy module not available')
def getPeerName(self):
if not self.handle:
raise TTransportException(TTransportException.NOT_OPEN,
'Transport not open')
return self.handle.getpeername()
def _frame_size_check(sz, set_max_size, header=True):
if sz > set_max_size or (not header and sz > MAX_FRAME_SIZE):
raise TTransportException(
TTransportException.INVALID_FRAME_SIZE,
"%s transport frame was too large" %
'Header' if header else 'Framed')
def getSocketName(self):
if not self.handles:
raise TTransportException(TTransportException.NOT_OPEN,
'Transport not open')
return next(iter(self.handles.values())).getsockname()
def ZstdDecompressor(self):
raise TTransportException(TTransportException.INVALID_TRANSFORM,
'zstd module not available')
def connection_lost(self, exc):
"""Implements asyncio.Protocol.connection_lost."""
te = TTransportException(type=TTransportException.END_OF_FILE,
message="Connection closed")
self.fail_all_futures(te)