async def write_message(iface: WireInterface, mtype: int,
mdata: bytes) -> None:
write = loop.wait(iface.iface_num() | io.POLL_WRITE)
# gather data from msg
msize = len(mdata)
# prepare the report buffer with header data
report = bytearray(_REP_LEN)
repofs = _REP_INIT_DATA
ustruct.pack_into(_REP_INIT, report, 0, _REP_MARKER, _REP_MAGIC,
_REP_MAGIC, mtype, msize)
nwritten = 0
while True:
# copy as much as possible to the report buffer
nwritten += utils.memcpy(report, repofs, mdata, nwritten)
# write the report
while True:
await write
n = iface.write(report)
if n == len(report):
break
# if we have more data to write, use continuation reports for it
if nwritten < msize:
repofs = _REP_CONT_DATA
else:
break
async def read_message(iface: WireInterface,
buffer: utils.BufferType) -> Message:
read = loop.wait(iface.iface_num() | io.POLL_READ)
# wait for initial report
report = await read
if report[0] != _REP_MARKER:
raise CodecError("Invalid magic")
_, magic1, magic2, mtype, msize = ustruct.unpack(_REP_INIT, report)
if magic1 != _REP_MAGIC or magic2 != _REP_MAGIC:
raise CodecError("Invalid magic")
read_and_throw_away = False
with buffer_lock:
if msize > len(buffer):
# allocate a new buffer to fit the message
try:
mdata: utils.BufferType = bytearray(msize)
except MemoryError:
mdata = bytearray(_REP_LEN)
read_and_throw_away = True
else:
# reuse a part of the supplied buffer
mdata = memoryview(buffer)[:msize]
# buffer the initial data
nread = utils.memcpy(mdata, 0, report, _REP_INIT_DATA)
while nread < msize:
# wait for continuation report
report = await read
if report[0] != _REP_MARKER:
raise CodecError("Invalid magic")
# buffer the continuation data
if read_and_throw_away:
nread += len(report) - 1
else:
nread += utils.memcpy(mdata, nread, report, _REP_CONT_DATA)
if read_and_throw_away:
raise CodecError("Message too large")
return Message(mtype, mdata)
async def handle_session(iface: WireInterface, session_id: int) -> None:
ctx = Context(iface, session_id)
next_reader = None # type: Optional[codec_v1.Reader]
res_msg = None
req_reader = None
req_type = None
req_msg = None
while True:
try:
if next_reader is None:
# We are not currently reading a message, so let's wait for one.
# If the decoding fails, exception is raised and we try again
# (with the same `Reader` instance, it's OK). Even in case of
# de-synchronized wire communication, report with a message
# header is eventually received, after a couple of tries.
req_reader = ctx.make_reader()
await req_reader.aopen()
if __debug__:
log.debug(
__name__,
"%s:%x receive: %s",
iface.iface_num(),
session_id,
messages.get_type(req_reader.type),
)
else:
# We have a reader left over from earlier. We should process
# this message instead of waiting for new one.
req_reader = next_reader
next_reader = None
# Now we are in a middle of reading a message and we need to decide
# what to do with it, based on its type from the message header.
# From this point on, we should take care to read it in full and
# send a response.
# Take a mark of modules that are imported at this point, so we can
# roll back and un-import any others. Should not raise.
modules = utils.unimport_begin()
# We need to find a handler for this message type. Should not
# raise.
handler = get_workflow_handler(req_reader)
if handler is None:
# If no handler is found, we can skip decoding and directly
# respond with failure, but first, we should read the rest of
# the message reports. Should not raise.
await read_and_throw_away(req_reader)
res_msg = unexpected_message()
else:
# We found a valid handler for this message type.
# Workflow task, declared for the `workflow.on_close` call later.
wf_task = None # type: Optional[loop.Task]
# Here we make sure we always respond with a Failure response
# in case of any errors.
try:
# Find a protobuf.MessageType subclass that describes this
# message. Raises if the type is not found.
req_type = messages.get_type(req_reader.type)
# Try to decode the message according to schema from
# `req_type`. Raises if the message is malformed.
req_msg = await protobuf.load_message(req_reader, req_type)
# At this point, message reports are all processed and
# correctly parsed into `req_msg`.
# Create the workflow task.
wf_task = handler(ctx, req_msg)
# Register the task into the workflow management system.
workflow.on_start(wf_task)
# Run the workflow task. Workflow can do more on-the-wire
# communication inside, but it should eventually return a
# response message, or raise an exception (a rather common
# thing to do). Exceptions are handled in the code below.
res_msg = await wf_task
except UnexpectedMessageError as exc:
# Workflow was trying to read a message from the wire, and
# something unexpected came in. See Context.read() for
# example, which expects some particular message and raises
# UnexpectedMessageError if another one comes in.
# In order not to lose the message, we pass on the reader
# to get picked up by the workflow logic in the beginning of
# the cycle, which processes it in the usual manner.
# TODO:
# We might handle only the few common cases here, like
# Initialize and Cancel.
next_reader = exc.reader
res_msg = None
except Exception as exc:
# Either:
# - the first workflow message had a type that has a
# registered handler, but does not have a protobuf class
# - the first workflow message was not a valid protobuf
# - workflow raised some kind of an exception while running
if __debug__:
if isinstance(exc, ActionCancelled):
log.debug(__name__,
"cancelled: {}".format(exc.message))
else:
log.exception(__name__, exc)
res_msg = failure(exc)
finally:
# De-register the task from the workflow system, if we
# registered it before.
if wf_task is not None:
workflow.on_close(wf_task)
# If a default workflow is on, make sure we do not race
# against the layout that is inside.
# TODO: this is very hacky and complects wire with the ui
if workflow.default_task is not None:
await ui.wait_until_layout_is_running()
if res_msg is not None:
# Either the workflow returned a response, or we created one.
# Write it on the wire. Possibly, the incoming message haven't
# been read in full. We ignore this case here and let the rest
# of the reports get processed while waiting for the message
# header.
# TODO: if the write fails, we do not unimport the loaded modules
await ctx.write(res_msg)
# Cleanup, so garbage collection triggered after un-importing can
# pick up the trash.
req_reader = None
req_type = None
req_msg = None
res_msg = None
handler = None
wf_task = None
# Unload modules imported by the workflow. Should not raise.
utils.unimport_end(modules)
except BaseException as exc:
# The session handling should never exit, just log and continue.
if __debug__:
log.exception(__name__, exc)
