diff --git a/common/protob/pb2py b/common/protob/pb2py index 9af4cefe46..e51bf5c3e8 100755 --- a/common/protob/pb2py +++ b/common/protob/pb2py @@ -349,6 +349,7 @@ class Descriptor: ] logging.debug(f"found {len(self.files)} bitcoin-only files") + # message_types is a nested dict: message_name -> wire_enum -> wire_id self.message_types, self.wire_enum_extensions = self.get_wire_enums() # find messages and enums @@ -458,6 +459,9 @@ class Descriptor: else: return None + def wire_enum_has_message(self, enum_name: str, message_name: str) -> bool: + return enum_name in self.message_types.get(message_name, {}) + class PythonRenderer: def __init__(self, descriptor: Descriptor, out_dir="", python_extension="py"): @@ -718,9 +722,12 @@ class RustBlobRenderer: for en in wire_enum_names: name_qstr = self.qstr_map[en] for m in wire_messages: - result.append( - WIRETYPE_ENTRY.build((name_qstr, m.wire_type, self.msg_map[m.name])) - ) + if self.descriptor.wire_enum_has_message(en, m.name): + result.append( + WIRETYPE_ENTRY.build( + (name_qstr, m.wire_type, self.msg_map[m.name]) + ) + ) return b"".join(result) diff --git a/core/embed/rust/librust_qstr.h b/core/embed/rust/librust_qstr.h index 948d23ffe5..74bee344fe 100644 --- a/core/embed/rust/librust_qstr.h +++ b/core/embed/rust/librust_qstr.h @@ -38,6 +38,7 @@ static void _librust_qstrs(void) { MP_QSTR_MAX; MP_QSTR_MESSAGE_NAME; MP_QSTR_MESSAGE_WIRE_TYPE; + MP_QSTR_MessageType; MP_QSTR_Msg; MP_QSTR_MsgDef; MP_QSTR_NONE; diff --git a/core/embed/rust/src/protobuf/defs.rs b/core/embed/rust/src/protobuf/defs.rs index 1771a11351..b934f80429 100644 --- a/core/embed/rust/src/protobuf/defs.rs +++ b/core/embed/rust/src/protobuf/defs.rs @@ -234,3 +234,17 @@ unsafe fn get_enum(enum_offset: u16) -> EnumDef { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::micropython::qstr::Qstr; + + #[test] + fn wire_type_basic_lookup() { + // SignTx: 15 + assert!(find_msg_offset_by_wire(Qstr::MP_QSTR_MessageType.to_u16(), 15).is_some()); + // Guaranteed empty. + assert!(find_msg_offset_by_wire(Qstr::MP_QSTR_CONFIRMED.to_u16(), 15).is_none()); + } +}