refactor(rust/trezor-thp): buffer size hint in PacketInResult::Route

diff --git a/rust/trezor-thp/src/channel/device.rs b/rust/trezor-thp/src/channel/device.rs
index 82b0fe50..f9669c57 100644
--- a/rust/trezor-thp/src/channel/device.rs
+++ b/rust/trezor-thp/src/channel/device.rs
@@ -7,12 +7,13 @@ use crate::{
         HANDSHAKE_BUFFER_DTH_LEN, HANDSHAKE_BUFFER_HTD_LEN, Nonce, PRIVKEY_LEN, PacketInResult,
         PairingState, ReceiveState, SendState, noise::NoiseHandshake,
     },
+    control_byte::ControlByte,
     credential::CredentialVerifier,
     error::TransportError,
     fragment::{Fragmenter, Reassembler},
     header::{
         BROADCAST_CHANNEL_ID, HandshakeMessage, Header, MAX_CHANNEL_ID, MIN_CHANNEL_ID,
-        channel_id_valid, parse_cb_channel,
+        channel_id_valid, parse_channel_length,
     },
     util::prepare_zeroed,
 };
@@ -191,19 +192,19 @@ where
     B: Backend,
 {
     fn packet_in(&mut self, packet_buffer: &[u8], _receive_buffer: &mut [u8]) -> PacketInResult {
-        let Ok((cb, channel_id, _rest)) = parse_cb_channel(packet_buffer) else {
-            // parse_cb_channel already writes to log
+        let Ok((cb, _)) = ControlByte::parse(packet_buffer) else {
+            // ControlByte::parse already writes to log
             return PacketInResult::ignore(Error::malformed_data());
         };
         if cb.is_codec_v1() {
             return self.handle_v1(packet_buffer);
         }
-        if !channel_id_valid(channel_id) {
-            log::warn!("Invalid channel id {:04x}.", channel_id);
+        let Ok((channel_id, len)) = parse_channel_length(cb, packet_buffer) else {
+            // parse_channel_length already writes to log
             return PacketInResult::ignore(Error::malformed_data());
-        }
+        };
         if channel_id != BROADCAST_CHANNEL_ID {
-            return PacketInResult::route(channel_id);
+            return PacketInResult::route(channel_id, len);
         }
         PacketInResult::from_result(self.handle_broadcast(packet_buffer).map(|is_allocation| {
             if is_allocation {
diff --git a/rust/trezor-thp/src/channel/host.rs b/rust/trezor-thp/src/channel/host.rs
index f14c9f14..7d629ddf 100644
--- a/rust/trezor-thp/src/channel/host.rs
+++ b/rust/trezor-thp/src/channel/host.rs
@@ -8,12 +8,10 @@ use crate::{
         MAX_DEVICE_PROPERTIES_LEN, Nonce, PacketInResult, PairingState, ReceiveState,
         noise::NoiseHandshake,
     },
+    control_byte::ControlByte,
     credential::CredentialStore,
     fragment::{Fragmenter, Reassembler},
-    header::{
-        BROADCAST_CHANNEL_ID, HandshakeMessage, Header, channel_id_valid, parse_cb_channel,
-        parse_u16,
-    },
+    header::{BROADCAST_CHANNEL_ID, HandshakeMessage, Header, parse_channel_length, parse_u16},
     util::prepare_zeroed,
 };
 
@@ -236,16 +234,20 @@ where
     B: Backend,
 {
     fn packet_in(&mut self, packet_buffer: &[u8], _receive_buffer: &mut [u8]) -> PacketInResult {
-        let Ok((cb, channel_id, _rest)) = parse_cb_channel(packet_buffer) else {
-            // parse_cb_channel already writes to log
+        let Ok((cb, _)) = ControlByte::parse(packet_buffer) else {
+            // ControlByte::parse already writes to log
             return PacketInResult::ignore(Error::malformed_data());
         };
-        if !channel_id_valid(channel_id) {
-            log::warn!("Invalid channel id {:04x}.", channel_id);
+        if cb.is_codec_v1() {
+            // host does not respond to v1 packets
             return PacketInResult::ignore(Error::malformed_data());
         }
-        if channel_id != BROADCAST_CHANNEL_ID && !cb.is_codec_v1() {
-            return PacketInResult::route(channel_id);
+        let Ok((channel_id, length)) = parse_channel_length(cb, packet_buffer) else {
+            // parse_channel_length already writes to log
+            return PacketInResult::ignore(Error::malformed_data());
+        };
+        if channel_id != BROADCAST_CHANNEL_ID {
+            return PacketInResult::route(channel_id, length);
         }
         PacketInResult::from_result(self.handle_broadcast(packet_buffer))
     }
diff --git a/rust/trezor-thp/src/channel/mod.rs b/rust/trezor-thp/src/channel/mod.rs
index f1b8ea71..6417970c 100644
--- a/rust/trezor-thp/src/channel/mod.rs
+++ b/rust/trezor-thp/src/channel/mod.rs
@@ -476,6 +476,15 @@ pub enum PacketInResult {
     Route {
         /// Channel id of the destination. Never a broadcast.
         channel_id: u16,
+        /// If Some, initiation packet was received for message with the given payload length.
+        /// Make sure you're calling [`ChannelIO::packet_in`] with receive buffer at least as large.
+        /// Please note the size is returned before checksum verification and can be unusually
+        /// large in presence of bit errors.
+        /// In particular, this is `None` for ACK messages, continuations and transport errors.
+        // NOTE: would be neat to get PacketInResult to fit in 32bit word but how? Buffer size can
+        // be e.g. 12 bit integer denoting 16-byte blocks but not sure how to pack it. The buffer
+        // size in `Accepted` is not very important and can probably be turned into a bool.
+        buffer_size: Option<NonZeroU16>,
     },
     /// Channel allocation request/response was received. Event loop should call
     /// [`Mux::channel_alloc`] to create new channel object. Only [`device::Mux`] and [`host::Mux`]
@@ -516,8 +525,18 @@ impl PacketInResult {
         Self::TransportError { error: e }
     }
 
-    const fn route(channel_id: u16) -> Self {
-        Self::Route { channel_id }
+    const fn route(channel_id: u16, payload_len: Option<u16>) -> Self {
+        let buffer_size = match payload_len {
+            // Using larger buffer should be safe. Empty messages are valid on the transport
+            // layer but they will fail the minimum app header length check anyway.
+            Some(0) => Some(NonZeroU16::new(1).unwrap()),
+            Some(n) => Some(NonZeroU16::new(n).unwrap()),
+            None => None,
+        };
+        Self::Route {
+            channel_id,
+            buffer_size,
+        }
     }
 
     const fn fail(error: Error) -> Self {
@@ -827,7 +846,11 @@ impl<R: Role, B: Backend> ChannelIO for Channel<R, B> {
             log::warn!("[{:04x}] Nothing to retransmit.", self.channel_id);
             return Ok(());
         };
-        log::debug!("[{:04x}] Retransmitting message.", self.channel_id);
+        log::debug!(
+            "[{:04x}] Retransmitting message, retry {}.",
+            self.channel_id,
+            retry
+        );
         fragmenter.reset();
         *retry = retry.saturating_add(1);
         Ok(())
diff --git a/rust/trezor-thp/src/channel/test.rs b/rust/trezor-thp/src/channel/test.rs
index 78468031..20a371d4 100644
--- a/rust/trezor-thp/src/channel/test.rs
+++ b/rust/trezor-thp/src/channel/test.rs
@@ -1114,9 +1114,21 @@ fn test_invalid_channel_id() -> Result<()> {
     let (mut hm, mut dm, _cids) = create_mux();
     // muxes return Route(cid) for valid non-broadcast channel
     let pir = dm.packet_in(&make_packet(66));
-    assert_eq!(pir, PacketInResult::Route { channel_id: 66 });
+    assert_eq!(
+        pir,
+        PacketInResult::Route {
+            channel_id: 66,
+            buffer_size: 16.try_into().ok(),
+        }
+    );
     let pir = hm.packet_in(&make_packet(66));
-    assert_eq!(pir, PacketInResult::Route { channel_id: 66 });
+    assert_eq!(
+        pir,
+        PacketInResult::Route {
+            channel_id: 66,
+            buffer_size: 16.try_into().ok(),
+        }
+    );
 
     // muxes return error for invalid channel ids
     for channel_id in INVALID {
diff --git a/rust/trezor-thp/src/control_byte.rs b/rust/trezor-thp/src/control_byte.rs
index 9d6bbd7b..88339fbe 100644
--- a/rust/trezor-thp/src/control_byte.rs
+++ b/rust/trezor-thp/src/control_byte.rs
@@ -131,6 +131,15 @@ impl ControlByte {
     pub fn with_sync_bits(self, sb: SyncBits) -> Self {
         Self(self.0 & !SYNC_MASK | <SyncBits as Into<u8>>::into(sb))
     }
+
+    pub fn parse(packet_buffer: &[u8]) -> Result<(Self, &[u8]), Error> {
+        let Some((first_byte, rest)) = packet_buffer.split_first() else {
+            log::error!("Packet is empty.");
+            return Err(Error::malformed_data());
+        };
+        let cb = Self::try_from(*first_byte)?;
+        Ok((cb, rest))
+    }
 }
 
 impl TryFrom<u8> for ControlByte {
diff --git a/rust/trezor-thp/src/header.rs b/rust/trezor-thp/src/header.rs
index 16539275..b036d6d0 100644
--- a/rust/trezor-thp/src/header.rs
+++ b/rust/trezor-thp/src/header.rs
@@ -74,13 +74,48 @@ pub(crate) fn parse_u16(buffer: &[u8]) -> Result<(u16, &[u8])> {
     Ok((u16::from_be_bytes(*bytes), rest))
 }
 
-pub(crate) fn parse_cb_channel(buffer: &[u8]) -> Result<(ControlByte, u16, &[u8])> {
-    let Some((cb, rest)) = buffer.split_first() else {
+pub(crate) fn parse_cb_channel(packet_buffer: &[u8]) -> Result<(ControlByte, u16, &[u8])> {
+    let (cb, rest) = ControlByte::parse(packet_buffer)?;
+    let (channel_id, rest) = parse_u16(rest)?;
+    Ok((cb, channel_id, rest))
+}
+
+fn parse_length(packet_buffer: &[u8]) -> Result<(u16, &[u8])> {
+    let (payload_len, rest) = parse_u16(packet_buffer)?;
+    if payload_len > MAX_PAYLOAD_LEN {
+        log::error!("Payload length exceeds {}.", MAX_PAYLOAD_LEN);
+        return Err(Error::malformed_data());
+    }
+    if payload_len < CHECKSUM_LEN {
+        log::error!("Payload length is less than {}.", CHECKSUM_LEN);
+        return Err(Error::malformed_data());
+    }
+    Ok((payload_len, rest))
+}
+
+// To be used by `Mux::packet_in` to determine whether to process this packet or route it
+// to unicast channel.
+// - channel id is validated (can be unicast or broadcast), even for codec_v1
+// - only unicast encrypted transport and handshakes have Some length
+pub(crate) fn parse_channel_length(
+    cb: ControlByte,
+    packet_buffer: &[u8],
+) -> Result<(u16, Option<u16>)> {
+    let Some((_cb, rest)) = packet_buffer.split_first() else {
         log::error!("Packet is empty.");
         return Err(Error::malformed_data());
     };
     let (channel_id, rest) = parse_u16(rest)?;
-    Ok((ControlByte::try_from(*cb)?, channel_id, rest))
+    if !channel_id_valid(channel_id) {
+        log::error!("Invalid channel id {:04x}.", channel_id);
+        return Err(Error::malformed_data());
+    }
+    if channel_id != BROADCAST_CHANNEL_ID && (cb.is_encrypted_transport() || cb.is_handshake()) {
+        let (payload_len, _rest) = parse_length(rest)?;
+        Ok((channel_id, Some(payload_len)))
+    } else {
+        Ok((channel_id, None))
+    }
 }
 
 impl<R: Role> Header<R> {
@@ -107,15 +142,7 @@ impl<R: Role> Header<R> {
         if cb.is_continuation() {
             return Ok((Header::Continuation { channel_id }, rest));
         }
-        let (payload_len, rest) = parse_u16(rest)?;
-        if payload_len > MAX_PAYLOAD_LEN {
-            log::error!("Payload length exceeds {}.", MAX_PAYLOAD_LEN);
-            return Err(Error::malformed_data());
-        }
-        if payload_len < CHECKSUM_LEN {
-            log::error!("Payload length is less than {}.", CHECKSUM_LEN);
-            return Err(Error::malformed_data());
-        }
+        let (payload_len, rest) = parse_length(rest)?;
         // strip padding if there is any
         let without_padding = rest.len().min(payload_len.into());
         let rest = &rest[..without_padding];