From 908a97800d64ee712995d7c3db7a3335e90232e8 Mon Sep 17 00:00:00 2001
From: Edgar Khanzadian <edgar.khanzadian@gmail.com>
Date: Fri, 2 Feb 2024 11:23:47 +0400
Subject: [PATCH] refactor: useLedgerSignTx should be chain agnostic

---
 .../ledger-sign-stacks-tx-container.tsx       | 20 ++++++++++++-
 .../tx-signing/use-ledger-sign-tx.ts          | 28 ++++++-------------
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/src/app/features/ledger/flows/stacks-tx-signing/ledger-sign-stacks-tx-container.tsx b/src/app/features/ledger/flows/stacks-tx-signing/ledger-sign-stacks-tx-container.tsx
index 0e8b9835f9..b7acdf2a36 100644
--- a/src/app/features/ledger/flows/stacks-tx-signing/ledger-sign-stacks-tx-container.tsx
+++ b/src/app/features/ledger/flows/stacks-tx-signing/ledger-sign-stacks-tx-container.tsx
@@ -1,5 +1,5 @@
 import { useEffect, useState } from 'react';
-import { Route, useLocation } from 'react-router-dom';
+import { Route, useLocation, useNavigate } from 'react-router-dom';
 
 import { deserializeTransaction } from '@stacks/transactions';
 import StacksApp, { LedgerError } from '@zondax/ledger-stacks';
@@ -15,6 +15,7 @@ import {
   connectLedgerStacksApp,
   getStacksAppVersion,
   isStacksAppOpen,
+  isVersionOfLedgerStacksAppWithContractPrincipalBug,
   signLedgerStacksTransaction,
   signStacksTransactionWithSignature,
 } from '@app/features/ledger/utils/stacks-ledger-utils';
@@ -43,6 +44,7 @@ function LedgerSignStacksTxContainer() {
   const account = useCurrentStacksAccount();
   const verifyLedgerPublicKey = useVerifyMatchingLedgerStacksPublicKey();
   const [unsignedTx, setUnsignedTx] = useState<null | string>(null);
+  const navigate = useNavigate();
 
   const chain = 'stacks' as const;
 
@@ -63,6 +65,22 @@ function LedgerSignStacksTxContainer() {
     isAppOpen: isStacksAppOpen,
     getAppVersion: getStacksAppVersion,
     connectApp: connectLedgerStacksApp,
+    async passesAdditionalVersionCheck(appVersion) {
+      if (appVersion.chain !== 'stacks') {
+        return true;
+      }
+
+      if (isVersionOfLedgerStacksAppWithContractPrincipalBug(appVersion)) {
+        navigate(RouteUrls.LedgerOutdatedAppWarning);
+        const response = await hasUserSkippedBuggyAppWarning.wait();
+
+        if (response === 'cancelled-operation') {
+          ledgerNavigate.cancelLedgerAction();
+        }
+        return false;
+      }
+      return true;
+    },
     async signTransactionWithDevice(stacksApp) {
       // TODO: need better handling
       if (!account) return;
diff --git a/src/app/features/ledger/generic-flows/tx-signing/use-ledger-sign-tx.ts b/src/app/features/ledger/generic-flows/tx-signing/use-ledger-sign-tx.ts
index b77dfa5eda..24e215f3e9 100644
--- a/src/app/features/ledger/generic-flows/tx-signing/use-ledger-sign-tx.ts
+++ b/src/app/features/ledger/generic-flows/tx-signing/use-ledger-sign-tx.ts
@@ -1,11 +1,9 @@
 import { useMemo, useState } from 'react';
-import { useNavigate } from 'react-router-dom';
 
 import StacksApp from '@zondax/ledger-stacks';
 import BitcoinApp from 'ledger-bitcoin';
 
 import { SupportedBlockchains } from '@shared/constants';
-import { RouteUrls } from '@shared/route-urls';
 
 import { delay } from '@app/common/utils';
 
@@ -16,10 +14,7 @@ import {
   checkLockedDeviceError,
   useLedgerResponseState,
 } from '../../utils/generic-ledger-utils';
-import {
-  StacksAppVersion,
-  isVersionOfLedgerStacksAppWithContractPrincipalBug,
-} from '../../utils/stacks-ledger-utils';
+import { StacksAppVersion } from '../../utils/stacks-ledger-utils';
 import { createWaitForUserToSeeWarningScreen } from './ledger-sign-tx.context';
 
 interface UseLedgerSignTxArgs<App extends BitcoinApp | StacksApp> {
@@ -27,6 +22,7 @@ interface UseLedgerSignTxArgs<App extends BitcoinApp | StacksApp> {
   isAppOpen({ name }: { name: string }): boolean;
   getAppVersion(app: App): Promise<StacksAppVersion> | Promise<BitcoinAppVersion>;
   connectApp(): Promise<App>;
+  passesAdditionalVersionCheck?(appVersion: StacksAppVersion | BitcoinAppVersion): Promise<boolean>;
   onSuccess?(): void;
   signTransactionWithDevice(app: App): Promise<void>;
 }
@@ -38,31 +34,24 @@ export function useLedgerSignTx<App extends StacksApp | BitcoinApp>({
   connectApp,
   onSuccess,
   signTransactionWithDevice,
+  passesAdditionalVersionCheck,
 }: UseLedgerSignTxArgs<App>) {
   const [outdatedAppVersionWarning, setAppVersionOutdatedWarning] = useState(false);
   const [latestDeviceResponse, setLatestDeviceResponse] = useLedgerResponseState();
   const [awaitingDeviceConnection, setAwaitingDeviceConnection] = useState(false);
   const ledgerNavigate = useLedgerNavigate();
   const hasUserSkippedBuggyAppWarning = useMemo(() => createWaitForUserToSeeWarningScreen(), []);
-  const navigate = useNavigate();
   async function checkCorrectAppIsOpenWithFailState(app: App) {
     const response = await getAppVersion(app);
     if (!isAppOpen({ name: response.name })) {
       setAwaitingDeviceConnection(false);
       throw new Error(LedgerConnectionErrors.AppNotOpen);
     }
-    if (response.chain === 'stacks') {
-      if (isVersionOfLedgerStacksAppWithContractPrincipalBug(response)) {
-        navigate(RouteUrls.LedgerOutdatedAppWarning);
-        const response = await hasUserSkippedBuggyAppWarning.wait();
-
-        if (response === 'cancelled-operation') {
-          ledgerNavigate.cancelLedgerAction();
-          return;
-        }
-      }
+    const passedAdditionalVersionCheck = await passesAdditionalVersionCheck?.(response);
+    if (passedAdditionalVersionCheck) {
+      return response;
     }
-    return response;
+    return;
   }
 
   async function signTransactionImpl() {
@@ -85,7 +74,8 @@ export function useLedgerSignTx<App extends StacksApp | BitcoinApp>({
       }
 
       ledgerNavigate.toErrorStep(chain);
-      return app?.transport.close();
+    } finally {
+      await app?.transport.close();
     }
   }