class QsfpServiceClient(QsfpService.Client): DEFAULT_PORT = 5910 DEFAULT_TIMEOUT = 10.0 # we ignore the value of port def __init__(self, host, port=None, timeout=None): # In a box with all 32 QSFP ports populated, it takes about 7.5s right # now to read all 32 QSFP ports. So, put the defaut timeout to 10s. self.host = host timeout = timeout or self.DEFAULT_TIMEOUT self._socket = TSocket(host, self.DEFAULT_PORT) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() QsfpService.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close()
def call_processor(self, input, client_type, protocol_type, client_principal): try: # The input string has already had the header removed, but # the python processor will expect it to be there. In # order to reconstitute the message with headers, we use # the THeaderProtocol object to write into a memory # buffer, then pass that buffer to the python processor. write_buf = TMemoryBuffer() trans = THeaderTransport(write_buf, client_types=[client_type]) trans.set_protocol_id(protocol_type) trans.write(input) trans.flush() prot_buf = TMemoryBuffer(write_buf.getvalue()) prot = THeaderProtocol(prot_buf) ctx = TCppConnectionContext(client_principal) self.processor.process(prot, prot, ctx) # And on the way out, we need to strip off the header, # because the C++ code will expect to add it. read_buf = TMemoryBuffer(prot_buf.getvalue()) trans = THeaderTransport(read_buf, client_types=[client_type]) trans.readFrame(0) return trans.cstringio_buf.read() except: # Don't let exceptions escape back into C++ traceback.print_exc()
def __init__(self, trans, client_types=None, fuzz_fields=None, fuzz_all_if_empty=True, verbose=False): """Create a TFuzzyHeaderTransport instance. @param fuzz_fields(iterable) Collection of string names of fields to fuzz. Each name must be a member of TFuzzyHeaderTransport.fuzzable_fields. If the iterable is empty and fuzz_all_if_empty is True, all fields will be fuzzed. If the iterable is empty and fuzz_all_if_empty is False, no fields will be fuzzed. @param verbose(bool) Whether to print fuzzed fields to stdout @param fuzz_all_if_empty(bool) Whether to fuzz eveyr field if the fuzz_fields iterable is None or empty. trans and client_type are forwarded to THeaderProtocol.__init__ """ cls = self.__class__ THeaderTransport.__init__(self, trans, client_types=client_types) if fuzz_fields is None: fuzz_fields = [] self._fuzz_fields = set() for field_name in fuzz_fields: if field_name not in cls.fuzzable_fields: raise NameError("Invalid fuzz field: %s" % field_name) self._fuzz_fields.add(field_name) if not self._fuzz_fields and fuzz_all_if_empty: # No fields were explicitly included self._fuzz_fields = cls.fuzzable_fields self._verbose = verbose
def setUp(self): self.trans = TTransport.TMemoryBuffer() self.trans = THeaderTransport(self.trans) self.prot = TBinaryProtocol.TBinaryProtocol(self.trans) self.x = Xtruct() self.x.string_thing = "Zero" self.x.byte_thing = 0
def call_processor(self, input, headers, client_type, protocol_type, context_data, callback): try: # The input string has already had the header removed, but # the python processor will expect it to be there. In # order to reconstitute the message with headers, we use # the THeaderProtocol object to write into a memory # buffer, then pass that buffer to the python processor. should_sample = self._shouldSample() timestamps = CallTimestamps() timestamps.processBegin = 0 timestamps.processEnd = 0 if self.observer and should_sample: timestamps.processBegin = int(time.time() * 10**6) write_buf = TMemoryBuffer() trans = THeaderTransport(write_buf) trans._THeaderTransport__client_type = client_type trans._THeaderTransport__write_headers = headers trans.set_protocol_id(protocol_type) trans.write(input) trans.flush() prot_buf = TMemoryBuffer(write_buf.getvalue()) prot = THeaderProtocol(prot_buf, client_types=[client_type]) ctx = TCppConnectionContext(context_data) ret = self.processor.process(prot, prot, ctx) done_callback = partial(_ProcessorAdapter.done, prot_buf=prot_buf, client_type=client_type, callback=callback) if self.observer: if should_sample: timestamps.processEnd = int(time.time() * 10**6) # This only bumps counters if `processBegin != 0` and # `processEnd != 0` and these will only be non-zero if # we are sampling this request. self.observer.callCompleted(timestamps) # This future is created by and returned from the processor's # ThreadPoolExecutor, which keeps a reference to it. So it is # fine for this future to end its lifecycle here. if isinstance(ret, Future): ret.add_done_callback(lambda x, d=done_callback: d()) else: done_callback() except: # Don't let exceptions escape back into C++ traceback.print_exc()
def __init__(self, host, port=None, timeout=2.0): self.host = host self._socket = TSocket(host, self.DEFAULT_PORT) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() QsfpService.Client.__init__(self, self._protocol)
def __init__(self, eden_dir=None, mounted_path=None): self._eden_dir = eden_dir if mounted_path: sock_path = os.path.join(mounted_path, '.eden', 'socket') else: sock_path = os.path.join(self._eden_dir, SOCKET_PATH) self._socket = TSocket(unix_socket=sock_path) self._socket.setTimeout(60000) # in milliseconds self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) super(EdenClient, self).__init__(self._protocol)
def __init__(self, host, port=None, timeout=5.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() PcapPushSubscriber.Client.__init__(self, self._protocol)
def __init__(self, host, port=None, timeout=5.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() NetlinkManagerService.Client.__init__(self, self._protocol)
def __init__(self, host, port=None, timeout=10.0): # In a box with all 32 QSFP ports populated, it takes about 7.5s right # now to read all 32 QSFP ports. So, put the defaut timeout to 10s. self.host = host self._socket = TSocket(host, self.DEFAULT_PORT) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() QsfpService.Client.__init__(self, self._protocol)
class FbossAgentClient(FbossCtrl.Client): DEFAULT_PORT = 5909 def __init__(self, host, port=None, timeout=5.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() FbossCtrl.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close() # # The getPortStats() thrift API was unfortunately renamed to getPortInfo(). # Here's a hacky workaround that tries to do the right thing regardless of # whether the switch we are talking to supports getPortStats() or # getPortInfo(). # def getPortStats(self, *args, **kwargs): return self.getPortInfo(*args, **kwargs) def getAllPortStats(self, *args, **kwargs): return self.getAllPortInfo(*args, **kwargs) def getPortInfo(self, *args, **kwargs): try: return FbossCtrl.Client.getPortInfo(self, *args, **kwargs) except TApplicationException as ex: if 'Method name getPortInfo not found' in str(ex): return FbossCtrl.Client.getPortStats(self, *args, **kwargs) raise def getAllPortInfo(self, *args, **kwargs): try: return FbossCtrl.Client.getAllPortInfo(self, *args, **kwargs) except TApplicationException as ex: if 'Method name getAllPortInfo not found' in str(ex): return FbossCtrl.Client.getAllPortStats(self, *args, **kwargs) raise
def __init__(self, eden_dir=None, mounted_path=None): self._eden_dir = eden_dir if mounted_path: sock_path = os.readlink( os.path.join(mounted_path, '.eden', 'socket')) else: sock_path = os.path.join(self._eden_dir, SOCKET_PATH) self._socket = TSocket(unix_socket=sock_path) # We used to set a timeout here, but picking the right duration is hard, # and safely retrying an arbitrary thrift call may not be safe. So we # just leave the client with no timeout. #self._socket.setTimeout(60000) # in milliseconds self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) super(EdenClient, self).__init__(self._protocol)
def __init__(self, eden_dir=None, socket_path=None): if socket_path is not None: self._socket_path = socket_path elif eden_dir is not None: self._socket_path = os.path.join(eden_dir, SOCKET_PATH) else: raise TypeError("one of eden_dir or socket_path is required") self._socket = TSocket(unix_socket=self._socket_path) # We used to set a timeout here, but picking the right duration is hard, # and safely retrying an arbitrary thrift call may not be safe. So we # just leave the client with no timeout. # self._socket.setTimeout(60000) # in milliseconds self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) super(EdenClient, self).__init__(self._protocol)
def done(prot_buf, client_type, callback): try: response = prot_buf.getvalue() if len(response) == 0: callback.call(response) else: # And on the way out, we need to strip off the header, # because the C++ code will expect to add it. read_buf = TMemoryBuffer(response) trans = THeaderTransport(read_buf, client_types=[client_type]) trans.readFrame(len(response)) callback.call(trans.cstringio_buf.read()) except: traceback.print_exc()
def _get_client_by_transport(self, options, transport, socket=None): # Create the protocol and client if options.json: protocol = TJSONProtocol.TJSONProtocol(transport) elif options.compact: protocol = TCompactProtocol.TCompactProtocol(transport) # No explicit option about protocol is specified. Try to infer. elif options.framed or options.unframed: protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport) elif socket is not None: # If json, compact, framed, and unframed are not specified, # THeaderProtocol is the default. Create a protocol using either # fuzzy or non-fuzzy transport depending on if options.fuzz is set. if options.fuzz is not None: transport = TFuzzyHeaderTransport(socket, fuzz_fields=options.fuzz, verbose=True) else: transport = THeaderTransport(socket) protocol = THeaderProtocol.THeaderProtocol(transport) else: self._exit(error_message=('No valid protocol ' 'specified for %s' % (type(self))), status=os.EX_USAGE) transport.open() self._transport = transport client = self.service_class.Client(protocol) return client
def create_thrift_client( eden_dir: "Optional[str]" = None, socket_path: "Optional[str]" = None, timeout: "Optional[float]" = None, ) -> "EdenClient": """ Construct a thrift client to speak to the running eden server instance associated with the specified mount point. @return Returns a context manager for EdenService.Client. """ if socket_path is not None: pass elif eden_dir is not None: socket_path = os.path.join(eden_dir, SOCKET_PATH) else: raise TypeError("one of eden_dir or socket_path is required") if sys.platform == "win32": socket = WinTSocket(unix_socket=socket_path) else: socket = TSocket(unix_socket=socket_path) # We used to set a default timeout here, but picking the right duration is hard, # and safely retrying an arbitrary thrift call may not be safe. So we # just leave the client with no timeout, unless one is given. if timeout is None: timeout_ms = None else: timeout_ms = timeout * 1000 socket.setTimeout(timeout_ms) transport = THeaderTransport(socket) protocol = THeaderProtocol(transport) return EdenClient(socket_path, transport, protocol)
def __init__(self, trans, strictRead=False, client_types=None, client_type=None): """Create a THeaderProtocol instance @param transport(TTransport) The underlying transport. @param strictRead(bool) Turn on strictRead if using TBinaryProtocol @param client_types([CLIENT_TYPE.HEADER, ...]) List of client types to support. Defaults to CLIENT_TYPE.HEADER only. """ if isinstance(trans, THeaderTransport): trans._THeaderTransport__supported_client_types = set( client_types or (CLIENT_TYPE.HEADER, )) if client_type is not None: trans._THeaderTransport__client_type = client_type htrans = trans else: htrans = THeaderTransport(trans, client_types, client_type) TProtocolBase.__init__(self, htrans) self.strictRead = strictRead self.reset_protocol()
def schedule_timeout(self, fname, seqid): timeout = self.timeouts[fname] if not timeout: return tmo = TMemoryBuffer() thp = THeaderTransport(tmo) oprot = THeaderProtocol(thp) exc = TApplicationException(TApplicationException.TIMEOUT, "Call to {} timed out".format(fname)) oprot.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid) exc.write(oprot) oprot.writeMessageEnd() thp.flush() timeout_task = self.loop.create_task( self.message_received(tmo.getvalue(), delay=timeout), ) self.update_pending_tasks(seqid, timeout_task)
class EdenClient(EdenService.Client): ''' EdenClient is a subclass of EdenService.Client that provides a few additional conveniences: - Smarter constructor - Implement the context manager __enter__ and __exit__ methods, so it can be used in with statements. ''' def __init__(self, eden_dir=None, mounted_path=None): self._eden_dir = eden_dir if mounted_path: sock_path = os.readlink( os.path.join(mounted_path, '.eden', 'socket')) else: sock_path = os.path.join(self._eden_dir, SOCKET_PATH) self._socket = TSocket(unix_socket=sock_path) # We used to set a timeout here, but picking the right duration is hard, # and safely retrying an arbitrary thrift call may not be safe. So we # just leave the client with no timeout. #self._socket.setTimeout(60000) # in milliseconds self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) super(EdenClient, self).__init__(self._protocol) def __enter__(self): self.open() return self def __exit__(self, exc_type, exc_value, exc_traceback): self.close() def open(self): try: self._transport.open() except TTransportException as ex: self.close() if ex.type == TTransportException.NOT_OPEN: raise EdenNotRunningError(self._eden_dir) raise def close(self): if self._transport is not None: self._transport.close() self._transport = None
def __init__(self, transport, allowed_client_types): # much of the actual work for THeaderProtocol happens down in # THeaderTransport since we need to do low-level shenanigans to detect # if the client is sending us headers or one of the headerless formats # we support. this wraps the real transport with the one that does all # the magic. if not isinstance(transport, THeaderTransport): transport = THeaderTransport(transport, allowed_client_types) super(THeaderProtocol, self).__init__(transport) self._set_protocol()
def setUp(self): """Create two buffers, transports, and protocols. self._h_trans uses THeaderTransport self._f_trans uses TFuzzyHeaderTransport """ cls = self.__class__ # THeaderTransport attributes self._h_buf = TMemoryBuffer() self._h_trans = THeaderTransport(self._h_buf) self._h_prot = THeaderProtocol(self._h_trans) # TFuzzyHeaderTransport attributes self._f_buf = TMemoryBuffer() self._f_trans = TFuzzyHeaderTransport( self._f_buf, fuzz_fields=cls.fuzz_fields, fuzz_all_if_empty=False, verbose=False) self._f_prot = THeaderProtocol(self._f_trans)
def schedule_timeout(self, fname, seqid): timeout = self.timeouts[fname] if not timeout: return tmo = TMemoryBuffer() thp = THeaderTransport(tmo) oprot = THeaderProtocol(thp) exc = TApplicationException( TApplicationException.TIMEOUT, "Call to {} timed out".format(fname) ) oprot.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid) exc.write(oprot) oprot.writeMessageEnd() thp.flush() timeout_task = self.loop.create_task( self.message_received(tmo.getvalue(), delay=timeout), ) self.update_pending_tasks(seqid, timeout_task)
def testOversizeFrameRecv(self): """Test that an oversize frame on recv gets a TTransportException""" self.trans.set_max_frame_size(200) self.x.write(self.prot) self.trans.flush() inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue()) inTrans = THeaderTransport(inTrans) inProt = TBinaryProtocol.TBinaryProtocol(inTrans) inTrans.set_max_frame_size(2) try: self.x.read(inProt) except TTransport.TTransportException: return self.fail("Should have gotten TTransportException")
class TestClient(TestService.Client): DEFAULT_PORT = fboss.system_tests.test.constants.DEFAULT_PORT def __init__(self, host, port=None, timeout=10.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() TestService.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close()
class NetlinkManagerClient(NetlinkManagerService.Client): DEFAULT_PORT = 5912 def __init__(self, host, port=None, timeout=5.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() NetlinkManagerService.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close()
class QsfpServiceClient(QsfpService.Client): DEFAULT_PORT = 5910 # we ignore the value of port def __init__(self, host, port=None, timeout=2.0): self.host = host self._socket = TSocket(host, self.DEFAULT_PORT) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() QsfpService.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close()
class PcapPushSubClient(PcapPushSubscriber.Client): DEFAULT_PORT = 5911 def __init__(self, host, port=None, timeout=5.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() PcapPushSubscriber.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close()
class PlainTextFbossAgentClientDontUseInFb(FbossCtrl.Client): DEFAULT_PORT = 5909 def __init__(self, host, port=None, timeout=5.0): self.host = host if port is None: port = self.DEFAULT_PORT self._socket = TSocket(host, port) # TSocket.setTimeout() takes a value in milliseconds self._socket.setTimeout(timeout * 1000) self._transport = THeaderTransport(self._socket) self._protocol = THeaderProtocol(self._transport) self._transport.open() FbossCtrl.Client.__init__(self, self._protocol) def __enter__(self): return self def __exit__(self, type, value, traceback): self._transport.close()
class TestEof(unittest.TestCase): def setUp(self): self.trans = TTransport.TMemoryBuffer() self.trans = THeaderTransport(self.trans) self.prot = TBinaryProtocol.TBinaryProtocol(self.trans) self.x = Xtruct() self.x.string_thing = "Zero" self.x.byte_thing = 0 def testOversizeFrameRecv(self): """Test that an oversize frame on recv gets a TTransportException""" self.trans.set_max_frame_size(200) self.x.write(self.prot) self.trans.flush() inTrans = TTransport.TMemoryBuffer( self.trans.getTransport().getvalue()) inTrans = THeaderTransport(inTrans) inProt = TBinaryProtocol.TBinaryProtocol(inTrans) inTrans.set_max_frame_size(2) try: self.x.read(inProt) except TTransport.TTransportException: return self.fail("Should have gotten TTransportException") def testOversizeFrameSend(self): """Test that an oversize frame on send gets a TTransportException""" self.trans.set_max_frame_size(2) self.x.write(self.prot) try: self.trans.flush() except TTransport.TTransportException: return self.fail("Should have gotten TTransportException")
def serialize_texception(cls, fname, seqid, exception): """This saves us a bit of processing time for timeout handling by reusing the Thrift structs involved in exception serialization. NOTE: this is not thread-safe nor it is meant to be. """ # the serializer is a singleton if cls._exception_serializer is None: buffer = TWriteOnlyBuffer() transport = THeaderTransport(buffer) cls._exception_serializer = THeaderProtocol(transport) else: transport = cls._exception_serializer.trans buffer = transport.getTransport() buffer.reset() serializer = cls._exception_serializer serializer.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid) exception.write(serializer) serializer.writeMessageEnd() serializer.trans.flush() return buffer.getvalue()
def serialize_texception(cls, fname, seqid, exception): """This saves us a bit of processing time for timeout handling by reusing the Thrift structs involved in exception serialization. NOTE: this is not thread-safe nor is it meant to be. """ # the serializer is a singleton if cls._exception_serializer is None: buffer = TWriteOnlyBuffer() transport = THeaderTransport(buffer) cls._exception_serializer = THeaderProtocol(transport) else: transport = cls._exception_serializer.trans buffer = transport.getTransport() buffer.reset() serializer = cls._exception_serializer serializer.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid) exception.write(serializer) serializer.writeMessageEnd() serializer.trans.flush() return buffer.getvalue()
class TestEof(unittest.TestCase): def setUp(self): self.trans = TTransport.TMemoryBuffer() self.trans = THeaderTransport(self.trans) self.prot = TBinaryProtocol.TBinaryProtocol(self.trans) self.x = Xtruct() self.x.string_thing = "Zero" self.x.byte_thing = 0 def testOversizeFrameRecv(self): """Test that an oversize frame on recv gets a TTransportException""" self.trans.set_max_frame_size(200) self.x.write(self.prot) self.trans.flush() inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue()) inTrans = THeaderTransport(inTrans) inProt = TBinaryProtocol.TBinaryProtocol(inTrans) inTrans.set_max_frame_size(2) try: self.x.read(inProt) except TTransport.TTransportException: return self.fail("Should have gotten TTransportException") def testOversizeFrameSend(self): """Test that an oversize frame on send gets a TTransportException""" self.trans.set_max_frame_size(2) self.x.write(self.prot) try: self.trans.flush() except TTransport.TTransportException: return self.fail("Should have gotten TTransportException")
def call_processor(self, input, headers, client_type, protocol_type, context_data, callback): try: # The input string has already had the header removed, but # the python processor will expect it to be there. In # order to reconstitute the message with headers, we use # the THeaderProtocol object to write into a memory # buffer, then pass that buffer to the python processor. write_buf = TMemoryBuffer() trans = THeaderTransport(write_buf) trans._THeaderTransport__client_type = client_type trans._THeaderTransport__write_headers = headers trans.set_protocol_id(protocol_type) trans.write(input) trans.flush() prot_buf = TMemoryBuffer(write_buf.getvalue()) prot = THeaderProtocol(prot_buf, client_types=[client_type]) ctx = TCppConnectionContext(context_data) ret = self.processor.process(prot, prot, ctx) done_callback = partial(_ProcessorAdapter.done, prot_buf=prot_buf, client_type=client_type, callback=callback) # This future is created by and returned from the processor's # ThreadPoolExecutor, which keeps a reference to it. So it is # fine for this future to end its lifecycle here. if isinstance(ret, Future): ret.add_done_callback(lambda x, d=done_callback: d()) else: done_callback() except: # Don't let exceptions escape back into C++ traceback.print_exc()
def __init__(self, trans, strictRead=False, client_types=None): """Create a THeaderProtocol instance @param transport(TTransport) The underlying transport. @param strictRead(bool) Turn on strictRead if using TBinaryProtocol @param client_types([THeaderTransport.HEADERS_CLIENT_TYPE, ...]) List of client types to support. Defaults to HEADERS_CLIENT_TYPE only. """ htrans = THeaderTransport(trans, client_types) TProtocolBase.__init__(self, htrans) self.strictRead = strictRead self.reset_protocol()
def call_processor(self, input, client_type, protocol_type, context_data): try: # The input string has already had the header removed, but # the python processor will expect it to be there. In # order to reconstitute the message with headers, we use # the THeaderProtocol object to write into a memory # buffer, then pass that buffer to the python processor. write_buf = TMemoryBuffer() trans = THeaderTransport(write_buf) trans._THeaderTransport__client_type = client_type trans.set_protocol_id(protocol_type) trans.write(input) trans.flush() prot_buf = TMemoryBuffer(write_buf.getvalue()) prot = THeaderProtocol(prot_buf, client_types=[client_type]) ctx = TCppConnectionContext(context_data) self.processor.process(prot, prot, ctx) # Check for empty result. If so, return an empty string # here. This is probably a oneway request, but we can't # reliably tell. The C++ code does basically the same # thing. response = prot_buf.getvalue() if len(response) == 0: return response # And on the way out, we need to strip off the header, # because the C++ code will expect to add it. read_buf = TMemoryBuffer(response) trans = THeaderTransport(read_buf, client_types=[client_type]) trans.readFrame(len(response)) return trans.cstringio_buf.read() except: # Don't let exceptions escape back into C++ traceback.print_exc()
class FuzzyTransportTest(object): """Test class that sets up a THeaderTransport and a TFuzzyHeaderTransport. Used for writing and comparing messages using both transports. """ fuzz_fields = [] # Sample TestService method arguments sampleListStruct = ttypes.ListStruct( a=[True, False], b=[1, 2, 3], c=[1.2, 3.4], d=["ab", "cd"], e=[list(sm.xrange(n)) for n in sm.xrange(20)], f=[{1: 2}, {3: 4, 5: 6}], g=[{"a", "b"}, {"c"}, set()] ) def setUp(self): """Create two buffers, transports, and protocols. self._h_trans uses THeaderTransport self._f_trans uses TFuzzyHeaderTransport """ cls = self.__class__ # THeaderTransport attributes self._h_buf = TMemoryBuffer() self._h_trans = THeaderTransport(self._h_buf) self._h_prot = THeaderProtocol(self._h_trans) # TFuzzyHeaderTransport attributes self._f_buf = TMemoryBuffer() self._f_trans = TFuzzyHeaderTransport( self._f_buf, fuzz_fields=cls.fuzz_fields, fuzz_all_if_empty=False, verbose=False) self._f_prot = THeaderProtocol(self._f_trans) def writeObject(self, obj=sampleListStruct): """Write an object to the test and reference protocols. Return the contents of both buffers. """ obj.write(self._h_prot) obj.write(self._f_prot) self._h_trans.flush() self._f_trans.flush() h_buf = self._h_buf.getvalue() f_buf = self._f_buf.getvalue() return h_buf, f_buf def differentIndices(self, header, fuzzy): """Return a list of byte positions at which two messages' bytes differ. Header should be the contents of self._h_buf Fuzzy should be the contents of self._f_buf """ indices = [] for i, (h, f) in enumerate(itertools.izip(header, fuzzy)): if h != f: indices.append(i) return indices def assertEqualsExceptIndices(self, header, fuzzy, indices): """Assert that the buffers `header` and `fuzzy` are equal, except possibly at the byte positions included in `indices`. This ensures that the message produced by TFuzzyHeaderProtocol (fuzzy) is equal to the message produced by THeaderProtocol (header), except at the byte positions that are expected to be fuzzed.""" self.assertEquals(len(header), len(fuzzy)) for diff in self.differentIndices(header, fuzzy): self.assertIn(diff, indices)
def round_robin(self, compress=None): original = b'A' * MAX_FRAME_SIZE mb = TMemoryBuffer() trans = THeaderTransport(mb, client_type=CLIENT_TYPE.HEADER) trans.set_max_frame_size(MAX_FRAME_SIZE + MIN_HEADER_SIZE) if compress: trans.add_transform(compress) trans.write(original) trans.flush() frame = mb.getvalue() # Cleanup the memory buffer mb.close() del mb if compress is None: # Partial Decode the frame and see if its correct size wise sz = struct.unpack('!I', frame[:4])[0] self.assertEqual(sz, BIG_FRAME_MAGIC) sz = struct.unpack('!Q', frame[4:12])[0] self.assertEqual(len(frame), sz + 12) # Read it back mb = TMemoryBuffer(frame) trans = THeaderTransport(mb, client_type=CLIENT_TYPE.HEADER) trans.set_max_frame_size(len(frame)) trans.readFrame(0) result = trans.read(MAX_FRAME_SIZE) mb.close() del mb self.assertEqual(result, original, 'round-robin different from original')