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)