fixup! WIP: switch THP to rust/trezor-thp

Author: mmilata

core/embed/rust/src/crypto/mod.rs

@@ -35,12 +35,10 @@ impl From<Error> for crate::error::Error {
     }
 }
 
-/// Constant time bytestring comparison. Panics if the parameters don't have the
-/// same length.
-pub fn consteq(a: &[u8], b: &[u8]) -> bool {
-    ensure!(a.len() == b.len(), "Invalid paramter lengths");
+/// Constant time bytestring comparison for two arrays of the same length.
+pub fn consteq<const N: usize>(a: &[u8; N], b: &[u8; N]) -> bool {
     let mut diff: u8 = 0;
-    for i in 0..a.len() {
+    for i in 0..N {
         diff |= a[i] ^ b[i];
     }
     diff == 0

core/embed/rust/src/thp/crypto.rs

@@ -132,8 +132,8 @@ impl Cipher for TrezorCryptoAesGcm {
         let mut full_nonce = [0u8; 12];
         full_nonce[4..].copy_from_slice(&nonce.to_be_bytes());
 
-        out.copy_from_slice(&ciphertext[..out.len()]);
-        let tag = &ciphertext[out.len()..];
+        let (ciphertext, tag) = unwrap!(ciphertext.split_last_chunk::<{ aesgcm::TAG_SIZE }>());
+        out.copy_from_slice(ciphertext);
 
         init_ctx!(aesgcm::AesGcm, ctx, key.as_slice(), &full_nonce);
         let mut ctx = unwrap!(ctx);
@@ -160,7 +160,8 @@ impl Cipher for TrezorCryptoAesGcm {
         let mut full_nonce = [0u8; 12];
         full_nonce[4..].copy_from_slice(&nonce.to_be_bytes());
 
-        let (in_out, tag) = in_out[..ciphertext_len].split_at_mut(ciphertext_len - 16);
+        let in_out = &mut in_out[..ciphertext_len];
+        let (in_out, tag) = unwrap!(in_out.split_last_chunk_mut::<{ aesgcm::TAG_SIZE }>());
 
         init_ctx!(aesgcm::AesGcm, ctx, key.as_slice(), &full_nonce);
         let mut ctx = unwrap!(ctx);

core/embed/rust/src/thp/mod.rs

@@ -405,6 +405,13 @@ impl InterfaceContext {
             return Err(CHANNEL_NOT_FOUND);
         };
         match retry {
+            None => {
+                log::error!(
+                    "[{:04x}] Requested to retransmit but not currently sending.",
+                    channel_id
+                );
+                Ok(true)
+            }
             Some(r) if r > MAX_RETRANSMISSION_COUNT => {
                 log::warn!(
                     "[{:04x}] Closing channel after too many retransmissions.",
@@ -413,15 +420,9 @@ impl InterfaceContext {
                 self.channel_close(channel_id);
                 Ok(false)
             }
-            None => {
-                log::error!(
-                    "[{:04x}] Requested to retransmit but not currently sending.",
-                    channel_id
-                );
-                Ok(true)
-            }
-            Some(r) => {
-                // TODO explain
+            Some(_) => {
+                // Update last write so that we don't end up retransmitting too quickly if
+                // writes are blocked for some reason.
                 if let Some(t) = self.timing.get_mut(&channel_id) {
                     t.update_last_write(Instant::now(), None);
                 }

core/src/trezor/wire/thp/channel.py

@@ -81,7 +81,8 @@ def __init__(
             ):
                 self.credential = decode_credential(inner)
         except Exception as e:
-            self._log(f"Cannot parse credential: {e}")
+            if __debug__:
+                self._log(f"cannot parse credential: {e}")
 
     def channel_id_bytes(self) -> bytes:
         return self.channel_id.to_bytes(2, "big")
@@ -91,7 +92,11 @@ def iface(self) -> WireInterface:
         return self.iface_ctx._iface
 
     def clear(self, exc: Exception | None = None) -> None:
-        self._log(f"Closing channel (exception: {exc is not None})")
+        if __debug__:
+            if exc:
+                self._log(f"closing channel (exception: {exc.__class__.__name__})")
+            else:
+                self._log("closing channel")
         clear_sessions_with_channel_id(self.channel_id_bytes())
         trezorthp.channel_close(self.iface.iface_num(), self.channel_id)
         self.expecting_message = False
@@ -126,7 +131,8 @@ def get_last_write(self) -> int | None:
             return None
 
     def set_channel_state(self, state: ChannelState) -> None:
-        self._log(f"set state {state}")
+        if __debug__:
+            self._log(f"set state {state}")
         self.state = state
 
     def end_pairing_and_replace(self) -> None:
@@ -163,7 +169,8 @@ async def read(self) -> tuple[int, Message]:
         finally:
             # wake up write loop to send ACKs or DECRYPTION_FAILED
             self.iface_ctx.request_write()
-        self._log("message is ready")
+        if __debug__ and _TRACE:
+            self._log("message is ready")
         message = Message(
             message_type,
             self.receive_buffer[3:][:message_bytes_len],
@@ -235,7 +242,7 @@ def write_packet(self, packet: AnyBuffer) -> bool:
             res = trezorthp.packet_out(
                 self.iface.iface_num(), self.channel_id, buffer, packet
             )
-            if res is False and self.send_buffer is not None:
+            if self.send_buffer:
                 self.expecting_ack = True
                 self.iface_ctx.request_read()
             return res

core/src/trezor/wire/thp/interface_context.py

@@ -66,10 +66,11 @@ def preempt_active_channel_if_stale(self) -> None:
         last_write_ms = self.active_channel.get_last_write()
         # TODO does not work well if we're stuck in retransmit loop
         if last_write_ms is None or last_write_ms > _PREEMPT_TIMEOUT_MS:
-            log.error(
-                __name__,
-                f"Interrupting channel {hex(self.active_channel.channel_id)} after {last_write_ms} ms",
-            )
+            if __debug__:
+                log.error(
+                    __name__,
+                    f"Interrupted channel {hex(self.active_channel.channel_id)[2:]} after {last_write_ms} ms",
+                )
             self.active_channel.clear(exc=CHANNEL_PREEMPTED)
 
     async def close(self) -> None:
@@ -189,7 +190,7 @@ def should_read(self) -> bool:
         """
         We want to avoid thefollowing sequence of events:
         - workflow triggered by a message has finished,
-        - message is reassembled before session.py restart,
+        - next message is reassembled before session.py restarts,
         - session is restarted, receive buffer is lost,
         - host has to resend message after a delay.
         To avoid unnecessarry delay, interface is only awaited when:
@@ -395,13 +396,14 @@ def handle_handshake_key(self, try_to_unlock: bool) -> None:
         elif not try_to_unlock:
             trezorthp.handshake_key(self._iface.iface_num(), None)
         elif self._handshake_key_task is None:
-            log.debug(
-                __name__,
-                "Static key needed but device is locked, spawning unlock dialog",
-                iface=self._iface,
-            )
+            if __debug__:
+                log.debug(
+                    __name__,
+                    "Static key needed but device is locked, spawning unlock dialog",
+                    iface=self._iface,
+                )
             self._handshake_key_task = loop.spawn(self.handshake_unlock())
-        else:
+        elif __debug__:
             log.debug(__name__, "Unlock task already running", iface=self._iface)
 
     async def handshake_unlock(self) -> None:

python/src/trezorlib/thp/channel.py

@@ -420,7 +420,7 @@ def should_back_off() -> bool:
             retries_left -= 1
             time.sleep(retry_backoff_time)
             retry_backoff_time *= 2
-            print("RETRY")
+            LOG.warning("RETRY %04x", self.channel_id)
             return True
 
         while True:
@@ -476,9 +476,6 @@ def _read_ack(self, message: Message) -> None:
                 # process ACK bit, return message in next _read
                 self._next_message = message
             else:
-                print(
-                    f"{self.is_ack_piggybacking_allowed} {expected_seq_bit} {message.ack_bit}"
-                )
                 LOG.error("Received message is not a valid ACK: %s", message)
                 # data messages and their acks should have been handled by _read()
                 continue