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 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 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 _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 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 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 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 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): 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 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 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 connection_lost(self, exc): """Implements asyncio.Protocol.connection_lost.""" te = TTransportException(type=TTransportException.END_OF_FILE, message="Connection closed") self.fail_all_futures(te)
def ZstdDecompressor(self): raise TTransportException(TTransportException.INVALID_TRANSFORM, 'zstd module not available')
def ZstdCompressor(self, write_content_size): raise TTransportException(TTransportException.INVALID_TRANSFORM, 'zstd module not available')
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 decompress(self, buf): raise TTransportException(TTransportException.INVALID_TRANSFORM, 'snappy module not available')
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 getPeerName(self): if not self.handle: raise TTransportException(TTransportException.NOT_OPEN, 'Transport not open') return self.handle.getpeername()
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 getSocketName(self): if not self.handles: raise TTransportException(TTransportException.NOT_OPEN, 'Transport not open') return next(iter(self.handles.values())).getsockname()