feat(protocol): encode/decode messageType as string

possible values: messageId (number) or messageName (string)

packages/protobuf/src/utils.ts

@@ -40,22 +40,31 @@ export function parseConfigure(data: protobuf.INamespace) {
 export const createMessageFromName = (messages: protobuf.Root, name: string) => {
     const Message = messages.lookupType(name);
     const MessageType = messages.lookupEnum('MessageType');
-    let messageType = MessageType.values[`MessageType_${name}`];
+    let messageTypeId = MessageType.values[`MessageType_${name}`];
 
-    if (!messageType && Message.options) {
-        messageType = Message.options['(wire_type)'];
+    if (typeof messageTypeId !== 'number' && Message.options) {
+        messageTypeId = Message.options['(wire_type)'];
     }
 
     return {
         Message,
-        messageType,
+        messageType: messageTypeId ?? name,
     };
 };
 
-export const createMessageFromType = (messages: protobuf.Root, typeId: number) => {
-    const MessageType = messages.lookupEnum('MessageType');
+export const createMessageFromType = (messages: protobuf.Root, messageType: number | string) => {
+    if (typeof messageType === 'string') {
+        const Message = messages.lookupType(messageType);
 
-    const messageName = MessageType.valuesById[typeId].replace(
+        return {
+            Message,
+            messageName: messageType as MessageFromTrezor['type'],
+        };
+    }
+
+    const messageTypes = messages.lookupEnum('MessageType');
+
+    const messageName = messageTypes.valuesById[messageType].replace(
         'MessageType_',
         '',
     ) as MessageFromTrezor['type'];

packages/protobuf/tests/messages.test.ts

@@ -1,6 +1,6 @@
 import * as protobuf from 'protobufjs/light';
 
-import { createMessageFromName } from '../src/utils';
+import { createMessageFromName, createMessageFromType } from '../src/utils';
 
 const json = {
     nested: {
@@ -34,6 +34,17 @@ const json = {
                                         },
                                     },
                                 },
+                                Initialize: {
+                                    fields: {
+                                        session_id: {
+                                            type: 'bytes',
+                                            id: 1,
+                                        },
+                                    },
+                                },
+                                MessageWithoutId: {
+                                    fields: {},
+                                },
                                 MessageType: {
                                     values: {
                                         MessageType_Initialize: 0,
@@ -55,4 +66,20 @@ describe('messages', () => {
 
         expect(() => createMessageFromName(messages, name)).not.toThrow();
     });
+
+    test('createMessageFromName (message without id)', () => {
+        const messages = protobuf.Root.fromJSON(json);
+        const name = 'MessageWithoutId'; // this message has no id defined in MessageType.MessageType_xxx
+        const result = createMessageFromName(messages, name);
+
+        expect(result.messageType).toEqual(name);
+    });
+
+    test('createMessageFromType (messageType as number and string)', () => {
+        const messages = protobuf.Root.fromJSON(json);
+        const asNumber = createMessageFromType(messages, 0);
+        const asString = createMessageFromType(messages, 'Initialize');
+
+        expect(asNumber.messageName).toEqual(asString.messageName);
+    });
 });

packages/protocol/src/protocol-bridge/encode.ts

@@ -6,6 +6,9 @@ import { TransportProtocolEncode } from '../types';
 // it is because bridge does some parts of the protocol itself (like chunking)
 export const encode: TransportProtocolEncode = (data, options) => {
     const { messageType } = options;
+    if (typeof messageType === 'string') {
+        throw new Error(`Unsupported message type ${messageType}`);
+    }
 
     const encodedBuffer = Buffer.alloc(HEADER_SIZE + data.length);
 

packages/protocol/src/protocol-v1/encode.ts

@@ -8,6 +8,10 @@ import { TransportProtocolEncode } from '../types';
 
 export const encode: TransportProtocolEncode = (data, options) => {
     const { messageType } = options;
+    if (typeof messageType === 'string') {
+        throw new Error(`Unsupported message type ${messageType}`);
+    }
+
     const fullSize = HEADER_SIZE + data.length;
     const chunkSize = options.chunkSize || BUFFER_SIZE;
 

packages/protocol/src/types.ts

@@ -1,11 +1,11 @@
 export type TransportProtocolDecode = (bytes: ArrayBuffer) => {
     length: number;
-    messageType: number;
+    messageType: number | string;
     payload: Buffer;
 };
 
 export interface TransportProtocolEncodeOptions {
-    messageType: number;
+    messageType: number | string;
     chunkSize?: number;
 }
 

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

@@ -14,6 +14,11 @@ describe('protocol-bridge', () => {
         expect(chunks[0].subarray(0, 6).toString('hex')).toEqual('003700000173');
         expect(chunks[0].readUint32BE(2)).toEqual(371);
         expect(chunks[0].length).toEqual(371 + 6);
+
+        // fail to encode unsupported messageType (string)
+        expect(() => bridge.encode(Buffer.alloc(64), { messageType: 'Initialize' })).toThrow(
+            'Unsupported message type Initialize',
+        );
     });
 
     it('decode', () => {

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

@@ -22,6 +22,11 @@ describe('protocol-v1', () => {
                 expect(chunk.subarray(0, 5).toString('hex')).toEqual('3f00000000');
             }
         });
+
+        // fail to encode unsupported messageType (string)
+        expect(() => v1.encode(Buffer.alloc(64), { messageType: 'Initialize' })).toThrow(
+            'Unsupported message type Initialize',
+        );
     });
 
     it('decode', () => {