refactor(rust/trezor-thp): separate handling for host-side broadcast messages

diff --git a/rust/trezor-thp/examples/host-cli/client.rs b/rust/trezor-thp/examples/host-cli/client.rs
index b7022c15..0c20d0a7 100644
--- a/rust/trezor-thp/examples/host-cli/client.rs
+++ b/rust/trezor-thp/examples/host-cli/client.rs
@@ -3,7 +3,7 @@ use std::net::{SocketAddr, UdpSocket};
 use std::time::Duration;
 
 use trezor_thp::{
-    Backend, ChannelIO, Error, channel::buffered::Buffered, channel::host::ChannelOpen,
+    Backend, ChannelIO, Error, channel::buffered::Buffered, channel::host::Mux,
     credential::CredentialStore,
 };
 
@@ -22,12 +22,12 @@ pub struct Client<C> {
     emu_addr: SocketAddr,
 }
 
-impl<C, B> Client<ChannelOpen<C, B>>
+impl<C, B> Client<Mux<C, B>>
 where
     B: Backend,
     C: CredentialStore,
 {
-    pub fn open(emu_addr: SocketAddr, channel: ChannelOpen<C, B>) -> Self {
+    pub fn open(emu_addr: SocketAddr, channel: Mux<C, B>) -> Self {
         let mut channel = Buffered::new(channel);
         channel.set_packet_len(PACKET_LEN);
         Client {
@@ -116,19 +116,18 @@ impl<C: ChannelIO> Client<C> {
 
     pub fn read(&mut self) -> (u8, u16, Vec<u8>) {
         let mut result: Option<(u8, u16, Vec<u8>)> = None;
+        let mut send_ack = true;
         while result.is_none() {
-            let mut message_ready = false;
-            while !message_ready {
+            let mut done = false;
+            while !done {
                 let Some(packet) = self.recv_from(READ_TIMEOUT) else {
                     log::error!("Timed out waiting for response for {:?}.", READ_TIMEOUT);
                     panic!();
                 };
-                message_ready = self
-                    .channel
-                    .packet_in(&packet)
-                    .check_failed()
-                    .unwrap()
-                    .got_message();
+                let pir = self.channel.packet_in(&packet).check_failed().unwrap();
+                assert!(!pir.got_transport_error());
+                done = pir.got_message() || pir.got_channel();
+                send_ack = !pir.got_channel();
             }
             result = match self.channel.message_out() {
                 Ok(r) => Some(r),
@@ -140,11 +139,14 @@ impl<C: ChannelIO> Client<C> {
                     log::error!("Cannot read message from channel: {:?}.", e);
                     panic!();
                 }
-            }
+            };
         }
         // Send ACK
-        let packet = self.write_ack();
-        self.send_to(&packet);
+        if send_ack {
+            log::trace!("Sending ACK.");
+            let packet = self.write_ack();
+            self.send_to(&packet);
+        }
         result.unwrap()
     }
 
diff --git a/rust/trezor-thp/examples/host-cli/main.rs b/rust/trezor-thp/examples/host-cli/main.rs
index ac495af7..0d77510b 100644
--- a/rust/trezor-thp/examples/host-cli/main.rs
+++ b/rust/trezor-thp/examples/host-cli/main.rs
@@ -7,7 +7,7 @@ use protobuf::Message;
 
 use trezor_thp::{
     Backend, Channel, Host,
-    channel::host::{ChannelOpen, ChannelPairing},
+    channel::host::{ChannelOpen, ChannelPairing, Mux},
     credential::{CredentialStore, NullCredentialStore},
 };
 
@@ -29,16 +29,21 @@ impl Backend for RustCrypto {
     }
 }
 
-fn do_handshake<C>(client: &mut Client<ChannelOpen<C, RustCrypto>>)
+fn do_allocation<C>(client: &mut Client<Mux<C, RustCrypto>>)
 where
     C: CredentialStore,
 {
-    // Handshake should finish within 3 request-response cycles.
-    // Device properties and channel id are available after the first one.
     client.call(0, &[]);
+}
+
+fn do_handshake<C>(client: &mut Client<ChannelOpen<C, RustCrypto>>)
+where
+    C: CredentialStore,
+{
     let device_properties =
         ThpDeviceProperties::parse_from_bytes(client.channel.device_properties()).unwrap();
     log::debug!("Device properties: {:?}.", device_properties);
+    // Handshake should finish within 2 request-response cycles.
     client.call(0, &[]);
     client.call(0, &[]);
 }
@@ -100,9 +105,14 @@ pub fn main() -> std::io::Result<()> {
     env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "info"));
 
     let cred_lookup = NullCredentialStore;
-    let channel = ChannelOpen::<_, RustCrypto>::new(false, cred_lookup).unwrap();
+    let mut channel = Mux::<_, RustCrypto>::new(cred_lookup);
+    channel.request_channel(false);
     let mut client = Client::open(get_address(), channel);
 
+    do_allocation(&mut client);
+    assert!(client.channel.channel_alloc_ready());
+    let mut client = client.map(|c| c.complete().unwrap());
+
     do_handshake(&mut client);
     assert!(client.channel.handshake_done());
     let mut client = client.map(|c| c.complete().unwrap());
diff --git a/rust/trezor-thp/src/channel/host.rs b/rust/trezor-thp/src/channel/host.rs
index ba2513f0..9ff716dc 100644
--- a/rust/trezor-thp/src/channel/host.rs
+++ b/rust/trezor-thp/src/channel/host.rs
@@ -4,10 +4,15 @@ use crate::{
     Backend, Channel, ChannelIO, Error, Host,
     channel::{ChannelState, Nonce, PacketInResult, PairingState, noise::NoiseHandshake},
     credential::CredentialStore,
-    header::{BROADCAST_CHANNEL_ID, HandshakeMessage, Header, parse_u16},
+    fragment::{Fragmenter, Reassembler},
+    header::{
+        BROADCAST_CHANNEL_ID, HandshakeMessage, Header, SyncBits, channel_id_valid,
+        parse_cb_channel, parse_u16,
+    },
+    util::prepare_zeroed,
 };
 
-use core::ops::ControlFlow;
+use core::marker::PhantomData;
 
 // Must fit any of:
 // - device_properties + overhead
@@ -17,10 +22,281 @@ const INTERNAL_BUFFER_LEN: usize = 192;
 const MAX_DEVICE_PROPERTIES_LEN: usize = 128;
 const MESSAGE_TYPE_END_RESPONSE: u16 = 1019; // ThpMessageType_ThpEndResponse
 
-#[derive(Copy, Clone)]
-enum HostHandshakeState {
+enum AllocationState {
+    None,
+    SendingRequest {
+        try_to_unlock: bool,
+    },
     /// `HH0`.
-    SentChannelRequest(Nonce),
+    SentRequest {
+        try_to_unlock: bool,
+        nonce: Nonce,
+    },
+    ReceivingResponse {
+        try_to_unlock: bool,
+        nonce: Nonce,
+        reassembler: Reassembler<Host>,
+    },
+    ReceivedId {
+        try_to_unlock: bool,
+        channel_id: u16,
+    },
+}
+
+/// Handles broadcast channel messages, notably channel allocation requests.
+/// Because host often only needs a single channel, you can throw away the Mux
+/// after allocating one, if you don't need the keep-alive functionality.
+pub struct Mux<C, B> {
+    cred_store: C,
+    internal_buffer: heapless::Vec<u8, MAX_DEVICE_PROPERTIES_LEN>,
+    channel_allocation: AllocationState,
+    ping: Option<(bool, Nonce)>,
+    _phantom: PhantomData<B>,
+}
+
+impl<C, B> Mux<C, B>
+where
+    C: CredentialStore,
+    B: Backend,
+{
+    pub fn new(cred_store: C) -> Self {
+        let mut internal_buffer = heapless::Vec::new();
+        prepare_zeroed(&mut internal_buffer);
+        Self {
+            cred_store,
+            internal_buffer,
+            channel_allocation: AllocationState::None,
+            ping: None,
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Enqueue a keep-alive message.
+    pub fn ping(&mut self) {
+        if self.ping.is_some() {
+            log::warn!("Dropping previous ping attempt.");
+        }
+        self.ping = Some((false, Nonce::random::<B>()));
+    }
+
+    /// Enqueue channel allocation request.
+    pub fn request_channel(&mut self, try_to_unlock: bool) {
+        if !matches!(self.channel_allocation, AllocationState::None) {
+            log::warn!("Abandoned previous channel allocation request.");
+        }
+        self.channel_allocation = AllocationState::SendingRequest { try_to_unlock };
+    }
+
+    /// Create new [`ChannelOpen`] after channel allocation response was received.
+    pub fn channel_alloc(&mut self) -> Result<ChannelOpen<C, B>, Error> {
+        let AllocationState::ReceivedId {
+            try_to_unlock,
+            channel_id,
+        } = self.channel_allocation
+        else {
+            return Err(Error::not_ready());
+        };
+        let ch = ChannelOpen::new(
+            channel_id,
+            self.cred_store.clone(),
+            &self.internal_buffer,
+            try_to_unlock,
+        )?;
+        self.channel_allocation = AllocationState::None;
+        prepare_zeroed(&mut self.internal_buffer);
+        Ok(ch)
+    }
+
+    /// True if [`Mux::channel_alloc`] can be called.
+    pub fn channel_alloc_ready(&self) -> bool {
+        matches!(self.channel_allocation, AllocationState::ReceivedId { .. })
+    }
+
+    /// Same as [`Mux::channel_alloc`] but destroys the [`Mux`],
+    /// like `complete()` does for other types.
+    pub fn complete(mut self) -> Result<ChannelOpen<C, B>, Error> {
+        self.channel_alloc()
+    }
+
+    fn handle_broadcast(&mut self, packet: &[u8]) -> Result<PacketInResult, Error> {
+        let (header, _rest) = Header::<Host>::parse(packet)?;
+        match header {
+            Header::Pong => {
+                let (_header, payload) = Reassembler::<Host>::single_inplace(packet)?;
+                let (nonce, _rest) = Nonce::parse(payload)?;
+                if Some((true, nonce)) != self.ping {
+                    log::warn!("Ignoring PONG with invalid nonce.");
+                    return Err(Error::malformed_data());
+                }
+                self.ping = None;
+                Ok(PacketInResult::pong())
+            }
+            Header::ChannelAllocationResponse { .. } => {
+                let (try_to_unlock, nonce) = match self.channel_allocation {
+                    AllocationState::SentRequest {
+                        try_to_unlock,
+                        nonce,
+                        ..
+                    } => (try_to_unlock, nonce),
+                    AllocationState::ReceivingResponse {
+                        try_to_unlock,
+                        nonce,
+                        ..
+                    } => (try_to_unlock, nonce),
+                    _ => return Err(Error::malformed_data()),
+                };
+                self.internal_buffer.fill(0u8);
+                let reassembler = Reassembler::<Host>::new(packet, &mut self.internal_buffer)?;
+                self.channel_allocation = AllocationState::ReceivingResponse {
+                    try_to_unlock,
+                    nonce,
+                    reassembler,
+                };
+                self.handle_allocation_response()
+            }
+            Header::Continuation {
+                channel_id: BROADCAST_CHANNEL_ID,
+            } => {
+                let AllocationState::ReceivingResponse { reassembler, .. } =
+                    &mut self.channel_allocation
+                else {
+                    return Err(Error::malformed_data());
+                };
+                reassembler.update(packet, &mut self.internal_buffer)?;
+                self.handle_allocation_response()
+            }
+            // No Header::TransportError for broadcast channel.
+            _ => {
+                log::debug!(
+                    "Broadcast channel: ignoring packet with control byte {}.",
+                    packet[0]
+                );
+                Err(Error::malformed_data())
+            }
+        }
+    }
+
+    fn handle_allocation_response(&mut self) -> Result<PacketInResult, Error> {
+        let AllocationState::ReceivingResponse {
+            try_to_unlock,
+            nonce,
+            reassembler,
+        } = &mut self.channel_allocation
+        else {
+            return Ok(PacketInResult::accept(false));
+        };
+        if !reassembler.is_done() {
+            return Ok(PacketInResult::accept(false));
+        }
+        let len = reassembler.verify(self.internal_buffer.as_slice())?;
+        let payload = &self.internal_buffer[..len];
+
+        let (received_nonce, payload) = Nonce::parse(payload)?;
+        if received_nonce != *nonce {
+            log::warn!("Received non matching channel request nonce.");
+            return Ok(PacketInResult::accept(false));
+        };
+        let (channel_id, device_properties) = parse_u16(payload)?;
+        let device_properties_len = device_properties.len();
+        // Shift internal buffer to only contain device_properties.
+        self.internal_buffer.copy_within(Nonce::LEN + 2..len, 0);
+        self.internal_buffer.truncate(device_properties_len);
+        self.channel_allocation = AllocationState::ReceivedId {
+            try_to_unlock: *try_to_unlock,
+            channel_id,
+        };
+        log::debug!("Got channel id {}.", channel_id);
+        Ok(PacketInResult::channel_allocation(channel_id))
+    }
+}
+
+impl<C, B> ChannelIO for Mux<C, B>
+where
+    C: CredentialStore,
+    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
+            return PacketInResult::ignore(Error::malformed_data());
+        };
+        if !channel_id_valid(channel_id) {
+            log::warn!("Invalid channel id {}.", channel_id);
+            return PacketInResult::ignore(Error::malformed_data());
+        }
+        if channel_id != BROADCAST_CHANNEL_ID {
+            return PacketInResult::route(channel_id);
+        }
+        PacketInResult::from_result(self.handle_broadcast(packet_buffer))
+    }
+
+    fn packet_in_ready(&self) -> bool {
+        true
+    }
+
+    fn packet_out(&mut self, packet_buffer: &mut [u8], _send_buffer: &[u8]) -> Result<(), Error> {
+        if let AllocationState::SendingRequest { try_to_unlock } = self.channel_allocation {
+            let nonce = Nonce::random::<B>();
+            Fragmenter::<Host>::single(
+                Header::new_channel_request(),
+                SyncBits::new(),
+                nonce.as_slice(),
+                packet_buffer,
+            )?;
+            self.channel_allocation = AllocationState::SentRequest {
+                try_to_unlock,
+                nonce,
+            };
+        } else if let Some((false, _)) = self.ping {
+            let nonce = Nonce::random::<B>();
+            Fragmenter::<Host>::single(
+                Header::new_ping(),
+                SyncBits::new(),
+                nonce.as_slice(),
+                packet_buffer,
+            )?;
+            self.ping = Some((true, nonce));
+        } else {
+            return Err(Error::not_ready());
+        }
+        Ok(())
+    }
+
+    fn packet_out_ready(&self) -> bool {
+        let send_ping = matches!(self.ping, Some((false, _)));
+        let send_channel_allocation = matches!(
+            self.channel_allocation,
+            AllocationState::SendingRequest { .. }
+        );
+        send_ping || send_channel_allocation
+    }
+
+    fn message_in(&mut self, _plaintext_len: usize, _send_buffer: &mut [u8]) -> Result<(), Error> {
+        Ok(())
+    }
+
+    fn message_in_ready(&self) -> bool {
+        true
+    }
+
+    fn message_out<'a>(
+        &mut self,
+        receive_buffer: &'a mut [u8],
+    ) -> Result<(u8, u16, &'a [u8]), Error> {
+        Ok((0, 0, &receive_buffer[..0]))
+    }
+
+    fn message_out_ready(&self) -> bool {
+        false
+    }
+
+    fn message_retransmit(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+#[derive(Copy, Clone)]
+enum HandshakeState {
     /// `HH1`.
     SentInitiationRequest,
     /// `HH2`.
@@ -33,70 +309,67 @@ enum HostHandshakeState {
 
 /// Open a [`Channel`] from the host side.
 ///
-/// - start by calling [`ChannelOpen::new`]
+/// - start by calling [`Mux::request_channel`]
+/// - perform [`ChannelIO`] with empty messages until [`PacketInResult::ChannelAllocation`] is returned
+/// - calling [`Mux::channel_alloc`] to obtain [`ChannelOpen`]
 /// - perform [`ChannelIO`] with empty messages until [`ChannelOpen::handshake_done`]
 /// - call [`ChannelOpen::complete`] to obtain [`ChannelPairing`]
 /// - perform [`ChannelIO`] until [`ChannelPairing::pairing_done`]
 /// - call [`ChannelPairing::complete`] to get established [`Channel`]
 pub struct ChannelOpen<C: CredentialStore, B: Backend> {
     channel: Channel<Host, B>,
-    state: HostHandshakeState,
-    noise: Option<NoiseHandshake<Host, B>>,
+    state: HandshakeState,
+    noise: NoiseHandshake<Host, B>,
     internal_buffer: heapless::Vec<u8, INTERNAL_BUFFER_LEN>,
     device_properties: heapless::Vec<u8, MAX_DEVICE_PROPERTIES_LEN>,
     cred_store: C,
-    try_to_unlock: bool,
 }
 
 impl<C: CredentialStore, B: Backend> ChannelOpen<C, B> {
-    pub fn new(try_to_unlock: bool, cred_store: C) -> Result<Self, Error> {
-        let nonce = Nonce::random::<B>();
-        let internal_buffer = heapless::Vec::from_slice(nonce.as_slice()).unwrap();
-        let mut channel = Channel::new(BROADCAST_CHANNEL_ID);
-        channel.raw_in(Header::new_channel_request(), &internal_buffer)?;
+    fn new(
+        channel_id: u16,
+        cred_store: C,
+        device_properties: &[u8],
+        try_to_unlock: bool,
+    ) -> Result<Self, Error> {
+        let device_properties = heapless::Vec::from_slice(device_properties)
+            .map_err(|_| Error::insufficient_buffer())?;
+
+        let mut internal_buffer = heapless::Vec::new();
+        prepare_zeroed(&mut internal_buffer);
+        let (hss, msg) = NoiseHandshake::initiation_request(
+            &device_properties,
+            try_to_unlock,
+            &mut internal_buffer,
+        )?;
+        let header = Header::new_handshake(channel_id, HandshakeMessage::InitiationRequest, msg)?;
+        let mut channel = Channel::new(channel_id);
+        channel.raw_in(header, msg)?;
+        let len = msg.len();
+        internal_buffer.truncate(len);
         let res = Self {
             channel,
-            state: HostHandshakeState::SentChannelRequest(nonce),
-            noise: None,
+            state: HandshakeState::SentInitiationRequest,
+            noise: hss,
             internal_buffer,
-            device_properties: heapless::Vec::new(),
+            device_properties,
             cred_store,
-            try_to_unlock,
         };
         Ok(res)
     }
 
-    pub fn is_broadcast(&self) -> bool {
-        self.channel.is_broadcast()
-    }
-
     fn incoming_internal(&mut self) -> Result<(), Error> {
         let (header, len) = self.channel.raw_out(&self.internal_buffer)?;
         self.internal_buffer.truncate(len);
 
         match (self.state, header.handshake_phase()) {
-            (HostHandshakeState::SentChannelRequest(nonce), _)
-                if header.is_channel_allocation_response() =>
-            {
-                if self.get_channel(&nonce)?.is_continue() {
-                    log::debug!("Got channel id {}.", self.channel.channel_id);
-                    self.start_handshake()?;
-                    self.state = HostHandshakeState::SentInitiationRequest;
-                }
-            }
-            (
-                HostHandshakeState::SentInitiationRequest,
-                Some(HandshakeMessage::InitiationResponse),
-            ) => {
+            (HandshakeState::SentInitiationRequest, Some(HandshakeMessage::InitiationResponse)) => {
                 self.continue_handshake()?;
-                self.state = HostHandshakeState::SentCompletionRequest;
+                self.state = HandshakeState::SentCompletionRequest;
             }
-            (
-                HostHandshakeState::SentCompletionRequest,
-                Some(HandshakeMessage::CompletionResponse),
-            ) => {
+            (HandshakeState::SentCompletionRequest, Some(HandshakeMessage::CompletionResponse)) => {
                 let device_state = self.finish_handshake()?;
-                self.state = HostHandshakeState::Finished(device_state);
+                self.state = HandshakeState::Finished(device_state);
             }
             _ => {
                 log::error!("Unexpected handshake state.");
@@ -106,49 +379,18 @@ impl<C: CredentialStore, B: Backend> ChannelOpen<C, B> {
         Ok(())
     }
 
-    fn get_channel(&mut self, expected_nonce: &Nonce) -> Result<ControlFlow<(), ()>, Error> {
-        let (nonce, payload) = Nonce::parse(&self.internal_buffer)?;
-        if nonce != *expected_nonce {
-            log::warn!("Received non matching channel request nonce.");
-            return Ok(ControlFlow::Break(()));
-        };
-        let (cid, device_properties) = parse_u16(payload)?;
-        self.channel.channel_id = cid;
-        self.device_properties = heapless::Vec::from_slice(device_properties)
-            .map_err(|_| Error::insufficient_buffer())?;
-        Ok(ControlFlow::Continue(()))
-    }
-
-    fn start_handshake(&mut self) -> Result<(), Error> {
-        self.zero_internal_buffer();
-        let (hss, msg) = NoiseHandshake::start_pairing(
-            &self.device_properties,
-            self.try_to_unlock,
-            &mut self.internal_buffer,
-        )?;
-        let header = Header::new_handshake(
-            self.channel.channel_id,
-            HandshakeMessage::InitiationRequest,
-            msg,
-        )?;
-        self.noise = Some(hss);
-        self.channel.raw_in(header, msg)?;
-        let len = msg.len();
-        self.internal_buffer.truncate(len);
-        Ok(())
-    }
-
     fn continue_handshake(&mut self) -> Result<(), Error> {
         let payload_len = self.internal_buffer.len();
         // Buffer used both for input and output - pad with zeros.
         self.internal_buffer
             .resize(self.internal_buffer.capacity(), 0u8)
             .unwrap();
-        let noise = self.noise.as_mut().ok_or_else(Error::unexpected_input)?;
-        let (nc, msg) =
-            noise.complete_pairing(&mut self.cred_store, &mut self.internal_buffer, payload_len)?;
+        let (nc, msg) = self.noise.completion_request(
+            &mut self.cred_store,
+            &mut self.internal_buffer,
+            payload_len,
+        )?;
         self.channel.noise = Some(nc);
-        self.noise = None;
         let header = Header::new_handshake(
             self.channel.channel_id,
             HandshakeMessage::CompletionRequest,
@@ -167,35 +409,28 @@ impl<C: CredentialStore, B: Backend> ChannelOpen<C, B> {
         PairingState::try_from(payload.as_slice())
     }
 
-    fn zero_internal_buffer(&mut self) {
-        self.internal_buffer.clear();
-        self.internal_buffer
-            .resize(self.internal_buffer.capacity(), 0u8)
-            .unwrap();
-    }
-
     pub fn device_properties(&self) -> &[u8] {
         self.device_properties.as_slice()
     }
 
     /// True if handshake finished and [`ChannelOpen::complete()`] can be called.
     pub fn handshake_done(&self) -> bool {
-        matches!(self.state, HostHandshakeState::Finished(_))
+        matches!(self.state, HandshakeState::Finished(_))
     }
 
     /// True if the handshake failed and the object should be discarded.
     pub fn handshake_failed(&self) -> bool {
-        matches!(self.state, HostHandshakeState::Failed)
+        matches!(self.state, HandshakeState::Failed)
     }
 
     /// Transition into the pairing phase.
     pub fn complete(self) -> Result<ChannelPairing<B>, Error> {
-        if self.is_broadcast() || self.channel.noise.is_none() {
+        if self.channel.noise.is_none() {
             return Err(Error::unexpected_input());
         }
         log::debug!("Handshake complete.");
         Ok(match self.state {
-            HostHandshakeState::Finished(ps) => ChannelPairing {
+            HandshakeState::Finished(ps) => ChannelPairing {
                 channel: self.channel,
                 device_properties: self.device_properties,
                 pairing_state: ps,
@@ -230,13 +465,13 @@ where
             return PacketInResult::ignore(Error::MalformedData);
         }
         if res.got_ack() {
-            self.zero_internal_buffer();
+            prepare_zeroed(&mut self.internal_buffer);
         }
         if res.got_message() {
             let handled = self.incoming_internal();
             if let Err(e) = handled {
                 if e != Error::InvalidChecksum {
-                    self.state = HostHandshakeState::Failed;
+                    self.state = HandshakeState::Failed;
                     return PacketInResult::fail(e);
                 }
             }
diff --git a/rust/trezor-thp/src/channel/mod.rs b/rust/trezor-thp/src/channel/mod.rs
index e37e11f7..3ddee90c 100644
--- a/rust/trezor-thp/src/channel/mod.rs
+++ b/rust/trezor-thp/src/channel/mod.rs
@@ -86,8 +86,8 @@ enum ChannelState<R: Role> {
 
 /// THP channel with established secure layer.
 ///
-/// There is no constructor, to obtain a channel please use [`host::ChannelOpen`]
-/// or `device::ChannelOpen`.
+/// There is no constructor, to obtain a channel please use [`host::Mux`]
+/// or [`device::Mux`].
 /// For actually sending and receiving messages please see [`ChannelIO`].
 pub struct Channel<R: Role, B: Backend> {
     channel_id: u16,
@@ -112,10 +112,6 @@ impl<R: Role, B: Backend> Channel<R, B> {
         self.noise.as_mut().ok_or_else(Error::unexpected_input)
     }
 
-    fn is_broadcast(&self) -> bool {
-        self.channel_id == BROADCAST_CHANNEL_ID
-    }
-
     pub fn channel_id(&self) -> u16 {
         self.channel_id
     }
@@ -128,18 +124,13 @@ impl<R: Role, B: Backend> Channel<R, B> {
         let ChannelState::Idle = self.state else {
             return Err(Error::not_ready());
         };
-        let sb = if self.is_broadcast() {
-            SyncBits::new()
-        } else {
-            self.sync.send_start().ok_or_else(Error::not_ready)?
-        };
+        let sb = self.sync.send_start().ok_or_else(Error::not_ready)?;
         let frag = Fragmenter::new(header, sb, send_buffer)?;
         self.state = ChannelState::Sending(frag);
         Ok(())
     }
 
     fn raw_out(&mut self, receive_buffer: &[u8]) -> Result<(Header<R>, usize)> {
-        let has_cid = !self.is_broadcast();
         let ChannelState::Receiving(r) = &mut self.state else {
             return Err(Error::not_ready());
         };
@@ -156,10 +147,8 @@ impl<R: Role, B: Backend> Channel<R, B> {
                 return Err(e);
             }
         };
-        if has_cid {
-            self.send_ack = Some(self.sync.receive_acknowledge());
-        }
-        let header = r.header();
+        self.send_ack = Some(self.sync.receive_acknowledge());
+        let header = r.header().clone();
         self.state = ChannelState::Idle;
         Ok((header, len))
     }
@@ -264,7 +253,7 @@ impl<R: Role, B: Backend> Channel<R, B> {
         receive_buffer: &mut [u8],
     ) -> Result<(bool, Option<u16>)> {
         let sb = SyncBits::try_from(packet_buffer)?;
-        if !self.is_broadcast() && !self.sync.receive_start(sb) {
+        if !self.sync.receive_start(sb) {
             // Bad sync bit, drop this packet and continuations.
             log::debug!("[{}] Bad sync bit, ignoring packet.", self.channel_id);
             self.state = ChannelState::Idle;
@@ -579,8 +568,10 @@ impl<R: Role, B: Backend> ChannelIO for Channel<R, B> {
             return Err(Error::not_ready());
         }
         if f.is_done() {
-            if self.is_broadcast() {
-                // No ACKs without channel ID, assume delivered.
+            if f.header().channel_id() == BROADCAST_CHANNEL_ID {
+                // This is a special case for `channel_allocation_response` which is the only
+                // message sent through Channel (by `device::ChannelOpen`) but does not
+                // wait for ACK because as it is sent on broadcast channel.
                 self.state = ChannelState::Idle;
             } else {
                 self.sync.send_finish();
diff --git a/rust/trezor-thp/src/channel/noise.rs b/rust/trezor-thp/src/channel/noise.rs
index b78bf556..b015b751 100644
--- a/rust/trezor-thp/src/channel/noise.rs
+++ b/rust/trezor-thp/src/channel/noise.rs
@@ -2,7 +2,7 @@ use trezor_noise_protocol::{
     Cipher, CipherState, DH, HandshakeState, Hash, U8Array, patterns::noise_xx,
 };
 
-use crate::{Error, Host, Role, credential::CredentialStore};
+use crate::{Error, Host, Role, credential::CredentialStore, util::prepare_zeroed};
 
 use core::marker::PhantomData;
 
@@ -138,7 +138,7 @@ impl<B: Backend> NoiseHandshake<Host, B> {
         DHPrivKey<B>: U8Array,
     {
         let mut buf = heapless::Vec::new();
-        buf.resize(buf.capacity(), 0u8).unwrap();
+        prepare_zeroed(&mut buf);
         let result = cs.lookup(re.as_slice(), rs.as_slice(), buf.as_mut_slice());
         if let Some(found) = result {
             let found_key = <DHPrivKey<B> as U8Array>::from_slice(found.local_static_privkey);
diff --git a/rust/trezor-thp/src/credential.rs b/rust/trezor-thp/src/credential.rs
index f690dea4..86fbfb0c 100644
--- a/rust/trezor-thp/src/credential.rs
+++ b/rust/trezor-thp/src/credential.rs
@@ -7,7 +7,7 @@ pub struct FoundCredential<'a> {
 
 /// Host-side credential store.
 /// Basically a set of (`remote_static_pubkey`, `local_static_privkey`, `auth_credential`).
-pub trait CredentialStore {
+pub trait CredentialStore: Clone {
     /// Find a record such that `remote_static_pubkey` satisfies
     /// `masked_static_pubkey == X25519(SHA-256(remote_static_pubkey || ephemeral_pubkey), remote_static_pubkey)`.
     /// If found, write `local_static_privkey` and `auth_credential` to `dest` and return the written subslices.
@@ -20,6 +20,7 @@ pub trait CredentialStore {
 }
 
 /// Never finds a matching credential.
+#[derive(Clone)]
 pub struct NullCredentialStore;
 
 impl CredentialStore for NullCredentialStore {
diff --git a/rust/trezor-thp/src/fragment.rs b/rust/trezor-thp/src/fragment.rs
index d5823fb1..cba52045 100644
--- a/rust/trezor-thp/src/fragment.rs
+++ b/rust/trezor-thp/src/fragment.rs
@@ -101,6 +101,10 @@ impl<R: Role> Fragmenter<R> {
         }
         Ok(())
     }
+
+    pub fn header(&self) -> &Header<R> {
+        &self.header
+    }
 }
 
 pub struct Reassembler<R: Role> {
@@ -194,8 +198,8 @@ impl<R: Role> Reassembler<R> {
         Ok(length_no_checksum)
     }
 
-    pub fn header(&self) -> Header<R> {
-        self.header.clone()
+    pub fn header(&self) -> &Header<R> {
+        &self.header
     }
 
     // Shortcut to deserialize single packet message.
@@ -209,6 +213,32 @@ impl<R: Role> Reassembler<R> {
         let header = reassembler.header;
         Ok((header, &dest[..reply_len]))
     }
+
+    pub fn single_inplace(buffer: &[u8]) -> Result<(Header<R>, &[u8])> {
+        let (header, after_header) = Header::parse(buffer)?;
+        if header.is_continuation() {
+            return Err(Error::unexpected_input());
+        }
+        let payload_len: usize = header.payload_len().into();
+        if payload_len != after_header.len() {
+            log::error!("Single packet message expected.");
+            return Err(Error::malformed_data());
+        }
+
+        let mut checksum = Crc32::new();
+        checksum.update(&buffer[..header.header_len()]);
+        let checksum_off = payload_len.saturating_sub(CHECKSUM_LEN);
+        checksum.update(&after_header[..checksum_off]);
+        let computed_checksum = checksum.finalize();
+
+        let received_checksum = *after_header
+            .last_chunk::<CHECKSUM_LEN>()
+            .ok_or_else(Error::invalid_checksum)?;
+        if computed_checksum != received_checksum {
+            return Err(Error::invalid_checksum());
+        }
+        Ok((header, &after_header[..checksum_off]))
+    }
 }
 
 #[cfg(test)]
@@ -276,6 +306,59 @@ mod test {
         }
     }
 
+    #[test]
+    fn test_reassemble_single_roundtrip() {
+        const DATA: &'static [u8] = b"The Quick Brown Fox Jumps Over the Lazy Dog The Quick Brown Fox Jumps Over the Lazy Dog";
+
+        for i in 0..DATA.len() {
+            let source = &DATA[..i];
+            let channel_id = 1 + i as u16;
+            let header = Header::<Host>::new_encrypted(channel_id, source).unwrap();
+            let header_len = header.header_len();
+
+            let packet_size = header_len + i + CHECKSUM_LEN;
+            let packets = fragment(header, SyncBits::new(), source, packet_size);
+            assert_eq!(packets.len(), 1);
+
+            let res = Reassembler::<Device>::single_inplace(&packets[0]).unwrap();
+            let expected_hex = hex::encode(source);
+            let received_hex = hex::encode(res.1);
+            assert_eq!(received_hex, expected_hex);
+        }
+    }
+
+    #[test]
+    fn test_reassemble_single_good() {
+        let empty = hex::decode(EMPTY_PAYLOAD_EXPECTED).unwrap();
+        let res = Reassembler::<Device>::single_inplace(&empty).unwrap();
+        assert_eq!(res.1, b"");
+        let res = Reassembler::<Host>::single_inplace(&empty).unwrap();
+        assert_eq!(res.1, b"");
+
+        let short = hex::decode(SHORT_PAYLOAD_EXPECTED).unwrap();
+        let res = Reassembler::<Device>::single_inplace(&short).unwrap();
+        assert_eq!(res.1, b"\x07");
+        let res = Reassembler::<Host>::single_inplace(&short).unwrap();
+        assert_eq!(res.1, b"\x07");
+    }
+
+    #[test]
+    fn test_reassemble_single_bad() {
+        // failure expected when more data follows
+        let incomplete = hex::decode(LONGER_PAYLOAD_EXPECTED[0]).unwrap();
+        let res = Reassembler::<Device>::single_inplace(&incomplete);
+        assert_eq!(res, Err(Error::MalformedData));
+        let res = Reassembler::<Host>::single_inplace(&incomplete);
+        assert_eq!(res, Err(Error::MalformedData));
+
+        // failure expected on continuations
+        let continuation = hex::decode(LONGER_PAYLOAD_EXPECTED[1]).unwrap();
+        let res = Reassembler::<Device>::single_inplace(&continuation);
+        assert_eq!(res, Err(Error::MalformedData));
+        let res = Reassembler::<Host>::single_inplace(&continuation);
+        assert_eq!(res, Err(Error::MalformedData));
+    }
+
     // follwing vectors and tests adapted from test_trezor.wire.thp.writer.py
 
     const EMPTY_PAYLOAD_EXPECTED: &str = "0412340004edbd479c00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000";
@@ -376,5 +459,8 @@ mod test {
         let mut received = [0u8; MAX_MESSAGE];
         let reassembler = Reassembler::<Device>::new(packet, &mut received);
         assert!(matches!(reassembler, Err(Error::MalformedData)));
+
+        let inplace = Reassembler::<Device>::single_inplace(packet);
+        assert!(matches!(inplace, Err(Error::MalformedData)));
     }
 }
diff --git a/rust/trezor-thp/src/lib.rs b/rust/trezor-thp/src/lib.rs
index d1ff02b4..b998fdf6 100644
--- a/rust/trezor-thp/src/lib.rs
+++ b/rust/trezor-thp/src/lib.rs
@@ -10,6 +10,7 @@ pub mod credential;
 pub mod error;
 pub mod fragment;
 pub mod header;
+mod util;
 
 pub use channel::{Backend, Channel, ChannelIO};
 pub use error::Error;
diff --git a/rust/trezor-thp/src/util.rs b/rust/trezor-thp/src/util.rs
new file mode 100644
index 00000000..97253f9f
--- /dev/null
+++ b/rust/trezor-thp/src/util.rs
@@ -0,0 +1,4 @@
+pub(crate) fn prepare_zeroed<const C: usize>(buf: &mut heapless::Vec<u8, C>) {
+    buf.clear();
+    buf.resize(buf.capacity(), 0u8).unwrap();
+}