refactor(rust/trezor-thp): always validate control byte

diff --git a/rust/trezor-thp/src/channel/mod.rs b/rust/trezor-thp/src/channel/mod.rs
index a4cf8008..e4b928d4 100644
--- a/rust/trezor-thp/src/channel/mod.rs
+++ b/rust/trezor-thp/src/channel/mod.rs
@@ -180,14 +180,10 @@ impl<R: Role, B: Backend> Channel<R, B> {
             return Ok(PacketInResult::transport_error(te));
         }
         let is_cont = cb.is_continuation();
-        if !(is_cont
-            || cb.is_channel_allocation_request()
-            || cb.is_channel_allocation_response()
-            || cb.is_handshake()
-            || cb.is_encrypted_transport())
-        {
+        if !(is_cont || cb.is_handshake() || cb.is_encrypted_transport()) {
+            // Channel allocation and codec v1 are handled by Mux.
             log::warn!(
-                "[{}] Invalid control byte {}.",
+                "[{}] Unexpected control byte {}.",
                 self.channel_id,
                 u8::from(cb)
             );
diff --git a/rust/trezor-thp/src/control_byte.rs b/rust/trezor-thp/src/control_byte.rs
index bcd8017c..e4edf0d1 100644
--- a/rust/trezor-thp/src/control_byte.rs
+++ b/rust/trezor-thp/src/control_byte.rs
@@ -1,25 +1,24 @@
-use crate::{alternating_bit::SyncBits, error::Error};
+use crate::{alternating_bit::SyncBits, error::Error, header::HandshakeMessage};
 
-pub const CODEC_V1: u8 = 0x3F;
-pub const CONTINUATION_PACKET: u8 = 0x80;
+const CODEC_V1: u8 = 0x3F;
+const CONTINUATION_PACKET: u8 = 0x80;
 
-pub const CONTINUATION_PACKET_MASK: u8 = 0x80;
+const CONTINUATION_PACKET_MASK: u8 = 0x80;
 
-pub const ACK_MASK: u8 = 0xF7;
-pub const ACK_MESSAGE: u8 = 0x20;
+const ACK_MASK: u8 = 0xF7;
+const ACK_MESSAGE: u8 = 0x20;
+const DATA_MASK: u8 = 0xE7;
+const HANDSHAKE_INIT_REQ: u8 = 0x00;
+const HANDSHAKE_INIT_RES: u8 = 0x01;
+const HANDSHAKE_COMP_REQ: u8 = 0x02;
+const HANDSHAKE_COMP_RES: u8 = 0x03;
+const ENCRYPTED_TRANSPORT: u8 = 0x04;
 
-pub const DATA_MASK: u8 = 0xE7;
-pub const HANDSHAKE_INIT_REQ: u8 = 0x00;
-pub const HANDSHAKE_INIT_RES: u8 = 0x01;
-pub const HANDSHAKE_COMP_REQ: u8 = 0x02;
-pub const HANDSHAKE_COMP_RES: u8 = 0x03;
-pub const ENCRYPTED_TRANSPORT: u8 = 0x04;
-
-pub const CHANNEL_ALLOCATION_REQ: u8 = 0x40;
-pub const CHANNEL_ALLOCATION_RES: u8 = 0x41;
-pub const ERROR: u8 = 0x42;
-pub const PING: u8 = 0x43;
-pub const PONG: u8 = 0x44;
+const CHANNEL_ALLOCATION_REQ: u8 = 0x40;
+const CHANNEL_ALLOCATION_RES: u8 = 0x41;
+const ERROR: u8 = 0x42;
+const PING: u8 = 0x43;
+const PONG: u8 = 0x44;
 
 pub const ACK_BIT: u8 = 0x08;
 pub const SEQ_BIT: u8 = 0x10;
@@ -29,12 +28,49 @@ pub const SYNC_MASK: u8 = ACK_BIT | SEQ_BIT;
 pub struct ControlByte(u8);
 
 impl ControlByte {
-    pub fn sync_bits(&self) -> SyncBits {
-        SyncBits::from(self.0 & SYNC_MASK)
+    pub const fn codec_v1() -> Self {
+        Self(CODEC_V1)
     }
 
-    pub fn with_sync_bits(self, sb: SyncBits) -> Self {
-        Self(self.0 & !SYNC_MASK | <SyncBits as Into<u8>>::into(sb))
+    pub const fn continuation() -> Self {
+        Self(CONTINUATION_PACKET)
+    }
+
+    pub const fn ack() -> Self {
+        Self(ACK_MESSAGE)
+    }
+
+    pub const fn encrypted_transport() -> Self {
+        Self(ENCRYPTED_TRANSPORT)
+    }
+
+    pub const fn ping() -> Self {
+        Self(PING)
+    }
+
+    pub const fn pong() -> Self {
+        Self(PONG)
+    }
+
+    pub const fn error() -> Self {
+        Self(ERROR)
+    }
+
+    pub const fn channel_allocation_request() -> Self {
+        Self(CHANNEL_ALLOCATION_REQ)
+    }
+
+    pub const fn channel_allocation_response() -> Self {
+        Self(CHANNEL_ALLOCATION_RES)
+    }
+
+    pub const fn handshake(phase: HandshakeMessage) -> Self {
+        match phase {
+            HandshakeMessage::InitiationRequest => Self(HANDSHAKE_INIT_REQ),
+            HandshakeMessage::InitiationResponse => Self(HANDSHAKE_INIT_RES),
+            HandshakeMessage::CompletionRequest => Self(HANDSHAKE_COMP_REQ),
+            HandshakeMessage::CompletionResponse => Self(HANDSHAKE_COMP_RES),
+        }
     }
 
     pub const fn is_ack(&self) -> bool {
@@ -74,26 +110,49 @@ impl ControlByte {
     }
 
     pub const fn is_handshake(&self) -> bool {
-        self.0 & DATA_MASK == HANDSHAKE_INIT_REQ
-            || self.0 & DATA_MASK == HANDSHAKE_INIT_RES
-            || self.0 & DATA_MASK == HANDSHAKE_COMP_REQ
-            || self.0 & DATA_MASK == HANDSHAKE_COMP_RES
+        self.handshake_phase().is_some()
     }
-}
 
-// Note consider TryFrom + validation
-impl From<u8> for ControlByte {
-    fn from(byte: u8) -> Self {
-        Self(byte)
+    pub const fn handshake_phase(&self) -> Option<HandshakeMessage> {
+        let res = match self.0 & DATA_MASK {
+            HANDSHAKE_INIT_REQ => HandshakeMessage::InitiationRequest,
+            HANDSHAKE_INIT_RES => HandshakeMessage::InitiationResponse,
+            HANDSHAKE_COMP_REQ => HandshakeMessage::CompletionRequest,
+            HANDSHAKE_COMP_RES => HandshakeMessage::CompletionResponse,
+            _ => return None,
+        };
+        Some(res)
+    }
+
+    pub fn sync_bits(&self) -> SyncBits {
+        SyncBits::from(self.0 & SYNC_MASK)
+    }
+
+    pub fn with_sync_bits(self, sb: SyncBits) -> Self {
+        Self(self.0 & !SYNC_MASK | <SyncBits as Into<u8>>::into(sb))
     }
 }
 
-impl TryFrom<&[u8]> for ControlByte {
+impl TryFrom<u8> for ControlByte {
     type Error = Error;
 
-    fn try_from(bytes: &[u8]) -> Result<Self, Error> {
-        let first_byte = bytes.first().ok_or_else(Error::malformed_data)?;
-        Ok(Self::from(*first_byte))
+    fn try_from(byte: u8) -> Result<Self, Error> {
+        let cb = Self(byte);
+        let valid = cb.is_continuation()
+            || cb.is_ack()
+            || cb.is_encrypted_transport()
+            || cb.is_error()
+            || cb.is_ping()
+            || cb.is_pong()
+            || cb.is_channel_allocation_request()
+            || cb.is_channel_allocation_response()
+            || cb.is_handshake()
+            || cb.is_codec_v1();
+        if !valid {
+            log::warn!("Invalid control byte {}.", byte);
+            return Err(Error::malformed_data());
+        }
+        Ok(cb)
     }
 }
 
diff --git a/rust/trezor-thp/src/header.rs b/rust/trezor-thp/src/header.rs
index efe2e910..a066bfb8 100644
--- a/rust/trezor-thp/src/header.rs
+++ b/rust/trezor-thp/src/header.rs
@@ -1,6 +1,6 @@
 use crate::Role;
 use crate::alternating_bit::SyncBits;
-use crate::control_byte::{self, ControlByte};
+use crate::control_byte::ControlByte;
 use crate::crc32;
 use crate::error::{Error, Result};
 
@@ -80,7 +80,7 @@ pub(crate) fn parse_cb_channel(buffer: &[u8]) -> Result<(ControlByte, u16, &[u8]
         return Err(Error::malformed_data());
     };
     let (channel_id, rest) = parse_u16(rest)?;
-    Ok((ControlByte::from(*cb), channel_id, rest))
+    Ok((ControlByte::try_from(*cb)?, channel_id, rest))
 }
 
 impl<R: Role> Header<R> {
@@ -177,7 +177,10 @@ impl<R: Role> Header<R> {
     }
 
     fn parse_unicast(cb: ControlByte, channel_id: u16, payload_len: u16) -> Option<Self> {
-        if let Some(phase) = HandshakeMessage::from_u8::<R>(cb.into()) {
+        if let Some(phase) = cb.handshake_phase() {
+            if !phase.valid_incoming::<R>() {
+                return None;
+            }
             return Some(Self::Handshake {
                 phase,
                 channel_id,
@@ -193,22 +196,25 @@ impl<R: Role> Header<R> {
         None
     }
 
-    pub(crate) fn control_byte(&self, sync_bits: SyncBits) -> Option<ControlByte> {
-        let cb = match self {
-            Self::Continuation { .. } => control_byte::CONTINUATION_PACKET,
-            Self::Ack { .. } => control_byte::ACK_MESSAGE,
-            Self::ChannelAllocationRequest if R::is_host() => control_byte::CHANNEL_ALLOCATION_REQ,
+    pub fn control_byte(&self, sync_bits: SyncBits) -> Option<ControlByte> {
+        let mut cb = match self {
+            Self::Continuation { .. } => ControlByte::continuation(),
+            Self::Ack { .. } => ControlByte::ack(),
+            Self::ChannelAllocationRequest if R::is_host() => {
+                ControlByte::channel_allocation_request()
+            }
             Self::ChannelAllocationResponse { .. } if !R::is_host() => {
-                control_byte::CHANNEL_ALLOCATION_RES
+                ControlByte::channel_allocation_response()
             }
-            Self::TransportError { .. } => control_byte::ERROR,
-            Self::Ping if R::is_host() => control_byte::PING,
-            Self::Pong if !R::is_host() => control_byte::PONG,
-            Self::Handshake { phase, .. } => phase.to_u8::<R>()?,
-            Self::Encrypted { .. } => control_byte::ENCRYPTED_TRANSPORT,
+            Self::TransportError { .. } => ControlByte::error(),
+            Self::Ping if R::is_host() => ControlByte::ping(),
+            Self::Pong if !R::is_host() => ControlByte::pong(),
+            Self::Handshake { phase, .. } if phase.valid_outgoing::<R>() => {
+                ControlByte::handshake(*phase)
+            }
+            Self::Encrypted { .. } => ControlByte::encrypted_transport(),
             _ => return None,
         };
-        let mut cb = ControlByte::from(cb);
         if self.is_ack() && sync_bits.seq_bit() {
             // ACK must have seq_bit=0
             return None;
@@ -224,7 +230,7 @@ impl<R: Role> Header<R> {
             return if dest.len() < Self::CONT_LEN {
                 None
             } else {
-                dest[0] = control_byte::CONTINUATION_PACKET;
+                dest[0] = ControlByte::continuation().into();
                 dest[1..3].copy_from_slice(&channel_id.to_be_bytes());
                 Some(Self::CONT_LEN)
             };
@@ -417,37 +423,20 @@ impl<R: Role> Header<R> {
 }
 
 impl HandshakeMessage {
-    pub fn from_u8<R: Role>(val: u8) -> Option<Self> {
-        let masked = val & control_byte::DATA_MASK;
-        Some(match masked {
-            control_byte::HANDSHAKE_INIT_REQ if !R::is_host() => {
-                HandshakeMessage::InitiationRequest
-            }
-            control_byte::HANDSHAKE_INIT_RES if R::is_host() => {
-                HandshakeMessage::InitiationResponse
-            }
-            control_byte::HANDSHAKE_COMP_REQ if !R::is_host() => {
-                HandshakeMessage::CompletionRequest
-            }
-            control_byte::HANDSHAKE_COMP_RES if R::is_host() => {
-                HandshakeMessage::CompletionResponse
-            }
-            _ => return None,
-        })
+    /// True if this is valid outgoing message for a given role.
+    fn valid_outgoing<R: Role>(&self) -> bool {
+        match self {
+            HandshakeMessage::InitiationRequest if R::is_host() => true,
+            HandshakeMessage::InitiationResponse if !R::is_host() => true,
+            HandshakeMessage::CompletionRequest if R::is_host() => true,
+            HandshakeMessage::CompletionResponse if !R::is_host() => true,
+            _ => false,
+        }
     }
 
-    pub fn to_u8<R: Role>(&self) -> Option<u8> {
-        Some(match self {
-            HandshakeMessage::InitiationRequest if R::is_host() => control_byte::HANDSHAKE_INIT_REQ,
-            HandshakeMessage::InitiationResponse if !R::is_host() => {
-                control_byte::HANDSHAKE_INIT_RES
-            }
-            HandshakeMessage::CompletionRequest if R::is_host() => control_byte::HANDSHAKE_COMP_REQ,
-            HandshakeMessage::CompletionResponse if !R::is_host() => {
-                control_byte::HANDSHAKE_COMP_RES
-            }
-            _ => return None,
-        })
+    /// True if this is valid outgoing message for a given role.
+    fn valid_incoming<R: Role>(&self) -> bool {
+        !self.valid_outgoing::<R>()
     }
 }