From e8fbc88177fef831d5d129f8ee4e93f0d5e6ac2a Mon Sep 17 00:00:00 2001
From: Jiri Zbytovsky <jiri.zbytovsky@satoshilabs.com>
Date: Wed, 4 Dec 2024 16:23:44 +0100
Subject: [PATCH] feat(suite): exonerate persisted failed FW by automatic hash
 check

---
 .../src/middlewares/wallet/storageMiddleware.ts     |  2 ++
 suite-common/firmware/src/firmwareReducer.ts        | 13 +++++++++++++
 suite-common/firmware/src/firmwareThunks.ts         |  2 +-
 suite-common/suite-utils/src/device.ts              |  6 +++++-
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/packages/suite/src/middlewares/wallet/storageMiddleware.ts b/packages/suite/src/middlewares/wallet/storageMiddleware.ts
index 653e3e436e9..6018a963b8d 100644
--- a/packages/suite/src/middlewares/wallet/storageMiddleware.ts
+++ b/packages/suite/src/middlewares/wallet/storageMiddleware.ts
@@ -311,6 +311,8 @@ const storageMiddleware = (api: MiddlewareAPI<Dispatch, AppState>) => {
                 case FORM_DRAFT.REMOVE_DRAFT:
                     storageActions.removeFormDraft(action.key);
                     break;
+
+                case deviceActions.connectDevice.type: // so that firmwareReducer.addCase for the same action is persisted
                 case firmwareActions.clearInvalidHash.type:
                 case firmwareActions.setHashInvalid.type:
                     api.dispatch(storageActions.saveFirmware());
diff --git a/suite-common/firmware/src/firmwareReducer.ts b/suite-common/firmware/src/firmwareReducer.ts
index 746ca580e17..c0328871c63 100644
--- a/suite-common/firmware/src/firmwareReducer.ts
+++ b/suite-common/firmware/src/firmwareReducer.ts
@@ -9,6 +9,7 @@ import {
 } from '@trezor/connect';
 import { createReducerWithExtraDeps } from '@suite-common/redux-utils';
 import { deviceActions } from '@suite-common/wallet-core';
+import { isDeviceKnown } from '@suite-common/suite-utils';
 
 import { firmwareActions } from './firmwareActions';
 
@@ -86,6 +87,18 @@ export const prepareFirmwareReducer = createReducerWithExtraDeps(initialState, (
             state.cachedDevice = payload;
         })
         .addCase(deviceActions.addButtonRequest, extra.reducers.addButtonRequestFirmware)
+        .addCase(deviceActions.connectDevice, (state, { payload: { device } }) => {
+            if (!isDeviceKnown(device)) return;
+
+            // use the automatic hash check to clear device if it hasn't passed hash check done after firmware update
+            // otherwise it'd be stuck in "error" state until next firmware update
+            // in `storageMiddleware` this is persisted into storage (on the same action type)
+            if (device.authenticityChecks?.firmwareHash?.success) {
+                state.firmwareHashInvalid = state.firmwareHashInvalid.filter(
+                    deviceId => deviceId !== device.id,
+                );
+            }
+        })
         .addMatcher(
             action => action.type === extra.actionTypes.storageLoad,
             extra.reducers.storageLoadFirmware,
diff --git a/suite-common/firmware/src/firmwareThunks.ts b/suite-common/firmware/src/firmwareThunks.ts
index 8f610e06ebe..c5e44e062a9 100644
--- a/suite-common/firmware/src/firmwareThunks.ts
+++ b/suite-common/firmware/src/firmwareThunks.ts
@@ -158,7 +158,7 @@ export const firmwareUpdate = createThunk<
                 // hash check was performed, and it does not match, so consider firmware counterfeit
                 dispatch(handleFwHashMismatch(device));
             } else if (check === 'other-error') {
-                // device failed to respond to the hash check, consider the firmware counterfeit
+                // device failed to respond, so display a warning in the modal
                 dispatch(
                     handleFwHashError({
                         errorMessage: firmwareUpdateResponse.payload.checkError,
diff --git a/suite-common/suite-utils/src/device.ts b/suite-common/suite-utils/src/device.ts
index eba60b1144f..a2a3bdecc6d 100644
--- a/suite-common/suite-utils/src/device.ts
+++ b/suite-common/suite-utils/src/device.ts
@@ -1,4 +1,4 @@
-import { Device, UnavailableCapability, DeviceModelInternal } from '@trezor/connect';
+import { Device, UnavailableCapability, DeviceModelInternal, KnownDevice } from '@trezor/connect';
 import { TrezorDevice, AcquiredDevice } from '@suite-common/suite-types';
 import * as URLS from '@trezor/urls';
 
@@ -119,6 +119,10 @@ export const isDeviceAccessible = (device?: TrezorDevice) => {
     return device.mode === 'normal' && device.firmware !== 'required';
 };
 
+// useful for working with `Device` type straight from connect, which is missing `ExtendedDevice` properties (required on `TrezorDevice`)
+export const isDeviceKnown = (device?: Device): device is KnownDevice =>
+    device?.type === 'acquired';
+
 export const isDeviceAcquired = (device?: TrezorDevice): device is AcquiredDevice =>
     !!device?.features;