refactor(transport): improve transports apis

trezor · Feb 18, 2025 · 52399a2 · 52399a2
1 parent a305fc4
commit 52399a2
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 73 deletions.
diff --git a/packages/transport-bridge/src/core.ts b/packages/transport-bridge/src/core.ts
@@ -11,7 +11,6 @@ import { BridgeProtocolMessage, PathInternal, Session } from '@trezor/transport/
 import { createProtocolMessage } from '@trezor/transport/src/utils/bridgeProtocolMessage';
 import { receive as receiveUtil } from '@trezor/transport/src/utils/receive';
 import { success, unknownError } from '@trezor/transport/src/utils/result';
-import { createChunks, sendChunks } from '@trezor/transport/src/utils/send';
 import { Log } from '@trezor/utils';
 
 export const createCore = (apiArg: 'usb' | 'udp' | AbstractApi, logger?: Log) => {
@@ -57,20 +56,17 @@ export const createCore = (apiArg: 'usb' | 'udp' | AbstractApi, logger?: Log) =>
     }) => {
         logger?.debug(`core: writeUtil protocol ${protocol.name}`);
         const buffer = Buffer.from(data, 'hex');
-        let encodedMessage;
-        let chunkHeader;
-        if (protocol.name === 'bridge') {
-            const { messageType, payload } = protocolBridge.decode(buffer);
-            encodedMessage = protocolV1.encode(payload, { messageType });
-            chunkHeader = protocolV1.getChunkHeader(encodedMessage);
-        } else {
-            encodedMessage = buffer;
-            chunkHeader = protocol.getChunkHeader(encodedMessage);
-        }
 
-        const chunks = createChunks(encodedMessage, chunkHeader, api.chunkSize);
-        const apiWrite = (chunk: Buffer) => api.write(path, chunk, signal);
-        const sendResult = await sendChunks(chunks, apiWrite);
+        // TODO does it handle both v1 and bridge correctly? Also there could be one extra decode
+        const { messageType, payload: message } = protocol.decode(buffer);
+
+        const sendResult = await api.send({
+            message,
+            messageType,
+            protocol: protocolV1,
+            path,
+            signal,
+        });
 
         return sendResult;
     };

diff --git a/packages/transport/src/api/abstract.ts b/packages/transport/src/api/abstract.ts
@@ -1,3 +1,4 @@
+import type { TransportProtocol } from '@trezor/protocol';
 import { TypedEmitter, getSynchronize } from '@trezor/utils';
 
 import * as ERRORS from '../errors';
@@ -10,6 +11,7 @@ import type {
     Success,
 } from '../types';
 import { error, success, unknownError } from '../utils/result';
+import { createChunks } from '../utils/send';
 
 export interface AbstractApiConstructorParams {
     logger?: Logger;
@@ -98,6 +100,32 @@ export abstract class AbstractApi extends TypedEmitter<{
         | typeof ERRORS.UNEXPECTED_ERROR
     >;
 
+    async send({
+        protocol: { encode, getChunkHeader },
+        messageType,
+        message,
+        signal,
+        path,
+    }: {
+        protocol: TransportProtocol;
+        messageType: number | string;
+        message: Buffer;
+        signal?: AbortSignal;
+        path: PathInternal;
+    }) {
+        const bytes = encode(message, { messageType });
+        const chunks = createChunks(bytes, getChunkHeader(bytes), this.chunkSize);
+
+        for (let i = 0; i < chunks.length; i++) {
+            const result = await this.write(path, chunks[i], signal);
+            if (!result.success) {
+                return result;
+            }
+        }
+
+        return { success: true as const, payload: undefined };
+    }
+
     /**
      * set device to the state when it is available to read/write
      */
@@ -132,7 +160,7 @@ export abstract class AbstractApi extends TypedEmitter<{
     /**
      * packet size for api
      */
-    public abstract chunkSize: number;
+    protected abstract chunkSize: number;
 
     protected success<T>(payload: T): Success<T> {
         return success(payload);

diff --git a/packages/transport/src/transports/abstractApi.ts b/packages/transport/src/transports/abstractApi.ts
@@ -13,7 +13,7 @@ import { SessionsClient } from '../sessions/client';
 import { SessionsBackgroundInterface } from '../sessions/types';
 import { Session } from '../types';
 import { receiveAndParse } from '../utils/receive';
-import { buildMessage, createChunks, sendChunks } from '../utils/send';
+import { buildMessage } from '../utils/send';
 
 interface ConstructorParams extends AbstractTransportParams {
     api: AbstractApi;
@@ -194,19 +194,15 @@ export abstract class AbstractApiTransport extends AbstractTransport {
                 const { path } = getPathBySessionResponse.payload;
 
                 const protocol = customProtocol || v1Protocol;
-                const bytes = buildMessage({
-                    messages: this.messages,
-                    name,
-                    data,
-                    encode: protocol.encode,
+                const { messages } = this;
+                const { messageType, message } = buildMessage({ messages, name, data });
+                const sendResult = await this.api.send({
+                    messageType,
+                    message,
+                    path,
+                    protocol,
+                    signal,
                 });
-                const chunks = createChunks(
-                    bytes,
-                    protocol.getChunkHeader(bytes),
-                    this.api.chunkSize,
-                );
-                const apiWrite = (chunk: Buffer) => this.api.write(path, chunk, signal);
-                const sendResult = await sendChunks(chunks, apiWrite);
 
                 if (!sendResult.success) {
                     handleError(sendResult.error);
@@ -242,17 +238,15 @@ export abstract class AbstractApiTransport extends AbstractTransport {
                     return this.error({ error: getPathBySessionResponse.error });
                 }
                 const { path } = getPathBySessionResponse.payload;
-
-                const { encode, getChunkHeader } = protocol || v1Protocol;
-                const bytes = buildMessage({
-                    messages: this.messages,
-                    name,
-                    data,
-                    encode,
+                const { messages } = this;
+                const { messageType, message } = buildMessage({ messages, name, data });
+                const sendResult = await this.api.send({
+                    messageType,
+                    message,
+                    path,
+                    protocol: protocol || v1Protocol,
+                    signal,
                 });
-                const chunks = createChunks(bytes, getChunkHeader(bytes), this.api.chunkSize);
-                const apiWrite = (chunk: Buffer) => this.api.write(path, chunk, signal);
-                const sendResult = await sendChunks(chunks, apiWrite);
 
                 if (!sendResult.success) {
                     if (sendResult.error === ERRORS.DEVICE_DISCONNECTED_DURING_ACTION) {

diff --git a/packages/transport/src/transports/bridge.ts b/packages/transport/src/transports/bridge.ts
@@ -237,12 +237,9 @@ export class BridgeTransport extends AbstractTransport {
         return this.scheduleAction(
             async signal => {
                 const protocol = this.getProtocol(customProtocol);
-                const bytes = buildMessage({
-                    messages: this.messages,
-                    name,
-                    data,
-                    encode: protocol.encode,
-                });
+                const { messages } = this;
+                const { message, messageType } = buildMessage({ messages, name, data });
+                const bytes = protocol.encode(message, { messageType });
                 const response = await this.post(`/call`, {
                     params: session,
                     body: this.getRequestBody(bytes, protocol),
@@ -272,12 +269,9 @@ export class BridgeTransport extends AbstractTransport {
         return this.scheduleAction(
             async signal => {
                 const protocol = this.getProtocol(customProtocol);
-                const bytes = buildMessage({
-                    messages: this.messages,
-                    name,
-                    data,
-                    encode: protocol.encode,
-                });
+                const { messages } = this;
+                const { message, messageType } = buildMessage({ messages, name, data });
+                const bytes = protocol.encode(message, { messageType });
                 const response = await this.post('/post', {
                     params: session,
                     body: this.getRequestBody(bytes, protocol),

diff --git a/packages/transport/src/utils/send.ts b/packages/transport/src/utils/send.ts
@@ -4,9 +4,6 @@
 import { Root } from 'protobufjs/light';
 
 import { createMessageFromName, encode as encodeProtobuf } from '@trezor/protobuf';
-import { TransportProtocolEncode } from '@trezor/protocol';
-
-import { AsyncResultWithTypedError } from '../types';
 
 export const createChunks = (data: Buffer, chunkHeader: Buffer, chunkSize: number) => {
     if (!chunkSize || data.byteLength <= chunkSize) {
@@ -35,28 +32,11 @@ interface BuildMessageProps {
     messages: Root;
     name: string;
     data: Record<string, unknown>;
-    encode: TransportProtocolEncode;
 }
 
-export const buildMessage = ({ messages, name, data, encode }: BuildMessageProps) => {
+export const buildMessage = ({ messages, name, data }: BuildMessageProps) => {
     const { Message, messageType } = createMessageFromName(messages, name);
-    const buffer = encodeProtobuf(Message, data);
-
-    return encode(buffer, {
-        messageType,
-    });
-};
-
-export const sendChunks = async <T, E>(
-    chunks: Buffer[],
-    apiWrite: (chunk: Buffer) => AsyncResultWithTypedError<T, E>,
-) => {
-    for (let i = 0; i < chunks.length; i++) {
-        const result = await apiWrite(chunks[i]);
-        if (!result.success) {
-            return result;
-        }
-    }
+    const message = encodeProtobuf(Message, data);
 
-    return { success: true as const, payload: undefined };
+    return { messageType, message };
 };