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