fixing bridge

trezor · Jul 30, 2024 · 6eff732 · 6eff732
1 parent 2b76552
commit 6eff732
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 6 deletions.
diff --git a/packages/transport/src/transports/bridge.ts b/packages/transport/src/transports/bridge.ts
@@ -271,7 +271,7 @@ export class BridgeTransport extends AbstractTransport {
     ) {
         if (this.legacyBridge && protocol?.name === 'v2') {
             // TODO: try catch + better condition?
-            const { body, state } = JSON.parse(response);
+            const { body, state } = response as any;
             protocolState?.deserialize(state);
 
             return body;
@@ -316,6 +316,7 @@ export class BridgeTransport extends AbstractTransport {
                     this.messages,
                     () => Promise.resolve(Buffer.from(payload, 'hex')),
                     protocol,
+                    protocolState,
                 );
 
                 if (protocolState?.shouldUpdateNonce(name)) {

diff --git a/packages/transport/src/utils/bridgeApiResult.ts b/packages/transport/src/utils/bridgeApiResult.ts
@@ -11,6 +11,10 @@ function isString(payload: UnknownPayload): payload is string {
     return typeof payload === 'string';
 }
 
+function isObject(payload: UnknownPayload): payload is string {
+    return typeof payload === 'object';
+}
+
 export function info(res: UnknownPayload) {
     if (isString(res)) {
         return error({ error: ERRORS.WRONG_RESULT_TYPE });
@@ -80,7 +84,7 @@ export function acquire(res: UnknownPayload) {
 }
 
 export function call(res: UnknownPayload) {
-    if (!isString(res)) {
+    if (!isString(res) && !isObject(res)) {
         return error({ error: ERRORS.WRONG_RESULT_TYPE });
     }
 

diff --git a/packages/transport/src/utils/receive.ts b/packages/transport/src/utils/receive.ts
@@ -1,7 +1,7 @@
 import { Root } from 'protobufjs/light';
 
 import { decode as decodeProtobuf, createMessageFromType } from '@trezor/protobuf';
-import { TransportProtocol } from '@trezor/protocol';
+import { TransportProtocol, TransportProtocolState, thp as protocolThp } from '@trezor/protocol';
 
 async function receiveRest(
     result: Buffer,
@@ -43,10 +43,37 @@ export async function receiveAndParse(
     messages: Root,
     receiver: () => Promise<Buffer>,
     protocol: TransportProtocol,
+    protocolState?: TransportProtocolState,
 ) {
-    const { messageType, payload } = await receive(receiver, protocol);
-    const { Message, messageName } = createMessageFromType(messages, messageType);
-    const message = decodeProtobuf(Message, payload);
+    const decoded = await receive(receiver, protocol);
+
+    const protobufDecoder = (protobufMessageType: string | number, protobufPayload: Buffer) => {
+        const { Message, messageName } = createMessageFromType(messages, protobufMessageType);
+        const message = decodeProtobuf(Message, protobufPayload);
+
+        return {
+            messageName,
+            message,
+        };
+    };
+
+    if (protocol.name === 'v2') {
+        // make sure that THP protobuf messages are loaded
+        protocolThp.loadProtobuf(messages);
+
+        const { messageName, message } = protocolThp.decode(
+            decoded,
+            protobufDecoder,
+            protocolState,
+        );
+
+        return {
+            message,
+            type: messageName,
+        };
+    }
+
+    const { messageName, message } = protobufDecoder(decoded.messageType, decoded.payload);
 
     return {
         message,