feat(connect): common function for DeviceEvent listeners registration

szymonlesisz · szymonlesisz · commit c5622d0c26ff · 2025-03-11T12:54:00.000+01:00
diff --git a/packages/connect/src/core/__tests__/onCallFirmwareUpdate.test.ts b/packages/connect/src/core/__tests__/onCallFirmwareUpdate.test.ts
@@ -193,6 +193,7 @@ const setupTest = () => {
         deviceList,
         postMessage,
         initDevice: () => Promise.resolve(deviceList.getAllDevices()[0]),
+        registerEvents: () => {},
         log: new Log('Test', false),
         abortSignal: new AbortController().signal,
     };
diff --git a/packages/connect/src/core/index.ts b/packages/connect/src/core/index.ts
@@ -639,21 +639,8 @@ const onCallDevice = async (
 
     // device is available
     // set public variables, listeners and run method
-    device.on(DEVICE.BUTTON, onDeviceButtonHandler(context, method));
-    device.on(DEVICE.PIN, onDevicePinHandler(context));
-    device.on(DEVICE.WORD, onDeviceWordHandler(context));
-    device.on(
-        DEVICE.PASSPHRASE,
-        (method.useEmptyPassphrase ? onEmptyPassphraseHandler : onDevicePassphraseHandler)(context),
-    );
-    device.on(DEVICE.PASSPHRASE_ON_DEVICE, () => {
-        sendCoreMessage(
-            createUiMessage(UI.REQUEST_PASSPHRASE_ON_DEVICE, { device: device.toMessageObject() }),
-        );
-    });
-    device.on(DEVICE.FIRMWARE_VERSION_CHANGED, payload => {
-        sendCoreMessage(createDeviceMessage(DEVICE.FIRMWARE_VERSION_CHANGED, payload));
-    });
+    registerDeviceEvents(context, method)(device);
+
     if (useCoreInPopup && env === 'webextension' && origin) {
         device.initStorage(new WebextensionStateStorage(origin));
     }
@@ -783,17 +770,17 @@ const closePopup = ({ popupPromise, sendCoreMessage }: CoreContext) => {
  * @memberof Core
  */
 const onDeviceButtonHandler =
-    (context: CoreContext, method: AbstractMethod<any>) =>
+    (context: CoreContext, method?: AbstractMethod<any>) =>
     async ({ device, payload: request }: DeviceEvents['button']) => {
         const { sendCoreMessage } = context;
         // wait for popup handshake
         const addressRequest = request.code === 'ButtonRequest_Address';
-        if (!addressRequest || (addressRequest && method.useUi)) {
+        if (!addressRequest || (addressRequest && method?.useUi)) {
             await waitForPopup(context);
         }
         const data =
-            typeof method.getButtonRequestData === 'function' && request.code
-                ? method.getButtonRequestData(request.code)
+            typeof method?.getButtonRequestData === 'function' && request.code
+                ? method?.getButtonRequestData(request.code)
                 : undefined;
         // interaction timeout
         startInteractionTimeout(context);
@@ -808,7 +795,7 @@ const onDeviceButtonHandler =
                 data,
             }),
         );
-        if (addressRequest && !method.useUi) {
+        if (addressRequest && !method?.useUi) {
             sendCoreMessage(createUiMessage(UI.ADDRESS_VALIDATION, data));
         }
     };
@@ -903,6 +890,30 @@ const onEmptyPassphraseHandler =
         callback({ value: '' });
     };
 
+const registerDeviceEvents =
+    (context: CoreContext, method?: AbstractMethod<any>) => (device: Device) => {
+        device.removeAllListeners();
+        device.on(DEVICE.BUTTON, onDeviceButtonHandler(context, method));
+        device.on(DEVICE.PIN, onDevicePinHandler(context));
+        device.on(DEVICE.WORD, onDeviceWordHandler(context));
+        device.on(
+            DEVICE.PASSPHRASE,
+            (method?.useEmptyPassphrase ? onEmptyPassphraseHandler : onDevicePassphraseHandler)(
+                context,
+            ),
+        );
+        device.on(DEVICE.PASSPHRASE_ON_DEVICE, () => {
+            context.sendCoreMessage(
+                createUiMessage(UI.REQUEST_PASSPHRASE_ON_DEVICE, {
+                    device: device.toMessageObject(),
+                }),
+            );
+        });
+        device.on(DEVICE.FIRMWARE_VERSION_CHANGED, payload => {
+            context.sendCoreMessage(createDeviceMessage(DEVICE.FIRMWARE_VERSION_CHANGED, payload));
+        });
+    };
+
 /**
  * Handle popup closed by user.
  * @returns {void}
@@ -1158,14 +1169,17 @@ export class Core extends EventEmitter {
                 // means that call immediately returns error.
                 if (message.payload.method === 'firmwareUpdate') {
                     assertDeviceListConnected(this.deviceList);
+
+                    const coreContext = this.getCoreContext();
                     onCallFirmwareUpdate({
                         params: message.payload,
                         context: {
                             deviceList: this.deviceList,
                             postMessage: this.sendCoreMessage.bind(this),
-                            initDevice: path => initDevice(this.getCoreContext(), { path }),
+                            initDevice: path => initDevice(coreContext, { path }),
                             log: _log,
                             abortSignal: this.abortController.signal,
+                            registerEvents: registerDeviceEvents(coreContext),
                         },
                     })
                         .then(payload => {
diff --git a/packages/connect/src/core/onCallFirmwareUpdate.ts b/packages/connect/src/core/onCallFirmwareUpdate.ts
@@ -13,24 +13,13 @@ import { ERRORS, PROTO } from '../constants';
 import { getReleases } from '../data/firmwareInfo';
 import type { Device } from '../device/Device';
 import { DeviceList } from '../device/DeviceList';
-import { CoreEventMessage, DEVICE, UI, createDeviceMessage, createUiMessage } from '../events';
+import { CoreEventMessage, UI, createUiMessage } from '../events';
 import { CommonParams, DeviceUniquePath } from '../types';
 import { FirmwareUpdateResponse } from '../types/api/firmwareUpdate';
 import type { Log } from '../utils/debug';
 
 type PostMessage = (message: CoreEventMessage) => void;
 
-const registerEvents = (device: Device, postMessage: PostMessage) => {
-    device.on(DEVICE.BUTTON, ({ device: _, payload }) => {
-        postMessage(
-            createDeviceMessage(DEVICE.BUTTON, {
-                code: payload.code,
-                device: device.toMessageObject(),
-            }),
-        );
-    });
-};
-
 type ReconnectParams = {
     bootloader: boolean;
     method: 'wait' | 'auto' | 'manual';
@@ -40,14 +29,15 @@ type ReconnectParams = {
 type ReconnectContext = {
     deviceList: DeviceList;
     device: Device;
+    registerEvents: (device: Device) => void;
     postMessage: PostMessage;
     log: Log;
     abortSignal: AbortSignal;
 };
 
 const waitForReconnectedDevice = async (
     { bootloader, method, intermediary }: ReconnectParams,
-    { deviceList, device, postMessage, log, abortSignal }: ReconnectContext,
+    { deviceList, device, registerEvents, postMessage, log, abortSignal }: ReconnectContext,
 ): Promise<Device> => {
     const target = intermediary || !bootloader ? 'normal' : 'bootloader';
 
@@ -112,6 +102,7 @@ const waitForReconnectedDevice = async (
                 'we were unable to read device.features on the first interaction after seeing it, retrying...',
             );
             try {
+                registerEvents(reconnectedDevice);
                 // todo: it keeps printing warning "Previous call is still running" on reconnect from bl to normal
                 await reconnectedDevice.run(undefined, {
                     skipFirmwareChecks: true,
@@ -147,7 +138,7 @@ const waitForReconnectedDevice = async (
         throw ERRORS.TypedError('Method_Interrupted');
     }
 
-    registerEvents(reconnectedDevice, postMessage);
+    registerEvents(reconnectedDevice);
     await reconnectedDevice.waitForFirstRun();
 
     if (!reconnectedDevice.isUsedHere()) {
@@ -292,6 +283,7 @@ export type Params = {
 
 type Context = {
     deviceList: DeviceList;
+    registerEvents: (device: Device) => void;
     postMessage: PostMessage;
     initDevice: (path?: DeviceUniquePath) => Promise<Device>;
     log: Log;
@@ -305,8 +297,9 @@ type OnCallFirmwareUpdateParams = {
 
 export const onCallFirmwareUpdate = async ({
     params,
-    context: { deviceList, postMessage, initDevice, log, abortSignal },
+    context,
 }: OnCallFirmwareUpdateParams): Promise<FirmwareUpdateResponse> => {
+    const { deviceList, registerEvents, postMessage, initDevice, log } = context;
     log.debug('onCallFirmwareUpdate with params: ', params);
 
     const device = await initDevice(params?.device?.path);
@@ -319,7 +312,7 @@ export const onCallFirmwareUpdate = async ({
 
     log.debug('onCallFirmwareUpdate', 'device', device);
 
-    registerEvents(device, postMessage);
+    registerEvents(device);
 
     const { manual, upgrade, language, btcOnly } = getInstallationParams(device, params);
     log.debug('onCallFirmwareUpdate', 'installation params', {
@@ -351,7 +344,7 @@ export const onCallFirmwareUpdate = async ({
 
         reconnectedDevice = await waitForReconnectedDevice(
             { bootloader: true, method: 'manual' },
-            { deviceList, device, log, postMessage, abortSignal },
+            { ...context, device },
         );
     } else {
         // Device supports automatic reboot to bootloader, load translation data and do it
@@ -422,7 +415,7 @@ export const onCallFirmwareUpdate = async ({
         }
         reconnectedDevice = await waitForReconnectedDevice(
             { bootloader: true, method: 'auto' },
-            { deviceList, device, log, postMessage, abortSignal },
+            { ...context, device },
         );
     }
 
@@ -449,7 +442,7 @@ export const onCallFirmwareUpdate = async ({
 
         reconnectedDevice = await waitForReconnectedDevice(
             { bootloader: true, method: 'manual', intermediary: true },
-            { deviceList, device: reconnectedDevice, log, postMessage, abortSignal },
+            { ...context, device: reconnectedDevice },
         );
 
         binaryInfo = await getBinaryHelper(reconnectedDevice, params, log, postMessage, btcOnly);
@@ -467,7 +460,7 @@ export const onCallFirmwareUpdate = async ({
 
     reconnectedDevice = await waitForReconnectedDevice(
         { bootloader: false, method: 'wait' },
-        { deviceList, device: reconnectedDevice, log, postMessage, abortSignal },
+        { ...context, device: reconnectedDevice },
     );
 
     // features.firmware_present non-null value implies that device was initially connected with
