refactor(connect): use improved message types

Author: marekrjpolak

packages/connect/src/device/DeviceCommands.ts

@@ -18,20 +18,16 @@ import { initLog } from '../utils/debug';
 import * as hdnodeUtils from '../utils/hdnodeUtils';
 import { getScriptType, getSerializedPath, isTaprootPath, toHardened } from '../utils/pathUtils';
 
-type MessageType = Messages.MessageType;
-type MessageKey = keyof MessageType;
-type TypedPayload<T extends MessageKey> = {
-    type: T;
-    message: MessageType[T];
-};
-type TypedCallResponseMap = {
-    [K in keyof MessageType]: TypedPayload<K>;
-};
-type DefaultPayloadMessage = TypedCallResponseMap[keyof MessageType];
+type TypedCall = Messages.TypedCall;
+
+export type { TypedCall };
 
 const logger = initLog('DeviceCommands');
 
-const assertType = (res: DefaultPayloadMessage, resType: MessageKey | MessageKey[]) => {
+const assertType = (
+    res: Messages.MessageResponse,
+    resType: Messages.MessageKey | Messages.MessageKey[],
+) => {
     const splitResTypes = Array.isArray(resType) ? resType : resType.split('|');
     if (!splitResTypes.includes(res.type)) {
         throw ERRORS.TypedError(
@@ -295,10 +291,10 @@ export class DeviceCommands {
     }
 
     // Sends an async message to the opened device.
-    private async call(
-        type: MessageKey,
-        msg: DefaultPayloadMessage['message'] = {},
-    ): Promise<DefaultPayloadMessage> {
+    private async call<T extends Messages.MessageKey>(
+        type: T,
+        msg: Messages.MessagePayload<T>,
+    ): Promise<Messages.MessageResponse> {
         logger.debug('Sending', type, filterForLog(type, msg));
 
         this.callPromise = this.transport.call({
@@ -327,32 +323,30 @@ export class DeviceCommands {
             filterForLog(res.payload.type, res.payload.message),
         );
 
-        // TODO: https://github.com/trezor/trezor-suite/issues/5301
-        // @ts-expect-error
         return res.payload;
     }
 
-    typedCall<T extends MessageKey, R extends MessageKey[]>(
+    typedCall<T extends Messages.MessageKey, R extends Messages.MessageKey[]>(
         type: T,
         resType: R,
-        msg?: MessageType[T],
-    ): Promise<TypedCallResponseMap[R[number]]>;
-    typedCall<T extends MessageKey, R extends MessageKey>(
+        msg?: Messages.MessagePayload<T>,
+    ): Promise<Messages.MessageResponse<R[number]>>;
+    typedCall<T extends Messages.MessageKey, R extends Messages.MessageKey>(
         type: T,
         resType: R,
-        msg?: MessageType[T],
-    ): Promise<TypedPayload<R>>;
+        msg?: Messages.MessagePayload<T>,
+    ): Promise<Messages.MessageResponse<R>>;
     async typedCall(
-        type: MessageKey,
-        resType: MessageKey | MessageKey[],
-        msg?: DefaultPayloadMessage['message'],
+        type: Messages.MessageKey,
+        resType: Messages.MessageKey | Messages.MessageKey[],
+        msg: Messages.MessagePayload = {},
     ) {
         if (this.disposed) {
             throw ERRORS.TypedError('Runtime', 'typedCall: DeviceCommands already disposed');
         }
         // Assert message type
         // msg is allowed to be undefined for some calls, in that case the schema is an empty object
-        Assert(Messages.MessageType.properties[type], msg ?? {});
+        Assert(Messages.MessageType.properties[type], msg);
         const response = await this._commonCall(type, msg);
         try {
             assertType(response, resType);
@@ -380,7 +374,7 @@ export class DeviceCommands {
         return response;
     }
 
-    async _commonCall(type: MessageKey, msg?: DefaultPayloadMessage['message']) {
+    async _commonCall<T extends Messages.MessageKey>(type: T, msg: Messages.MessagePayload<T>) {
         if (this.disposed) {
             throw ERRORS.TypedError('Runtime', 'typedCall: DeviceCommands already disposed');
         }
@@ -389,7 +383,7 @@ export class DeviceCommands {
         return this._filterCommonTypes(resp);
     }
 
-    _filterCommonTypes(res: DefaultPayloadMessage): Promise<DefaultPayloadMessage> {
+    _filterCommonTypes(res: Messages.MessageResponse): Promise<Messages.MessageResponse> {
         this.device.clearCancelableAction();
 
         if (res.type === 'Failure') {
@@ -589,5 +583,3 @@ export class DeviceCommands {
         }
     }
 }
-
-export type TypedCall = DeviceCommands['typedCall'];

packages/connect/src/device/prompts.ts

@@ -35,6 +35,9 @@ export const cancelPrompt = (device: Device, expectResponse = true) => {
     return expectResponse ? device.transport.call(cancelArgs) : device.transport.send(cancelArgs);
 };
 
+const extractMessage = (payload?: Messages.MessageResponse) =>
+    (payload && 'message' in payload.message && payload.message.message) || '';
+
 const prompt = <E extends PromptEvents>(event: E, { device, ...rest }: DeviceEventArgs<E>) =>
     // return non nullable first arg of PromptCallback<E>
     new Promise<PromptReturnType<E>>(resolve => {
@@ -43,8 +46,8 @@ const prompt = <E extends PromptEvents>(event: E, { device, ...rest }: DeviceEve
                 response.success
                     ? resolve({
                           success: false,
-                          error: error || (response.payload?.message.message as string),
-                          message: response.payload?.message.message as string,
+                          error: error || extractMessage(response.payload),
+                          message: extractMessage(response.payload),
                           isTransportError: !response.success,
                       })
                     : resolve({

packages/protobuf/src/decode.ts

@@ -1,5 +1,6 @@
 import { Field, Message as MessageType, Type } from 'protobufjs/light';
 
+import type { MessageResponse } from './messages';
 import { createMessageFromType, isPrimitiveField } from './utils';
 
 const transform = (field: Field, value: any) => {
@@ -84,5 +85,5 @@ export const decodeMessage = (
     const { Message, messageName } = createMessageFromType(messages, messageType);
     const message = decode(Message, data);
 
-    return { messageName, message };
+    return { type: messageName, message } as MessageResponse;
 };

packages/protobuf/src/index.ts

@@ -17,7 +17,6 @@ export const { parseConfigure, decodeMessage, encodeMessage } = (() => {
     return { parseConfigure: parse, decodeMessage: decode, encodeMessage: encode };
 })();
 
-export * from './types';
 export * as Messages from './messages';
 export { loadDefinitions } from './load-definitions';
 export * as MessagesSchema from './messages-schema';

packages/protobuf/src/types.ts

@@ -2,7 +2,7 @@
 
 import * as protobuf from 'protobufjs/light';
 
-import type { MessageFromTrezor } from './types';
+import type { MessageKey } from './messages';
 
 const primitiveTypes = [
     'bool',
@@ -58,13 +58,13 @@ export const createMessageFromType = (messages: protobuf.Root, messageType: numb
 
         return {
             Message,
-            messageName: messageType as MessageFromTrezor['type'],
+            messageName: messageType as MessageKey,
         };
     }
 
     const messageTypes = messages.lookupEnum('MessageType');
 
-    const messageName = messageTypes.valuesById[messageType] as MessageFromTrezor['type'];
+    const messageName = messageTypes.valuesById[messageType] as MessageKey;
 
     const Message = messages.lookupType(messageName);
 

packages/protobuf/src/utils.ts

@@ -1,4 +1,4 @@
-import { MessageFromTrezor, loadDefinitions, parseConfigure } from '@trezor/protobuf';
+import { Messages, loadDefinitions, parseConfigure } from '@trezor/protobuf';
 import { PROTOCOL_MALFORMED, TransportProtocol } from '@trezor/protocol';
 import { ScheduleActionParams, ScheduledAction, TypedEmitter, scheduleAction } from '@trezor/utils';
 
@@ -260,7 +260,7 @@ export abstract class AbstractTransport extends TransportEmitter {
             session: Session;
             protocol?: TransportProtocol;
         } & AbortableParam,
-    ): AsyncResultWithTypedError<MessageFromTrezor, ReadWriteError>;
+    ): AsyncResultWithTypedError<Messages.MessageResponse, ReadWriteError>;
 
     /**
      * Send and read after that
@@ -272,7 +272,7 @@ export abstract class AbstractTransport extends TransportEmitter {
             data: Record<string, unknown>;
             protocol?: TransportProtocol;
         } & AbortableParam,
-    ): AsyncResultWithTypedError<MessageFromTrezor, ReadWriteError>;
+    ): AsyncResultWithTypedError<Messages.MessageResponse, ReadWriteError>;
 
     /**
      * Stop transport = remove all listeners + try to release session + cancel all requests

packages/transport/src/transports/abstract.ts

@@ -44,7 +44,7 @@ export async function receiveAndParse<T extends () => ReturnType<AbstractApi['re
     if (!readResult.success) return readResult;
 
     const { messageType, payload } = readResult.payload;
-    const { messageName, message } = decodeMessage(messages, messageType, payload);
+    const message = decodeMessage(messages, messageType, payload);
 
-    return success({ message, type: messageName });
+    return success(message);
 }