def handle_ready_msg(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Process READY message received from client. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. Note: All exceptions are handled by `handle_exception`. """ if not session.await_ready: # Transmission in progress, READY is out of band raise StopError("Out of band READY message", code=ErrorCode.PROTOCOL_VIOLATION) session.await_ready = False if msg.type_data == 0: # Client either confirmed our zero, or is not ready yet. self.on_schedule_ready(channel, session) else: # All green to transmit DATA session.transmit = msg.type_data if session.socket is PipeSocket.OUTPUT: # Initiate transfer to output (via I/O loop) channel.set_wait_out(True, session)
def handle_ready_msg(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Process READY message received from server. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. Note: All exceptions are handled by `handle_exception`. """ if session.transmit is not None: # Transmission in progress, READY is out of band raise StopError("Out of band READY message", code=ErrorCode.PROTOCOL_VIOLATION) if msg.type_data > 0: # Server is ready batch_size = self.on_server_ready(channel, session, msg.type_data) result = max( 0, min(msg.type_data, self.batch_size if batch_size == -1 else batch_size)) self.send_ready(channel, session, result) if result > 0: # We are ready to transmit as well session.transmit = result if session.socket is PipeSocket.INPUT: # Initiate transfer to server (via I/O loop) channel.set_wait_out(True, session) else: # Server is not ready, but we must send READY(0) back to confirm we've got it! self.send_ready(channel, session, 0)
def handle_data_msg(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Process DATA message received from client. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. Note: All exceptions are handled by `handle_exception`. """ if session.socket is self._flow_in_socket: # DATA flow to us (INPUT for server context, OUTPUT for client context) if session.transmit is None: # Transmission not started, DATA out of band raise StopError("Out of band DATA message", code=ErrorCode.PROTOCOL_VIOLATION) # ACK before processing? if msg.has_ack_req() and not self.confirm_processing: # We must create reply message directly to keep received message reply = FBDPMessage() reply.msg_type = msg.msg_type reply.type_data = msg.type_data reply.set_flag(MsgFlag.ACK_REPLY) if channel.send(msg, session) != 0: raise StopError("ACK-REPLY send failed", code=ErrorCode.ERROR) # Process incoming data self.on_accept_data(channel, session, msg.data_frame) # ACK after processing? if msg.has_ack_req() and self.confirm_processing: if channel.send(self.create_ack_reply(msg), session) != 0: raise StopError("ACK-REPLY send failed", code=ErrorCode.ERROR) session.transmit -= 1 if session.transmit == 0: self._init_new_batch(channel, session) else: # DATA flow from us (OUTPUT for server context, INPUT for client context) if msg.has_ack_reply(): if (session.transmit > 0) and self.send_after_confirmed: # Re-Initiate transfer to output (via I/O loop) if data are available if not self.on_get_data.is_set() or self.on_get_data( channel, session): channel.set_wait_out(True, session) self.on_data_confirmed(channel, session, msg.type_data) else: # Only client attached to PIPE_INPUT can send DATA messages socket: PipeSocket = PipeSocket.OUTPUT \ if self._flow_in_socket is PipeSocket.INPUT else PipeSocket.INPUT raise StopError(f"DATA message sent to {socket.name} socket", code=ErrorCode.PROTOCOL_VIOLATION)
def handle_noop_msg(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Process NOOP message received from peer. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. """ if msg.has_ack_req(): channel.send(self.create_ack_reply(msg), session) self.on_noop(channel, session)
def _send_data(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Sends next DATA message to the client attached to PIPE_OUTPUT. """ error_code = None exc = None try: self.on_produce_data(channel, session, msg) channel.send(msg, session) session.transmit -= 1 if session.transmit > 0: if msg.has_ack_req() and self.send_after_confirmed: channel.set_wait_out(False, session) elif self.on_get_data.is_set(): if not self.on_get_data(channel, session): channel.set_wait_out(False, session) else: channel.set_wait_out(False, session) self._init_new_batch(channel, session) except StopError as err: error_code = getattr(err, 'code', ErrorCode.ERROR) if error_code is not ErrorCode.OK: exc = err except Exception as err: error_code = ErrorCode.INTERNAL_ERROR exc = err if error_code is not None: self.send_close(channel, session, error_code, exc)
def open(self, channel: Channel, address: ZMQAddress, agent: AgentDescriptor, peer_uid: UUID) -> None: """Open connection to Firebird Butler service. Arguments: channel: Channel used for communication with service. address: Service endpoint address. agent: Client agent identification. peer_uid: Client peer ID. """ assert isinstance(channel.protocol, FBSPClient) self.channel = channel self.session = channel.connect(address) self.protocol = channel.protocol self.protocol.send_hello(channel, self.session, agent, PeerDescriptor(peer_uid, os.getpid(), platform.node()), self.token.next()) msg = self.channel.receive(self.timeout) if isinstance(msg, ErrorMessage): raise self.protocol.exception_for(msg) elif msg is TIMEOUT: raise TimeoutError() elif msg is INVALID: raise Error("Invalid response from service") elif not isinstance(msg, WelcomeMessage): raise Error(f"Unexpected {msg.msg_type.name} message from service")
def handle_request(self, channel: Channel, session: Session, msg: ICCPMessage) -> None: """Process REQUEST message received from controller. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. """ result = self.ok_msg() try: if msg.data is Request.CONFIGURE: self.on_config_request(msg.config) except Exception as exc: result = self.error_msg(exc) if not (err_code := channel.send(result, session)): raise StopError("Send to controller failed", err_code=err_code)
def handle_roman(self, channel: Channel, session: FBSPSession, msg: FBSPMessage, protocol: FBSPService) -> None: """Handle REQUEST/ROMAN message. Data frames must contain strings as UTF-8 encoded bytes. We'll send them back in REPLY with Arabic numbers replaced with Roman ones. """ if msg.has_ack_req(): channel.send(protocol.create_ack_reply(msg), session) reply = protocol.create_reply_for(msg) try: for data in msg.data: line = data.decode('utf8') reply.data.append(arabic2roman(line)) channel.send(reply, session) except UnicodeDecodeError: protocol.send_error(session, msg, ErrorCode.BAD_REQUEST, "Data must be UTF-8 bytestrings")
def handle_close_msg(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Process CLOSE message received from client. Calls `on_pipe_closed` and then discards the session. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. """ try: self.on_pipe_closed(channel, session, msg) except: # We don't want to handle this via `handle_exception` and we're closing # the pipe anyway pass finally: channel.discard_session(session)
def send_close(self, channel: Channel, session: FBDPSession, error_code: ErrorCode, exc: Exception = None) -> None: """Sends `CLOSE` message, calls `on_pipe_closed` and then discards the session. Arguments: channel: Channel associate with data pipe. session: Session associated with transmission. error_code: Error code. exc: Exception that caused the error. """ msg = self.create_message_for(MsgType.CLOSE, error_code) if exc: msg.note_exception(exc) try: channel.send(msg, session) self.on_pipe_closed(channel, session, msg, exc) finally: channel.discard_session(session)
def send_open(self, channel: Channel, session: FBDPSession, data_pipe: str, pipe_socket: PipeSocket, data_format: str, parameters: Dict = None) -> None: """Sends `OPEN` message. Arguments: channel: Channel associated with data pipe. session: Session associated with transmission. data_pipe: Data pipe identification. pipe_socket: Connected pipe socket. data_format: Required data format. parameters: Data pipe parameters. Raises: StopError: When sending message fails. """ msg = self.create_message_for(MsgType.OPEN) msg.data_frame.data_pipe = data_pipe msg.data_frame.pipe_socket = pipe_socket.value msg.data_frame.data_format = data_format if parameters: msg.data_frame.parameters.CopyFrom(dict2struct(parameters)) if channel.send(msg, session) != 0: raise StopError("Broken pipe, can't send OPEN message", code=ErrorCode.ERROR) channel.on_output_ready = self._on_output_ready session.pipe = data_pipe session.socket = pipe_socket session.data_format = data_format if parameters: session.params.update(parameters) self.on_init_session(channel, session)
def send_ready(self, channel: Channel, session: FBDPSession, batch_size: int) -> None: """Sends `READY` message. Arguments: channel: Channel associate with data pipe. session: Session associated with transmission. batch_size: Requested data transmission batch size. Raises: StopError: When sending message fails. """ msg = self.create_message_for(MsgType.READY, batch_size) if channel.send(msg, session) != 0: raise StopError("Broken pipe, can't send READY message", code=ErrorCode.ERROR)
def handle_open_msg(self, channel: Channel, session: FBDPSession, msg: FBDPMessage) -> None: """Process OPEN message received from client. Arguments: channel: Channel that received the message. session: Session instance. msg: Received message. Note: All exceptions are handled by `handle_exception`. """ if session.pipe is not None: # Client already attached to data pipe, OPEN out of band raise StopError("Out of band OPEN message", code=ErrorCode.PROTOCOL_VIOLATION) socket = PipeSocket(msg.data_frame.pipe_socket) session.pipe = msg.data_frame.data_pipe session.socket = socket session.data_format = msg.data_frame.data_format session.params.update(struct2dict(msg.data_frame.parameters)) self.on_accept_client(channel, session) self._init_new_batch(channel, session) channel.on_output_ready = self._on_output_ready