def from_bytes(klass, data): """Returns thrift object from a string, using standard binary representation.""" obj = klass() b = TMemoryBuffer(data) p = TBinaryProtocol(b) obj.read(p) return obj
def make_broad_cast_node_info(node_id, name): m = TMemoryBuffer() p = TBinaryProtocol(m) msg_name = SERVICE_BROADCAST_REQ nonce = get_random_id() head = msg_header(get_magic(), msg_name, nonce) head.write(p) node_info = node_service_info() node_info.name = name node_info.time_stamp = int(time.time()) service_list = [] service_list.append("ai_training") node_info.service_list = service_list kvs = {} kvs["gpu"] = "1 * GeForce940MX" kvs["state"] = "idle" node_info.kvs = kvs node_map = {} node_map[node_id] = node_info req = service_broadcast_req_body(node_map) req.write(p) p.writeMessageEnd() m.flush() return pack_head(m)
def test_with_headers(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocol(client_memory_trans) client_header_trans = client_prot.trans client_header_trans.set_header("Trace", "1234") client_header_trans.set_header("Parent", "2345") client_header_trans.set_header("Span", "3456") client = BaseplateService.Client(client_prot) try: client.is_healthy() except: pass # we don't have a test response for the client self.itrans._readBuffer = StringIO(client_memory_trans.getvalue()) self.processor.process(self.iprot, self.oprot, self.server_context) self.assertEqual(self.observer.on_server_span_created.call_count, 1) context, server_span = self.observer.on_server_span_created.call_args[0] self.assertEqual(server_span.trace_id, 1234) self.assertEqual(server_span.parent_id, 2345) self.assertEqual(server_span.id, 3456) self.assertTrue(self.server_observer.on_start.called) self.assertTrue(self.server_observer.on_finish.called)
def test_primitive_serialization(self): val = int(12) buf = TMemoryBuffer() proto = TCompactProtocol(buf) proto.writeI32(val) reader = TCompactProtocol(StringIO(buf.getvalue())) self.assertEqual(reader.readI32(), val)
async def _on_message_callback(self, message): """The function to be executed when a message is received.""" if not message.reply: logger.warn('no reply present, discarding message') return frame_size = struct.unpack('!I', message.data[:4])[0] if frame_size > _NATS_MAX_MESSAGE_SIZE - 4: logger.warning('frame size too large, dropping message') return # process frame, first four bytes are the frame size iprot = self._protocol_factory.get_protocol( TMemoryBuffer(message.data[4:])) otrans = TMemoryOutputBuffer(_NATS_MAX_MESSAGE_SIZE) oprot = self._protocol_factory.get_protocol(otrans) try: await self._processor.process(iprot, oprot) except TApplicationException: # Continue so the exception is sent to the client pass except Exception: return if len(otrans) == 4: return await self._nats_client.publish(message.reply, otrans.getvalue())
async def request(self, context: FContext, payload) -> TTransportBase: """ Write the current buffer payload over the network and return the response. """ self._preflight_request_check(payload) encoded = base64.b64encode(payload) status, text = await self._make_request(context, encoded) if status == 413: raise TTransportException( type=TTransportExceptionType.RESPONSE_TOO_LARGE, message='response was too large for the transport') if status >= 300: raise TTransportException( type=TTransportExceptionType.UNKNOWN, message='request errored with code {0} and message {1}'.format( status, str(text))) decoded = base64.b64decode(text) if len(decoded) < 4: raise TTransportException(type=TTransportExceptionType.UNKNOWN, message='invalid frame size') if len(decoded) == 4: if any(decoded): raise TTransportException(type=TTransportExceptionType.UNKNOWN, message='missing data') # One-way method, drop response return return TMemoryBuffer(decoded[4:])
def request(self, context, payload): self._preflight_request_check(payload) op_id = str(context._get_op_id()) future = gen.Future() with (yield self._futures_lock.acquire()): if op_id in self._futures: raise TTransportException( type=TTransportExceptionType.UNKNOWN, message="request already in flight for context") self._futures[op_id] = future try: @gen.coroutine def flush_and_wait(): yield self.flush(payload) data = yield future raise gen.Return(data) data = yield gen.with_timeout( timedelta(milliseconds=context.timeout), flush_and_wait()) raise gen.Return(TMemoryBuffer(data)) except gen.TimeoutError: raise TTransportException(type=TTransportExceptionType.TIMED_OUT, message='request timed out') finally: with (yield self._futures_lock.acquire()): del self._futures[op_id]
def get_doc_hot_by_thread(docid_list): docid_hot_list = [] r = redis.Redis(host='***', port=1680, password='******', charset='gbk') redis_result_list = [] try: redis_result_list = r.mget(docid_list) if redis_result_list is not None and len(redis_result_list) == len( docid_list): for i in range(len(redis_result_list)): if redis_result_list[i] is None or len( redis_result_list[i]) == 0: continue docid = docid_list[i] tMemory_o = TMemoryBuffer(redis_result_list[i]) tBinaryProtocol_o = TBinaryProtocol(tMemory_o) hot_feature = HotFeature() hot_feature.read(tBinaryProtocol_o) pv = hot_feature.app_show_num if hot_feature.app_show_num is not None else 0 cl = hot_feature.app_read_num if hot_feature.app_read_num is not None else 0 readtime = hot_feature.app_read_duration_double if hot_feature.app_read_duration_double is not None else 0 coldstart_pv = hot_feature.coldstart_show_num if hot_feature.coldstart_show_num is not None else 0 coldstart_cl = hot_feature.coldstart_read_num if hot_feature.coldstart_read_num is not None else 0 docid_hot_list.append( (docid, pv, cl, readtime, coldstart_pv, coldstart_cl)) except: traceback.print_exc() return docid_hot_list
def readStructFromJSON(jstr, struct_type): stuff = struct_type() trans = TMemoryBuffer(jstr) proto = TSimpleJSONProtocol.TSimpleJSONProtocol(trans, struct_type.thrift_spec) stuff.read(proto) return stuff
def make_get_peer_nodes_resp(dport): m = TMemoryBuffer() p = TBinaryProtocol(m) msg_name = P2P_GET_PEER_NODES_RESP nonce = get_random_id() head = msg_header(get_magic(), msg_name, nonce) head.write(p) print("nonce:%s" % (nonce)) # self.peer_node_id = peer_node_id # self.core_version = core_version # self.protocol_version = protocol_version # self.live_time_stamp = live_time_stamp # self.addr = addr node_list = [] node_info = peer_node_info() node_info.addr = network_address("10.10.254.198", dport) node_info.peer_node_id = gen_node_id() node_info.core_version = core_version node_info.protocol_version = pro_version node_info.live_time_stamp = int(time.time()) node_list.append(node_info) req = get_peer_nodes_resp_body(node_list) req.write(p) p.writeMessageEnd() m.flush() return pack_head(m)
def make_stop_training_req(in_task_id): m = TMemoryBuffer() p = TBinaryProtocol(m) msg_name = STOP_TRAINING_REQ nonce = get_random_id() session_id = get_random_id() head = msg_header(get_magic(), msg_name, nonce, session_id) print("stop req.nonce:%s, task_id:%s" %(nonce, in_task_id)) req=stop_training_req_body(in_task_id) message = in_task_id + nonce sign_algo = "ecdsa" origin = get_node_id() exten_info = {} exten_info["origin_id"] = origin exten_info["sign_algo"] = sign_algo exten_info["sign"] = dbc_sign(message) head.exten_info = exten_info head.write(p) req.write(p) p.writeMessageEnd() m.flush() return pack_head(m)
def write(self, buf): """Bounded write to buffer""" if len(self) + len(buf) > self._limit > 0: self._buffer = TMemoryBuffer() raise TTransportException( type=TTransportExceptionType.REQUEST_TOO_LARGE, message="Buffer size reached {}".format(self._limit)) self._buffer.write(buf)
def to_LinkAttr(self, body): link_info = LinkAttr() try: tMemory_o = TMemoryBuffer(body) tBinaryProtocol_o = TBinaryProtocol(tMemory_o) link_info.read(tBinaryProtocol_o) except: print traceback.format_exc()
def test_unexpected_exception_passed_to_server_span_finish(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans) client = TestService.Client(client_prot) try: client.example_throws(crash=True) except (TTransportException, EOFError): pass # we don't have a test response for the client itrans = TMemoryBuffer(client_memory_trans.getvalue()) iprot = THeaderProtocolFactory().getProtocol(itrans) self.processor.process(iprot, self.oprot) self.assertEqual(self.server_observer.on_start.call_count, 1) self.assertEqual(self.server_observer.on_finish.call_count, 1) _, captured_exc, _ = self.server_observer.on_finish.call_args[0][0] self.assertIsInstance(captured_exc, UnexpectedException)
def test_tree(): t = branch(branch( leaf("aa"), leaf("ab"), ), branch( leaf("ba"), leaf("bb"), )) transport_out = TMemoryBuffer() protocol_out = TJSONProtocol(transport_out) t.write(protocol_out) json_out = transport_out.getvalue() transport_in = TMemoryBuffer(json_out) protocol_in = TJSONProtocol(transport_in) u = Tree() u.read(protocol_in) assert t == u
def service_core_json(request): try: iprot = TJSONProtocol.TJSONProtocol(TMemoryBuffer(request.body)) oprot = TJSONProtocol.TJSONProtocol(TMemoryBuffer()) core_processor.process(iprot, oprot) resp = HttpResponse(oprot.trans.getvalue(), content_type="application/json") try: resp['Access-Control-Allow-Origin'] = request.META['HTTP_ORIGIN'] except: pass return resp except Exception, e: return HttpResponse(json.dumps({'error': str(e)}), content_type="application/json")
def test_client_proxy_flow(self): client_memory_trans = TMemoryBuffer() client_prot = THeaderProtocolFactory().getProtocol(client_memory_trans) class Pool(object): @contextlib.contextmanager def connection(self): yield client_prot client_factory = ThriftContextFactory(Pool(), TestService.Client) span = mock.MagicMock() child_span = span.make_child().__enter__() child_span.trace_id = 1 child_span.parent_id = 1 child_span.id = 1 child_span.sampled = True child_span.flags = None edge_context = self.edge_context_factory.from_upstream( SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH) edge_context.attach_context(child_span.context) client = client_factory.make_object_for_context("test", span) try: client.example_simple() except (TTransportException, EOFError): pass # we don't have a test response for the client itrans = TMemoryBuffer(client_memory_trans.getvalue()) iprot = THeaderProtocolFactory().getProtocol(itrans) self.processor.process(iprot, self.oprot) context, _ = self.observer.on_server_span_created.call_args[0] try: self.assertEqual(context.request_context.user.id, "t2_example") self.assertEqual(context.request_context.user.roles, set()) self.assertEqual(context.request_context.user.is_logged_in, True) self.assertEqual(context.request_context.user.loid, "t2_deadbeef") self.assertEqual(context.request_context.user.cookie_created_ms, 100000) self.assertEqual(context.request_context.oauth_client.id, None) self.assertFalse( context.request_context.oauth_client.is_type("third_party")) self.assertEqual(context.request_context.session.id, "beefdead") except jwt.exceptions.InvalidAlgorithmError: raise unittest.SkipTest("cryptography is not installed")
def setUp(self): self.itrans = TMemoryBuffer() self.iprot = THeaderProtocol(self.itrans) self.otrans = TMemoryBuffer() self.oprot = THeaderProtocol(self.otrans) self.observer = mock.Mock(spec=BaseplateObserver) self.server_observer = mock.Mock(spec=ServerSpanObserver) def _register_mock(context, server_span): server_span.register(self.server_observer) self.observer.on_server_span_created.side_effect = _register_mock self.logger = mock.Mock(spec=logging.Logger) self.server_context = TRpcConnectionContext(self.itrans, self.iprot, self.oprot) mock_filewatcher = mock.Mock(spec=FileWatcher) mock_filewatcher.get_data.return_value = { "secrets": { "jwt/authentication/secret": { "type": "simple", "value": self.TOKEN_SECRET, }, }, "vault": { "token": "test", "url": "http://vault.example.com:8200/", } } secrets = store.SecretsStore("/secrets") secrets._filewatcher = mock_filewatcher baseplate = Baseplate() baseplate.register(self.observer) event_handler = BaseplateProcessorEventHandler( self.logger, baseplate, auth_factory=AuthenticationContextFactory(secrets)) handler = TestHandler() self.processor = TestService.ContextProcessor(handler) self.processor.setEventHandler(event_handler)
def to_string(self, download_req_info): str_req = None try: tMemory_b = TMemoryBuffer() tBinaryProtocol_b = TBinaryProtocol(tMemory_b) download_req_info.write(tBinaryProtocol_b) str_req = tMemory_b.getvalue() except EOFError, e: self.log.warning("cann't write DownLoadReq to string")
def to_string(self, link_info): str_entity = None try: tMemory_b = TMemoryBuffer() tBinaryProtocol_b = TBinaryProtocol.TBinaryProtocol(tMemory_b) link_info.write(tBinaryProtocol_b) str_entity = tMemory_b.getvalue() except EOFError, e: self.log.warning("can't write LinkAttr to string")
def to_string(page_info): str_page_info = None try: tMemory_b = TMemoryBuffer() tBinaryProtocol_b = TBinaryServerProtocol(tMemory_b) page_info.write(tBinaryProtocol_b) str_page_info = tMemory_b.getvalue() except EOFError, e: pass
def write_communication_to_buffer(comm): ''' Serialize communication to buffer (binary string) and return buffer. ''' transport = TMemoryBuffer() protocol = factory.createProtocol(transport) comm.write(protocol) return transport.getvalue()
def to_string(log, page_info): str_page_info = None try: tMemory_b = TMemoryBuffer() tBinaryProtocol_b = TBinaryServerProtocol(tMemory_b) page_info.write(tBinaryProtocol_b) str_page_info = tMemory_b.getvalue() except EOFError, e: log.warning("cann't write DownLoadRsp to string")
def flush(self): msg = self._writeBuffer.getvalue() tmi = TMemoryBuffer(msg) iprot = THeaderProtocol(tmi) fname, mtype, seqid = iprot.readMessageBegin() fname = fname.decode() self._proto.schedule_timeout(fname, seqid) self._trans.send_message(msg) self._writeBuffer = BytesIO()
def call_processor(self, input, headers, 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._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) 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()
def test_forward_compatibility_nested_nested(self): obj = OldStructureNested() obj.features = [{}] obj.features[0][1] = 314 obj.features[0][2] = 271 objN = OldStructureNestedNested() objN.field = obj trans = TMemoryBuffer() proto = self.createProto(trans) objN.write(proto) obj_new = NewStructureNestedNested() trans = TMemoryBuffer(trans.getvalue()) proto = proto.__class__(trans) fastproto.decode( obj_new, trans, [obj_new.__class__, obj_new.thrift_spec, obj_new.isUnion()], utf8strings=0, protoid=self.PROTO, forward_compatibility=True) self.assertAlmostEqual(obj_new.field.features[0][1], 314.0) self.assertAlmostEqual(obj_new.field.features[0][2], 271.0) trans2 = TMemoryBuffer() proto2 = self.createProto(trans2) obj_new.write(proto2) obj_new2 = NewStructureNestedNested() trans2 = TMemoryBuffer(trans2.getvalue()) proto2 = proto2.__class__(trans2) fastproto.decode( obj_new2, trans2, [obj_new2.__class__, obj_new2.thrift_spec, obj_new2.isUnion()], utf8strings=0, protoid=self.PROTO) self.assertAlmostEqual(obj_new2.field.features[0][1], 314.0) self.assertAlmostEqual(obj_new2.field.features[0][2], 271.0)
def decode(self, recvdata): lenth = len(recvdata) if lenth <= 0: return None, None m = TMemoryBuffer(recvdata) p = TBinaryProtocol(m) packet_header_len, protocol_type = self.decode_packet_header(p) try: h = msg_header() h.read(p) except EOFError: print("Error: msg header decode failure") return None, None # print("package_len:") # print(packet_header_len) # print "msg header: " # pprint(vars(h), indent=4) msg_name = h.msg_name print("recv msg name:", msg_name) if msg_name is None: print("big error ", binascii.hexlify(recvdata)) return 0 try: s = h.msg_name + "_body" if s in globals(): t = globals()[s] else: t = empty body = t() body.read(p) except EOFError: print("Error: msg body decode failure") return None, None # print("body: ") # pprint(vars(body), indent=4, width=24) if msg_name == "ver_resp": self.deal_ver_resp() if msg_name != "shake_hand_resp": print("------------------------------------------------") print("msg header: ") pprint(vars(h), indent=4) print("body: ") pprint(vars(body), indent=4, width=24) message = inner_message() message.header.msg_name = str(msg_name) message.header.src_id = self.get_sid() message.msg_head = h message.msg_content = body return msg_name, message return None, None
def do_task(self, body): link_info = LinkAttr() try: tMemory_o = TMemoryBuffer(body) tBinaryProtocol_o = TBinaryProtocol(tMemory_o) link_info.read(tBinaryProtocol_o) self.scheduler.schedule_task(link_info) except EOFError, e: self.log.warning("can't read LinkAttr from string") return None
def to_string(self, page_info): str_page_info = None try: tMemory_b = TMemoryBuffer() tBinaryProtocol_b = TBinaryServerProtocol(tMemory_b) page_info.write(tBinaryProtocol_b) str_page_info = tMemory_b.getvalue() except EOFError as e: self._log.warning("cann't write data to string") return str_page_info
def to_string(self, download_rsp): str_rsq = None try: tMemory_b = TMemoryBuffer() tBinaryProtocol_b = TBinaryProtocol(tMemory_b) download_rsp.write(tBinaryProtocol_b) str_rsq = tMemory_b.getvalue() self.log.info('data-length is {}'.format(str(len(str_rsq)))) except EOFError, e: self.log.warning("cann't write PageParseInfo to string")