async def read_any(self, allowed_types: Iterable[int]) -> MessageType:
reader = self.make_reader()
if __debug__:
log.debug(
__name__,
"%s:%x expect: %s",
self.iface.iface_num(),
self.sid,
allowed_types,
)
await reader.aopen() # wait for the message header
# if we got a message with unexpected type, raise the reader via
# `UnexpectedMessageError` and let the session handler deal with it
if reader.type not in allowed_types:
raise UnexpectedMessageError(reader)
# find the protobuf type
exptype = messages.get_type(reader.type)
if __debug__:
log.debug(__name__, "%s:%x read: %s", self.iface.iface_num(),
self.sid, exptype)
# parse the message and return it
return await protobuf.load_message(reader, exptype)
async def write(self,
msg: protobuf.MessageType,
field_cache: protobuf.FieldCache = None) -> None:
if __debug__:
log.debug(__name__, "%s:%x write: %s", self.iface.iface_num(),
self.sid, msg)
if field_cache is None:
field_cache = self._field_cache
# write the message
msg_size = protobuf.count_message(msg, field_cache)
# prepare buffer
if msg_size <= len(self.buffer_writer.buffer):
# reuse preallocated
buffer_writer = self.buffer_writer
else:
# message is too big, we need to allocate a new buffer
buffer_writer = utils.BufferWriter(bytearray(msg_size))
buffer_writer.seek(0)
protobuf.dump_message(buffer_writer, msg, field_cache)
await codec_v1.write_message(
self.iface,
msg.MESSAGE_WIRE_TYPE,
memoryview(buffer_writer.buffer)[:msg_size],
)
# make sure we don't keep around fields of all protobuf types ever
self._field_cache.clear()
def closedefault():
global default
if default:
default.close()
default = None
log.debug(__name__, 'closedefault')
def _step(task: Task, value: Any) -> None:
try:
if isinstance(value, BaseException):
result = task.throw(value) # type: ignore
# error: Argument 1 to "throw" of "Coroutine" has incompatible type "Exception"; expected "Type[BaseException]"
# rationale: In micropython, generator.throw() accepts the exception object directly.
else:
result = task.send(value)
except StopIteration as e: # as e:
if __debug__:
log.debug(__name__, "finish: %s", task)
finalize(task, e.value)
except Exception as e:
if __debug__:
log.exception(__name__, e)
finalize(task, e)
else:
if isinstance(result, Syscall):
result.handle(task)
elif result is None:
schedule(task)
else:
if __debug__:
log.error(__name__, "unknown syscall: %s", result)
if after_step_hook:
after_step_hook()
async def read(
self, expected_type: type[LoadedMessageType]) -> LoadedMessageType:
if __debug__:
log.debug(
__name__,
"%s:%x expect: %s",
self.iface.iface_num(),
self.sid,
expected_type.MESSAGE_NAME,
)
# Load the full message into a buffer, parse out type and data payload
msg = await self.read_from_wire()
# If we got a message with unexpected type, raise the message via
# `UnexpectedMessageError` and let the session handler deal with it.
if msg.type != expected_type.MESSAGE_WIRE_TYPE:
raise UnexpectedMessageError(msg)
if __debug__:
log.debug(
__name__,
"%s:%x read: %s",
self.iface.iface_num(),
self.sid,
expected_type.MESSAGE_NAME,
)
workflow.idle_timer.touch()
# look up the protobuf class and parse the message
return _wrap_protobuf_load(msg.data, expected_type)
def _cbor_encode(value):
if isinstance(value, int):
yield _header(_CBOR_UNSIGNED_INT, value)
elif isinstance(value, bytes):
yield _header(_CBOR_BYTE_STRING, len(value))
yield value
elif isinstance(value, bytearray):
yield _header(_CBOR_BYTE_STRING, len(value))
yield bytes(value)
elif isinstance(value, list):
# definite-length valued list
yield _header(_CBOR_ARRAY, len(value))
for x in value:
yield from _cbor_encode(x)
elif isinstance(value, dict):
yield _header(_CBOR_MAP, len(value))
for k, v in value.items():
yield from _cbor_encode(k)
yield from _cbor_encode(v)
elif isinstance(value, Tagged):
yield _header(_CBOR_TAG, value.tag)
yield from _cbor_encode(value.value)
elif isinstance(value, IndefiniteLengthArray):
yield bytes([_CBOR_ARRAY + 31])
for x in value.array:
yield from _cbor_encode(x)
yield bytes([_CBOR_PRIMITIVE + 31])
elif isinstance(value, Raw):
yield value.value
else:
if __debug__:
log.debug(__name__, "not implemented (encode): %s", type(value))
raise NotImplementedError()
async def read_any(
self, expected_wire_types: Iterable[int]) -> protobuf.MessageType:
if __debug__:
log.debug(
__name__,
"%s:%x expect: %s",
self.iface.iface_num(),
self.sid,
expected_wire_types,
)
# Load the full message into a buffer, parse out type and data payload
msg = await self.read_from_wire()
# If we got a message with unexpected type, raise the message via
# `UnexpectedMessageError` and let the session handler deal with it.
if msg.type not in expected_wire_types:
raise UnexpectedMessageError(msg)
# find the protobuf type
exptype = protobuf.type_for_wire(msg.type)
if __debug__:
log.debug(
__name__,
"%s:%x read: %s",
self.iface.iface_num(),
self.sid,
exptype.MESSAGE_NAME,
)
workflow.idle_timer.touch()
# parse the message and return it
return _wrap_protobuf_load(msg.data, exptype)
async def read(
self, expected_type: Type[protobuf.LoadedMessageType]
) -> protobuf.LoadedMessageType:
reader = self.make_reader()
if __debug__:
log.debug(
__name__,
"%s:%x expect: %s",
self.iface.iface_num(),
self.sid,
expected_type,
)
# Wait for the message header, contained in the first report. After
# we receive it, we have a message type to match on.
await reader.aopen()
# If we got a message with unexpected type, raise the reader via
# `UnexpectedMessageError` and let the session handler deal with it.
if reader.type != expected_type.MESSAGE_WIRE_TYPE:
raise UnexpectedMessageError(reader)
if __debug__:
log.debug(
__name__,
"%s:%x read: %s",
self.iface.iface_num(),
self.sid,
expected_type,
)
# parse the message and return it
return await protobuf.load_message(reader, expected_type)
def closedefault():
global default
if default:
loop.close(default)
default = None
log.debug(__name__, 'closedefault')
async def write(self, msg: protobuf.MessageType) -> None:
if __debug__:
log.debug(
__name__,
"%s:%x write: %s",
self.iface.iface_num(),
self.sid,
msg.MESSAGE_NAME,
)
# cannot write message without wire type
assert msg.MESSAGE_WIRE_TYPE is not None
msg_size = protobuf.encoded_length(msg)
if msg_size <= len(self.buffer):
# reuse preallocated
buffer = self.buffer
else:
# message is too big, we need to allocate a new buffer
buffer = bytearray(msg_size)
msg_size = protobuf.encode(buffer, msg)
await codec_v1.write_message(
self.iface,
msg.MESSAGE_WIRE_TYPE,
memoryview(buffer)[:msg_size],
)
async def read_any(
self, expected_wire_types: Iterable[int]) -> protobuf.MessageType:
reader = self.make_reader()
if __debug__:
log.debug(
__name__,
"%s:%x expect: %s",
self.iface.iface_num(),
self.sid,
expected_wire_types,
)
# Wait for the message header, contained in the first report. After
# we receive it, we have a message type to match on.
await reader.aopen()
# If we got a message with unexpected type, raise the reader via
# `UnexpectedMessageError` and let the session handler deal with it.
if reader.type not in expected_wire_types:
raise UnexpectedMessageError(reader)
# find the protobuf type
exptype = messages.get_type(reader.type)
if __debug__:
log.debug(__name__, "%s:%x read: %s", self.iface.iface_num(),
self.sid, exptype)
# parse the message and return it
return await protobuf.load_message(reader, exptype)
def _finalize_default(task: loop.Task, value: Any) -> None:
"""Finalizer for the default task. Cleans up globals and restarts the default
in case no other task is running."""
global default_task
if default_task is task:
if __debug__:
log.debug(__name__, "default closed: %s", task)
default_task = None
if not tasks:
# No registered workflows are running and we are in the default task
# finalizer, so when this function finished, nothing will be running.
# We must schedule a new instance of the default now.
if default_constructor is not None:
start_default(default_constructor)
else:
raise RuntimeError # no tasks and no default constructor
else:
if __debug__:
log.warning(
__name__,
"default task does not match: task=%s, default_task=%s",
task,
default_task,
)
async def sign_tx(ctx, received_msg, keychain):
state = State(ctx)
mods = utils.unimport_begin()
# Splitting ctx.call() to write() and read() helps to reduce memory fragmentation
# between calls.
while True:
if __debug__:
log.debug(__name__, "#### F: %s, A: %s", gc.mem_free(),
gc.mem_alloc())
gc.collect()
gc.threshold(gc.mem_free() // 4 + gc.mem_alloc())
result_msg, accept_msgs = await sign_tx_dispatch(
state, received_msg, keychain)
if accept_msgs is None:
break
await ctx.write(result_msg)
del (result_msg, received_msg)
utils.unimport_end(mods)
received_msg = await ctx.read(accept_msgs)
utils.unimport_end(mods)
return result_msg
def close_default() -> None:
"""Explicitly close the default workflow task."""
if default_task:
if __debug__:
log.debug(__name__, "close default")
# We let the `_finalize_default` reset the global.
loop.close(default_task)
async def _sync_step(s, ctx, tds):
if not tds.tdis:
raise wire.DataError("Empty")
kis = []
buff = bytearray(32 * 3)
buff_mv = memoryview(buff)
await confirms.keyimage_sync_step(ctx, s.current_output, s.num_outputs)
for td in tds.tdis:
s.current_output += 1
if s.current_output >= s.num_outputs:
raise wire.DataError("Too many outputs")
if __debug__:
log.debug(__name__, "ki_sync, step i: %d", s.current_output)
# Update the control hash
s.hasher.update(key_image.compute_hash(td))
# Compute keyimage + signature
ki, sig = key_image.export_key_image(s.creds, s.subaddresses, td)
# Serialize into buff
crypto.encodepoint_into(buff_mv[0:32], ki)
crypto.encodeint_into(buff_mv[32:64], sig[0][0])
crypto.encodeint_into(buff_mv[64:], sig[0][1])
# Encrypt with enc_key
nonce, ciph, _ = chacha_poly.encrypt(s.enc_key, buff)
kis.append(MoneroExportedKeyImage(iv=nonce, blob=ciph))
return MoneroKeyImageSyncStepAck(kis=kis)
async def protobuf_workflow(ctx: Context, reader: codec_v1.Reader,
handler: Handler, *args: Any) -> None:
from trezor.messages.Failure import Failure
req = await protobuf.load_message(reader, messages.get_type(reader.type))
if __debug__:
log.debug(__name__, "%s:%x request: %s", ctx.iface.iface_num(),
ctx.sid, req)
try:
res = await handler(ctx, req, *args)
except UnexpectedMessageError:
# session handler takes care of this one
raise
except Error as exc:
# respond with specific code and message
await ctx.write(Failure(code=exc.code, message=exc.message))
raise
except Exception as e:
# respond with a generic code and message
message = "Firmware error"
if __debug__:
message = "{}: {}".format(type(e), e)
await ctx.write(
Failure(code=FailureType.FirmwareError, message=message))
raise
if res:
# respond with a specific response
await ctx.write(res)
def _dispatch_reports():
while True:
report, = yield loop.Select(_interface)
report = memoryview(report)
if __debug__:
log.debug(__name__, 'read report %s', ubinascii.hexlify(report))
sessions.dispatch(report, _session_open, _session_close,
_session_unknown)
async def dispatch_DebugLinkWatchLayout(
ctx: wire.Context, msg: DebugLinkWatchLayout
) -> Success:
global watch_layout_changes
layout_change_chan.putters.clear()
watch_layout_changes = bool(msg.watch)
log.debug(__name__, "Watch layout changes: {}".format(watch_layout_changes))
return Success()
def _on_start(workflow: loop.spawn) -> None:
"""
Called after creating a workflow task, but before running it.
"""
# Take note that this workflow task is running.
if __debug__:
log.debug(__name__, "start: %s", workflow.task)
idle_timer.touch()
tasks.add(workflow)
async def button_request(
ctx: wire.GenericContext,
br_type: str,
code: ButtonRequestType = ButtonRequestType.Other,
) -> None:
if __debug__:
log.debug(__name__, "ButtonRequest.type=%s", br_type)
workflow.close_others()
await ctx.call(ButtonRequest(code=code), ButtonAck)
def log_trace(x=None) -> None:
log.debug(
__name__,
"Log trace %s, ... F: %s A: %s, S: %s",
x,
gc.mem_free(),
gc.mem_alloc(),
micropython.stack_use(),
)
async def dispatch_DebugLinkWatchLayout(
ctx: wire.Context, msg: DebugLinkWatchLayout
) -> Success:
from trezor import ui
layout_change_chan.putters.clear()
await ui.wait_until_layout_is_running()
storage.watch_layout_changes = bool(msg.watch)
log.debug(__name__, "Watch layout changes: %s", storage.watch_layout_changes)
return Success()
async def interact(
ctx: wire.GenericContext,
layout: LayoutType,
brtype: str,
brcode: EnumTypeButtonRequestType = ButtonRequestType.Other,
) -> Any:
log.debug(__name__, "ButtonRequest.type={}".format(brtype))
workflow.close_others()
await ctx.call(ButtonRequest(code=brcode), ButtonAck)
return await ctx.wait(layout)
def startdefault(handler):
global default
global default_handler
if not default:
default_handler = handler
default = handler()
loop.schedule(default)
ui.display.backlight(ui.BACKLIGHT_NORMAL)
log.debug(__name__, 'startdefault')
def close(self) -> None:
"""Shut down the spawned task.
If another caller is awaiting its result it will get a TaskClosed exception.
If the task was already finished, the call has no effect.
"""
if not self.finished:
if __debug__:
log.debug(__name__, "close spawned task: %s", self.task)
close(self.task)
def __init__(self, task: Task) -> None:
self.task = task
self.callback: Optional[Task] = None
self.finalizer_callback: Optional[Callable[["spawn"], None]] = None
self.finished = False
self.return_value: Any = None
# schedule task immediately
if __debug__:
log.debug(__name__, "spawn new task: %s", task)
schedule(task, finalizer=self._finalize)
def mem_trace(self, x=None, collect=False):
if __debug__:
log.debug(
__name__,
"Log trace: %s, ... F: %s A: %s",
x,
gc.mem_free(),
gc.mem_alloc(),
)
if collect:
gc.collect()
async def inner(*args, **kwargs):
mods = set(sys.modules)
try:
ret = await genfunc(*args, **kwargs)
finally:
for mod in sys.modules:
if mod not in mods:
log.debug(__name__, 'unimport %s', mod)
del sys.modules[mod]
gc.collect()
return ret
def _cbor_encode(value: Value) -> Iterator[bytes]:
if isinstance(value, int):
if value >= 0:
yield _header(_CBOR_UNSIGNED_INT, value)
else:
yield _header(_CBOR_NEGATIVE_INT, -1 - value)
elif isinstance(value, bytes):
yield _header(_CBOR_BYTE_STRING, len(value))
yield value
elif isinstance(value, bytearray):
yield _header(_CBOR_BYTE_STRING, len(value))
yield bytes(value)
elif isinstance(value, str):
encoded_value = value.encode()
yield _header(_CBOR_TEXT_STRING, len(encoded_value))
yield encoded_value
elif isinstance(value, (list, tuple)):
# definite-length valued list
yield _header(_CBOR_ARRAY, len(value))
for x in value:
yield from _cbor_encode(x)
elif isinstance(value, dict):
yield _header(_CBOR_MAP, len(value))
sorted_map = sorted((encode(k), v) for k, v in value.items())
for k, v in sorted_map:
yield k
yield from _cbor_encode(v)
elif isinstance(value, OrderedMap):
yield _header(_CBOR_MAP, len(value))
for k, v in value:
yield encode(k)
yield from _cbor_encode(v)
elif isinstance(value, Tagged):
yield _header(_CBOR_TAG, value.tag)
yield from _cbor_encode(value.value)
elif isinstance(value, IndefiniteLengthArray):
yield bytes([_CBOR_ARRAY + 31])
for x in value.array:
yield from _cbor_encode(x)
yield bytes([_CBOR_PRIMITIVE + 31])
elif isinstance(value, bool):
if value:
yield bytes([_CBOR_PRIMITIVE + _CBOR_TRUE])
else:
yield bytes([_CBOR_PRIMITIVE + _CBOR_FALSE])
elif isinstance(value, Raw):
yield value.value
elif value is None:
yield bytes([_CBOR_PRIMITIVE + _CBOR_NULL])
else:
if __debug__:
log.debug(__name__, "not implemented (encode): %s", type(value))
raise NotImplementedError
def perf_info_debug():
while True:
queue_len = len(loop._scheduled_tasks)
delay_avg = sum(loop.log_delay_rb) / loop.log_delay_rb_len
delay_last = loop.log_delay_rb[loop.log_delay_pos]
mem_alloc = gc.mem_alloc()
gc.collect()
log.debug(__name__, "mem_alloc: %s/%s, delay_avg: %d, delay_last: %d, queue_len: %d",
mem_alloc, gc.mem_alloc(), delay_avg, delay_last, queue_len)
yield loop.Sleep(1000000)