chore(transport): use shared method for sending chunks to api

trezor · Jul 18, 2024 · f8e4ec4 · f8e4ec4
1 parent 8f68b00
commit f8e4ec4
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 31 deletions.
diff --git a/packages/transport-bridge/src/core.ts b/packages/transport-bridge/src/core.ts
@@ -2,7 +2,7 @@ import { WebUSB } from 'usb';
 
 import { v1 as protocolV1, bridge as protocolBridge } from '@trezor/protocol';
 import { receive as receiveUtil } from '@trezor/transport/src/utils/receive';
-import { createChunks } from '@trezor/transport/src/utils/send';
+import { createChunks, sendChunks } from '@trezor/transport/src/utils/send';
 import { SessionsBackground } from '@trezor/transport/src/sessions/background';
 import { SessionsClient } from '@trezor/transport/src/sessions/client';
 import { UsbApi } from '@trezor/transport/src/api/usb';
@@ -66,22 +66,15 @@ export const createApi = (apiArg: 'usb' | 'udp' | AbstractApi, logger?: Log) =>
         );
 
         const encodedMessage = protocolV1.encode(payload, { messageType });
-        const buffers = createChunks(
+        const chunks = createChunks(
             encodedMessage,
             protocolV1.getChunkHeader(encodedMessage),
             api.chunkSize,
         );
+        const apiWrite = (chunk: Buffer) => api.write(path, chunk, signal);
+        const sendResult = await sendChunks(chunks, apiWrite);
 
-        for (let i = 0; i < buffers.length; i++) {
-            const bufferSegment = buffers[i];
-
-            const result = await api.write(path, bufferSegment, signal);
-            if (!result.success) {
-                return result;
-            }
-        }
-
-        return { success: true as const };
+        return sendResult;
     };
 
     const readUtil = async ({ path, signal }: { path: string; signal: AbortSignal }) => {

diff --git a/packages/transport/src/transports/abstractApi.ts b/packages/transport/src/transports/abstractApi.ts
@@ -7,7 +7,7 @@ import {
     AbstractTransportMethodParams,
 } from './abstract';
 import { AbstractApi } from '../api/abstract';
-import { buildMessage, createChunks } from '../utils/send';
+import { buildMessage, createChunks, sendChunks } from '../utils/send';
 import { receiveAndParse } from '../utils/receive';
 import { SessionsClient } from '../sessions/client';
 import * as ERRORS from '../errors';
@@ -195,19 +195,15 @@ export abstract class AbstractApiTransport extends AbstractTransport {
                         data,
                         encode: protocol.encode,
                     });
-                    const buffers = createChunks(
+                    const chunks = createChunks(
                         bytes,
                         protocol.getChunkHeader(bytes),
                         this.api.chunkSize,
                     );
-                    for (let i = 0; i < buffers.length; i++) {
-                        const chunk = buffers[i];
-
-                        await this.api.write(path, chunk, signal).then(result => {
-                            if (!result.success) {
-                                throw new Error(result.error);
-                            }
-                        });
+                    const apiWrite = (chunk: Buffer) => this.api.write(path, chunk, signal);
+                    const sendResult = await sendChunks(chunks, apiWrite);
+                    if (!sendResult.success) {
+                        throw new Error(sendResult.error);
                     }
 
                     const message = await receiveAndParse(
@@ -260,15 +256,11 @@ export abstract class AbstractApiTransport extends AbstractTransport {
                         data,
                         encode,
                     });
-                    const buffers = createChunks(bytes, getChunkHeader(bytes), this.api.chunkSize);
-                    for (let i = 0; i < buffers.length; i++) {
-                        const chunk = buffers[i];
-
-                        await this.api.write(path, chunk, signal).then(result => {
-                            if (!result.success) {
-                                throw new Error(result.error);
-                            }
-                        });
+                    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) {
+                        throw new Error(sendResult.error);
                     }
 
                     return this.success(undefined);

diff --git a/packages/transport/src/utils/send.ts b/packages/transport/src/utils/send.ts
@@ -5,6 +5,8 @@ import { Root } from 'protobufjs/light';
 import { encode as encodeProtobuf, createMessageFromName } 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) {
         const buffer = Buffer.alloc(Math.max(chunkSize, data.byteLength));
@@ -43,3 +45,17 @@ export const buildMessage = ({ messages, name, data, encode }: BuildMessageProps
         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;
+        }
+    }
+
+    return { success: true as const };
+};