From 26daf9d39f6cc1a7239ab217cc0896a9a3affcdf Mon Sep 17 00:00:00 2001
From: Szymon Lesisz <sz.lesisz@gmail.com>
Date: Wed, 23 Oct 2024 19:33:11 +0200
Subject: [PATCH] tmp

---
 packages/connect/src/device/Device.ts                | 5 ++---
 packages/protocol/src/protocol-thp/messages/index.ts | 2 +-
 packages/transport/src/thp/receive.ts                | 3 ---
 packages/transport/src/transports/abstract.ts        | 6 +++++-
 packages/transport/src/utils/receive.ts              | 3 ---
 packages/transport/src/utils/send.ts                 | 3 ---
 6 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/packages/connect/src/device/Device.ts b/packages/connect/src/device/Device.ts
index 2da3b58394ae..f228e52baf49 100644
--- a/packages/connect/src/device/Device.ts
+++ b/packages/connect/src/device/Device.ts
@@ -446,12 +446,11 @@ export class Device extends TypedEmitter<DeviceEvents> {
                     return this._runInner(() => Promise.resolve({}), options);
                 }
                 if (error.code === 'Failure_InvalidProtocol') {
-                    if (this.protocol.name === 'v2') {
-                        // TODO: is this possible?
-                    }
                     _log.info('Changing device protocol to THP');
                     // switch to THP and try again
                     this.protocol = protocolV2;
+                    // TODO: this can throw?
+                    await this.transport.loadMessages('thp', protocolThp.getThpProtobufMessages);
 
                     return this._runInner(fn, options);
                 }
diff --git a/packages/protocol/src/protocol-thp/messages/index.ts b/packages/protocol/src/protocol-thp/messages/index.ts
index 3c6230e12c60..18e53d19ffe9 100644
--- a/packages/protocol/src/protocol-thp/messages/index.ts
+++ b/packages/protocol/src/protocol-thp/messages/index.ts
@@ -1,3 +1,3 @@
-export * from './loadProtobuf';
+export { getThpProtobufMessages } from './protobufDefinitions';
 export * from './messageTypes';
 export * from './protobufTypes';
diff --git a/packages/transport/src/thp/receive.ts b/packages/transport/src/thp/receive.ts
index 096fcf0cc11c..a897b0da7080 100644
--- a/packages/transport/src/thp/receive.ts
+++ b/packages/transport/src/thp/receive.ts
@@ -100,9 +100,6 @@ export type ParseThpMessageProps = {
 export const parseThpMessage = ({ decoded, messages, protocolState }: ParseThpMessageProps) => {
     const isAckExpected = protocolThp.isAckExpected(protocolState?.expectedResponses || []);
 
-    // make sure that THP protobuf messages are loaded
-    protocolThp.loadProtobuf(messages);
-
     const protobufDecoder = (protobufMessageType: string | number, protobufPayload: Buffer) => {
         const { Message, messageName } = createMessageFromType(messages, protobufMessageType);
         const message = decodeProtobuf(Message, protobufPayload);
diff --git a/packages/transport/src/transports/abstract.ts b/packages/transport/src/transports/abstract.ts
index b4587a79a235..40a5c85fb4ce 100644
--- a/packages/transport/src/transports/abstract.ts
+++ b/packages/transport/src/transports/abstract.ts
@@ -13,7 +13,7 @@ import {
     TransportProtocolState,
     thp as protocolThp,
 } from '@trezor/protocol';
-import type { MessageFromTrezor as ProtobufMessageType } from '@trezor/protobuf';
+import { type MessageFromTrezor as ProtobufMessageType, loadProtobuf } from '@trezor/protobuf';
 
 import {
     Session,
@@ -447,6 +447,10 @@ export abstract class AbstractTransport extends TransportEmitter {
         this.messages = protobuf.Root.fromJSON(messages);
     }
 
+    public loadMessages(packageName: string, packageLoader: Parameters<typeof loadProtobuf>[2]) {
+        return loadProtobuf(this.messages, packageName, packageLoader);
+    }
+
     protected success<T>(payload: T): Success<T> {
         return success(payload);
     }
diff --git a/packages/transport/src/utils/receive.ts b/packages/transport/src/utils/receive.ts
index 70cd6cd4d369..019d599a6748 100644
--- a/packages/transport/src/utils/receive.ts
+++ b/packages/transport/src/utils/receive.ts
@@ -62,9 +62,6 @@ export async function receiveAndParse<T extends () => ReturnType<AbstractApi['re
     };
 
     if (protocol.name === 'v2') {
-        // make sure that THP protobuf messages are loaded
-        protocolThp.loadProtobuf(messages);
-
         const { messageName, message } = protocolThp.decode(
             readResult.payload,
             protobufDecoder,
diff --git a/packages/transport/src/utils/send.ts b/packages/transport/src/utils/send.ts
index 47de01926587..b32e844d1d2a 100644
--- a/packages/transport/src/utils/send.ts
+++ b/packages/transport/src/utils/send.ts
@@ -55,9 +55,6 @@ export const buildMessage = ({
     };
 
     if (protocol.name === 'v2') {
-        // make sure that THP protobuf messages are loaded
-        protocolThp.loadProtobuf(messages);
-
         return protocolThp.encode({
             data,
             messageType: name,