def send(self, msg: HorseMessage, additional_headers: List[MessageHeader] = None) -> HorseResult: """ Sends a raw message to server. Returns true if all data sent over network. :param msg: Sending message :param additional_headers: Additional message headers :return: Successful if the message sent over network. Otherwise returns failed. """ try: writer = ProtocolWriter() if msg.source_len == 0: msg.source = self.id if not msg.message_id: msg.message_id = unique_generator.create() bytes = writer.write(msg, additional_headers) self.__socket.sendall(bytes.getbuffer()) result = HorseResult() result.code = ResultCode.Ok return result except: self.disconnect() result = HorseResult() result.code = ResultCode.Failed return result
def __read_content(self, sock: socket, length: int, msg: HorseMessage): """ Reads message content """ msg.reset_content_stream() content = msg.get_content_stream() left = length buf = bytearray(255) while left > 0: read_count = len(buf) if left < read_count: read_count = left read_done = sock.recv_into(buf, read_count) if read_done == 0: raise IOError() left -= read_done if read_done == len(buf): content.write(buf) else: writebuf = buf[:read_done] content.write(writebuf) msg.recalculate_length()
async def send_get_ack( self, msg: HorseMessage, additional_headers: List[MessageHeader] = None ) -> HorseResult: # Awaitable[HorseResult]: """ Sends a message and waits for acknowledge :param msg: Sending message :param additional_headers: Additional message headers :return: Returns a result after acknowledge received or timed out """ future: asyncio.Future = None try: writer = ProtocolWriter() if msg.source_len == 0: msg.source = self.id if not msg.message_id: msg.message_id = unique_generator.create() msg.pending_response = False if not msg.pending_acknowledge: msg.pending_acknowledge = True tracking = await self.__tracker.track(msg, self.ack_timeout) bytes = writer.write(msg, additional_headers) self.__socket.sendall(bytes.getbuffer()) while not tracking.future.done(): time.sleep(0.001) resp: HorseMessage = await tracking.future result = HorseResult() if resp is None: result.code = ResultCode.RequestTimeout result.reason = "timeout" else: nack_value = resp.get_header( HorseHeaders.NEGATIVE_ACKNOWLEDGE_REASON) if nack_value is None: result.code = ResultCode.Ok result.reason = "" else: result.code = ResultCode.Failed result.reason = nack_value return result except: self.disconnect() if future is not None: await self.__tracker.forget(msg) result = HorseResult() result.code = ResultCode.SendError result.reason = "" return result
async def publish_router( self, router: str, content_type: int, message: str, wait_ack: bool, additional_headers: List[MessageHeader] = None) -> HorseResult: """ Publishes a message to a router :param router: Router name :param content_type: Message content type. It's used in router. Real queue content type may be different (if overwritten in server) :param message: String message content :param wait_ack: If true, waits for acknowledge :param additional_headers: Additional message headers :return: If operation successful, returns Ok """ msg = HorseMessage() msg.type = MessageType.Router msg.content_type = router msg.target = content_type msg.set_content(message) if wait_ack: msg.pending_acknowledge = True return await self.send_get_ack(msg, additional_headers) else: msg.pending_acknowledge = False return self.send(msg, additional_headers)
async def push_queue( self, channel: str, queue: int, message: str, wait_ack: bool, additional_headers: List[MessageHeader] = None) -> HorseResult: """ Pushes a message into a queue :param channel: Channel name of the queue :param queue: Queue Id :param message: String message content :param wait_ack: If true, waits for acknowledge :param additional_headers: Additional message headers :return: If operation successful, returns Ok """ msg = HorseMessage() msg.type = MessageType.QueueMessage msg.content_type = queue msg.target = channel msg.set_content(message) if wait_ack: msg.pending_acknowledge = True return await self.send_get_ack(msg, additional_headers) else: msg.pending_acknowledge = False return self.send(msg, additional_headers)
async def send_direct( self, target: str, content_type: int, message: str, wait_ack: bool, additional_headers: List[MessageHeader] = None) -> HorseResult: """ Sends a direct message to a receiver :param target: Unique Id of the client or @name:name_of_client or @type:type_of_client :param content_type: Message content type :param message: String message content :param wait_ack: If true, message will wait for acknowledge :param additional_headers: Additional message headers :return: """ msg = HorseMessage() msg.type = MessageType.DirectMessage msg.content_type = content_type msg.target = target msg.set_content(message) if wait_ack: msg.pending_acknowledge = True return await self.send_get_ack(msg, additional_headers) else: msg.pending_acknowledge = False return self.send(msg, additional_headers)
async def request( self, msg: HorseMessage, additional_headers: List[MessageHeader] = None) -> HorseResult: """ Sends a request and waits for response :param msg: Request message :param additional_headers: Additional headers :return: Response message is message variable of Horse result """ future: asyncio.Future = None try: writer = ProtocolWriter() if msg.source_len == 0: msg.source = self.id msg.pending_acknowledge = False if not msg.pending_response: msg.pending_response = True tracking = await self.__tracker.track(msg, self.request_timeout) bytes = writer.write(msg, additional_headers) self.__socket.sendall(bytes.getbuffer()) while not tracking.future.done(): time.sleep(0.001) resp: HorseMessage = await tracking.future result = HorseResult() if resp is None: result.code = ResultCode.RequestTimeout result.reason = "timeout" else: result.code = resp.content_type result.message = resp result.reason = resp.get_header(HorseHeaders.REASON) return result except: self.disconnect() if future is not None: await self.__tracker.forget(msg) result = HorseResult() result.code = ResultCode.SendError return result
async def leave(self, channel: str, wait_ack: bool = False) -> HorseResult: """ Leavess from a channel :param channel: Channel name :param wait_ack: If true, waits for acknowledge from server :return: If waits for ack, ack result. Otherview Ok if message is sent successfuly """ msg = HorseMessage() msg.type = MessageType.Server msg.content_type = KnownContentTypes.LEAVE.value msg.target = channel msg.pending_response = wait_ack result: HorseResult if wait_ack: msg.message_id = unique_generator.create() result = await self.request(msg) else: result = self.send(msg) # remove channel from joined channel list if result.code == ResultCode.Ok: self.__joined_channels.remove(channel) return result
async def join(self, channel: str, wait_ack: bool = False) -> HorseResult: """ Joins to a channel :param channel: Channel name :param wait_ack: If true, waits for acknowledge from server :return: If waits for ack, ack result. Otherview Ok if message is sent successfuly """ msg = HorseMessage() msg.type = MessageType.Server msg.content_type = KnownContentTypes.JOIN.value msg.target = channel msg.pending_response = wait_ack result: HorseResult if wait_ack: msg.message_id = unique_generator.create() result = await self.request(msg) else: result = self.send(msg) # add channel to joined list (if not already added) if result.code == ResultCode.Ok: has = next((x for x in self.__joined_channels if x == channel), None) if not has: self.__joined_channels.append(channel) return result
def __read_headers(self, sock: socket, msg: HorseMessage): """ Reads message headers """ # read unsigned int 16 for headers length header_bytes = self.__read_certain(sock, 2) header_length = int.from_bytes(header_bytes, byteorder='little', signed=False) buffer = self.__read_certain(sock, header_length) header_str: str = buffer.decode('UTF-8') # split header data with CR LF, split per line with : and get key value pair lines = header_str.split('\r\n') for line in lines: if len(line) < 2: continue index = line.index(':') if index < 1: continue key = line[:index].strip() value = line[index + 1:].strip() msg.add_header(key, value)
def read(self, sock: socket) -> HorseMessage: """ Reads a Horse message from socket. :returns None if frame is corrupted or connection is lost while reading """ # magic num 8: protocol frame must have at least 8 bytes frame_bytes = self.__read_certain(sock, 8) if frame_bytes is None: return None msg = HorseMessage() message_len = self.__read_frame(frame_bytes, sock, msg) if msg.has_header: self.__read_headers(sock, msg) if message_len > 0: self.__read_content(sock, message_len, msg) return msg
def __handshake(self) -> bool: """ Sends HMQP handshake message and reads handshake response. :returns true if handshake is successful """ self.__socket.sendall("HMQP/2.1".encode('UTF-8')) # create handshake message properties content = 'CONNECT /\r\n' if self.id: content += HorseHeaders.create_line(HorseHeaders.CLIENT_ID, self.id) if self.name: content += HorseHeaders.create_line(HorseHeaders.CLIENT_NAME, self.name) if self.type: content += HorseHeaders.create_line(HorseHeaders.CLIENT_TYPE, self.type) if self.token: content += HorseHeaders.create_line(HorseHeaders.CLIENT_TOKEN, self.token) if self.headers: for h in self.headers: content += HorseHeaders.create_line(h.key, h.value) msg = HorseMessage() msg.type = MessageType.Server msg.content_type = KnownContentTypes.HELLO.value msg.set_content(content) sent = self.send(msg) if not sent: return False hr_result = self.__read_certain(8) if hr_result is None: return False hs_response = hr_result.decode('UTF-8') return hs_response == "HMQP/2.1"
def __read_frame(self, array: bytearray, sock: socket, msg: HorseMessage) -> int: """ Reads frame data """ first = array[0] second = array[1] if first >= 128: msg.first_acquirer = True first -= 128 if first >= 64: msg.high_priority = True first -= 64 msg.type = first if second >= 128: msg.pending_response = True second -= 128 if second >= 64: msg.pending_acknowledge = True second -= 64 if second >= 32 and msg.type != MessageType.Ping.value and msg.type != MessageType.Pong.value: msg.has_header = True second -= 32 msg.ttl = second id_len = array[2] source_len = array[3] target_len = array[4] size_bytes = [array[5], array[6]] msg.content_type = int.from_bytes(size_bytes, byteorder='little', signed=False) message_len = 0 # length is unsigned in 16 if array[7] == 253: short_len = self.__read_certain(sock, 2) message_len = int.from_bytes(short_len, byteorder='little', signed=False) # length is unsigned int 32 elif array[7] == 254: int_len = self.__read_certain(sock, 4) message_len = int.from_bytes(int_len, byteorder='little', signed=False) # length is unsigned int 64 elif array[7] == 255: long_len = self.__read_certain(sock, 8) message_len = int.from_bytes(long_len, byteorder='little', signed=False) # length is byte else: message_len = array[7] if id_len > 0: id_bytes = self.__read_certain(sock, id_len) msg.message_id = id_bytes.decode('UTF-8') if source_len > 0: source_bytes = self.__read_certain(sock, source_len) msg.source = source_bytes.decode('UTF-8') if target_len > 0: target_bytes = self.__read_certain(sock, target_len) msg.target = target_bytes.decode('UTF-8') return message_len
async def pull( self, request: PullRequest, each_msg_func: Callable[[int, HorseMessage], None]) -> PullContainer: msg = HorseMessage() msg.type = MessageType.QueuePullRequest msg.message_id = unique_generator.create() msg.target = request.channel msg.content_type = request.queue_id msg.add_header(HorseHeaders.COUNT, str(request.count)) if request.clear_after == ClearDecision.AllMessages.value: msg.add_header(HorseHeaders.CLEAR, "all") elif request.clear_after == ClearDecision.PriorityMessages.value: msg.add_header(HorseHeaders.CLEAR, "High-Priority") elif request.clear_after == ClearDecision.Messages.value: msg.add_header(HorseHeaders.CLEAR, "Default-Priority") if request.get_counts: msg.add_header(HorseHeaders.INFO, "yes") if request.order == MessageOrder.LIFO.value: msg.add_header(HorseHeaders.ORDER, HorseHeaders.LIFO) if request.request_headers: for header in request.request_headers: msg.add_header(header.key, header.value) container = PullContainer() container.request_id = msg.message_id container.request_count = request.count container.received_count = 0 container.status = PullProcess.Receiving container.messages = [] container.each_msg_func = each_msg_func container.last_received = datetime.utcnow() container.future = asyncio.Future() self.__pull_containers[msg.message_id] = container send_result = self.send(msg) if send_result.code != ResultCode.Ok: self.__pull_containers.pop(msg.message_id) return container while not container.future.done(): time.sleep(0.001) diff = datetime.utcnow() - container.last_received if diff > self.request_timeout: self.__pull_containers.pop(msg.message_id) container.future.set_result(None) break await container.future return container
def response( self, request_msg: HorseMessage, status: ResultCode, response_content: str = None, additional_headers: List[MessageHeader] = None) -> HorseResult: """ Sends a response message to a request :param request_msg: Request message :param status: Response message status :param response_content: Response message content :param additional_headers: Additional message headers :return: Returns ok, if sent successfully """ msg = HorseMessage() msg.type = MessageType.Response msg.content_type = status.value msg.message_id = request_msg.message_id msg.high_priority = request_msg.high_priority msg.first_acquirer = request_msg.first_acquirer if request_msg.type == MessageType.QueueMessage.value: msg.target = request_msg.target else: msg.target = request_msg.source if not response_content is None: msg.set_content(response_content) return self.send(msg, additional_headers)
def __send_ack(self, message: HorseMessage, reason: str) -> HorseResult: """ Sends an acknowledge message :param message: Message that will be acknowledged :param reason: Negative acknowledge reason. If None, acknowledge is positive :return: Returns Ok if sent successfuly """ msg = HorseMessage() msg.type = MessageType.Acknowledge msg.content_type = message.content_type msg.message_id = message.message_id msg.first_acquirer = message.first_acquirer if message.type == MessageType.DirectMessage.value: msg.high_priority = True msg.source = message.target msg.target = message.source else: msg.high_priority = False msg.target = message.target if reason: msg.add_header(HorseHeaders.NEGATIVE_ACKNOWLEDGE_REASON, reason) return self.send(msg)