fix(transport): allow skipping ThpAck from device

packages/protocol/src/protocol-thp/decode.ts

@@ -17,6 +17,7 @@ import { getHandshakeHash, getTrezorState } from './crypto/pairing';
 import { getIvFromNonce } from './crypto/tools';
 import { ThpDeviceProperties, ThpError, ThpMessageResponse } from './messages';
 import { clearControlBit, readThpHeader } from './utils';
+import { MESSAGE_TYPE } from '../protocol-v2/constants';
 
 type ThpMessage = ReturnType<TransportProtocolDecode> & {
     magic: number;
@@ -120,7 +121,7 @@ const readProtobufMessage = (
     return protobufDecoder(messageType, messagePayload) as ThpMessageResponse;
 };
 
-const decodeReadAck = (): ThpMessageResponse => ({
+const decodeReadAck = (): Extract<ThpMessageResponse, { type: 'ThpAck' }> => ({
     type: 'ThpAck',
     message: {},
 });
@@ -131,7 +132,7 @@ const decodeReadAck = (): ThpMessageResponse => ({
 // [magic | channel | len* | error | crc     ]
 // [42    | 1222    | 0005 | 02    | 70303cfa]
 // *len includes error+crc
-const decodeThpError = (payload: Buffer): ThpMessageResponse => {
+const decodeThpError = (payload: Buffer): Extract<ThpMessageResponse, { type: 'ThpError' }> => {
     const [errorType] = payload;
 
     const error = (() => {
@@ -154,10 +155,7 @@ const decodeThpError = (payload: Buffer): ThpMessageResponse => {
         message: error ?? `Unknown ThpError ${errorType}`,
     };
 
-    return {
-        type: 'ThpError',
-        message,
-    };
+    return { type: 'ThpError', message };
 };
 
 const validateCrc = (decodedMessage: ReturnType<TransportProtocolDecode>) => {
@@ -186,13 +184,19 @@ const validateCrc = (decodedMessage: ReturnType<TransportProtocolDecode>) => {
 
 // Decode protocol-v2 message from thp send process: ThpAck or ThpError
 export const decodeSendAck = (decodedMessage: MessageV2) => {
-    validateCrc(decodedMessage);
-
     const header = readThpHeader(decodedMessage.header);
     const magic = clearControlBit(header.magic);
+
+    if (magic === THP_CONTROL_BYTE_ENCRYPTED) {
+        return { type: MESSAGE_TYPE } as const;
+    }
+
+    validateCrc(decodedMessage);
+
     if (magic === THP_ERROR_HEADER_BYTE) {
         return decodeThpError(decodedMessage.payload);
     }
+
     if (magic === THP_READ_ACK_HEADER_BYTE) {
         return decodeReadAck();
     }

packages/protocol/tests/protocol-thp/protocol-thp.test.ts

@@ -105,7 +105,7 @@ describe('protocol-thp', () => {
         const unexpected = decodeSendAck(decodeV2(Buffer.from('40ffff0004f9215951', 'hex')));
         expect(unexpected).toBe(undefined);
 
-        expect(() => decodeSendAck(decodeV2(Buffer.from('40ffff000499999999', 'hex')))).toThrow(
+        expect(() => decodeSendAck(decodeV2(Buffer.from('2812340004e98c8598', 'hex')))).toThrow(
             'Invalid CRC',
         );
     });

packages/transport-bridge/src/core.ts

@@ -340,7 +340,10 @@ export const createCore = (apiArg: 'usb' | 'udp' | AbstractApi, logger?: Log) =>
                 return writeResult;
             }
 
-            return createProtocolMessageResponse(writeResult, protocolName);
+            return createProtocolMessageResponse(
+                { success: true, payload: undefined } as const,
+                protocolName,
+            );
         }
 
         const writeResult = await writeUtil({ path, data, signal, protocol });

packages/transport/src/thp/call.ts

@@ -26,6 +26,7 @@ export const callThpMessage = async ({
 
     // read and send ThpAck
     const receiveResult = await receiveThpMessage({
+        firstReadResult: sendResult.payload,
         thpState,
         apiWrite,
         apiRead,

packages/transport/src/thp/receive.ts

@@ -17,9 +17,11 @@ export type ReceiveThpMessageProps = {
     signal?: AbortSignal;
     graceful?: boolean;
     logger?: Logger;
+    firstReadResult?: Awaited<ReturnType<AbstractApi['read']>>;
 };
 
 export const receiveThpMessage = async ({
+    firstReadResult,
     thpState,
     skipAck,
     apiRead,
@@ -40,9 +42,19 @@ export const receiveThpMessage = async ({
             graceful,
             logger,
         });
+
+        let firstRead = !!firstReadResult;
+        const getFirstResult = () => {
+            if (firstRead && firstReadResult) {
+                firstRead = false;
+
+                return Promise.resolve(firstReadResult);
+            }
+        };
+
         const expectedHeaders = protocolThp.getExpectedHeaders(thpState);
         const message = await receive(
-            () => apiReadWithExpectedHeaders(expectedHeaders),
+            () => getFirstResult() ?? apiReadWithExpectedHeaders(expectedHeaders),
             protocolV2,
         );
         if (!message.success) {

packages/transport/src/thp/send.ts

@@ -42,7 +42,7 @@ export const sendThpMessage = async ({
 
     // ThpAck is expected.
     // set expectedResponses to ThpAck
-    thpState.setExpectedResponses([0x20]); // THP_READ_ACK_HEADER_BYTE
+    thpState.setExpectedResponses([0x20, ...expectedResponses]); // THP_READ_ACK_HEADER_BYTE
 
     let attempt = 0;
 
@@ -114,7 +114,10 @@ export const sendThpMessage = async ({
         // set expectedResponses as they will be used in receiveThpMessage
         thpState.setExpectedResponses(expectedResponses);
 
-        return success(undefined);
+        const firstReadResult =
+            decodedResult?.type === 'TrezorHostProtocolMessage' ? result : undefined;
+
+        return success(firstReadResult);
     } catch (err) {
         logger?.error(`sendThpMessage error ${err.message}`);
 

packages/transport/src/transports/abstractApi.ts

@@ -373,9 +373,11 @@ export abstract class AbstractApiTransport extends AbstractTransport {
                     if (sendResult.error === ERRORS.DEVICE_DISCONNECTED_DURING_ACTION) {
                         this.enumerate();
                     }
+
+                    return sendResult;
                 }
 
-                return sendResult;
+                return { success: true, payload: undefined } as const;
             },
             { signal, graceful: true, timeout },
         );