From 6513a3b61c81ae10154d8ec0c716ff1048a8112d Mon Sep 17 00:00:00 2001
From: katspaugh <katspaugh@gmail.com>
Date: Wed, 14 Jun 2023 12:47:19 +0200
Subject: [PATCH] Onclose callback

---
 .../common/NewModalDialog/index.tsx           |  7 ++--
 src/components/safe-apps/AppFrame/index.tsx   | 23 +++++-----
 src/components/tx-flow/index.tsx              | 42 +++++++++++--------
 3 files changed, 40 insertions(+), 32 deletions(-)

diff --git a/src/components/common/NewModalDialog/index.tsx b/src/components/common/NewModalDialog/index.tsx
index 0e0d54c8f9..e559b293d2 100644
--- a/src/components/common/NewModalDialog/index.tsx
+++ b/src/components/common/NewModalDialog/index.tsx
@@ -1,11 +1,10 @@
-import { type ReactElement, type ReactNode, useContext } from 'react'
+import { type ReactElement, type ReactNode } from 'react'
 import { IconButton, type ModalProps } from '@mui/material'
 import { Dialog, DialogTitle, type DialogProps, useMediaQuery } from '@mui/material'
 import { useTheme } from '@mui/material/styles'
 import ChainIndicator from '@/components/common/ChainIndicator'
 import CloseIcon from '@mui/icons-material/Close'
 import css from './styles.module.css'
-import { TxModalContext } from '@/components/tx-flow'
 
 interface ModalDialogProps extends DialogProps {
   dialogTitle?: React.ReactNode
@@ -47,12 +46,12 @@ const NewModalDialog = ({
   dialogTitle,
   hideChainIndicator,
   children,
+  onClose,
   fullScreen = false,
   ...restProps
 }: ModalDialogProps): ReactElement => {
   const theme = useTheme()
   const isSmallScreen = useMediaQuery(theme.breakpoints.down('md'))
-  const { setTxFlow } = useContext(TxModalContext)
 
   return (
     <Dialog
@@ -85,7 +84,7 @@ const NewModalDialog = ({
         <span style={{ flex: 1 }} />
         <IconButton
           aria-label="close"
-          onClick={() => setTxFlow(undefined)}
+          onClick={(e) => onClose?.(e, 'backdropClick')}
           size="small"
           sx={{
             ml: 2,
diff --git a/src/components/safe-apps/AppFrame/index.tsx b/src/components/safe-apps/AppFrame/index.tsx
index d91dbc00ef..fdc1ed9883 100644
--- a/src/components/safe-apps/AppFrame/index.tsx
+++ b/src/components/safe-apps/AppFrame/index.tsx
@@ -89,7 +89,17 @@ const AppFrame = ({ appUrl, allowedFeaturesList }: AppFrameProps): ReactElement
   const { getPermissions, hasPermission, permissionsRequest, setPermissionsRequest, confirmPermissionRequest } =
     useSafePermissions()
   const appName = useMemo(() => (remoteApp ? remoteApp.name : appUrl), [appUrl, remoteApp])
-  const { txFlow, setTxFlow } = useContext(TxModalContext)
+  const { setTxFlow } = useContext(TxModalContext)
+
+  const onTxFlowClose = () => {
+    setCurrentRequestId((prevId) => {
+      if (prevId) {
+        communicator?.send(CommunicatorMessages.REJECT_TRANSACTION_MESSAGE, prevId, true)
+        trackSafeAppEvent(SAFE_APPS_EVENTS.PROPOSE_TRANSACTION_REJECTED, appName)
+      }
+      return undefined
+    })
+  }
 
   const communicator = useAppCommunicator(iframeRef, remoteApp || safeAppFromManifest, chain, {
     onConfirmTransactions: (txs: BaseTransaction[], requestId: RequestId, params?: SendTransactionRequestParams) => {
@@ -102,8 +112,7 @@ const AppFrame = ({ appUrl, allowedFeaturesList }: AppFrameProps): ReactElement
       }
 
       setCurrentRequestId(requestId)
-
-      setTxFlow(<SafeAppsTxFlow data={data} />)
+      setTxFlow(<SafeAppsTxFlow data={data} />, onTxFlowClose)
     },
     onSignMessage: (
       message: string | EIP712TypedData,
@@ -177,14 +186,6 @@ const AppFrame = ({ appUrl, allowedFeaturesList }: AppFrameProps): ReactElement
     },
   })
 
-  // @TODO: figure out how to handle this
-  const onSafeAppsModalClose = () => {
-    if (currentRequestId) {
-      communicator?.send(CommunicatorMessages.REJECT_TRANSACTION_MESSAGE, currentRequestId, true)
-      trackSafeAppEvent(SAFE_APPS_EVENTS.PROPOSE_TRANSACTION_REJECTED, appName)
-    }
-  }
-
   const onAcceptPermissionRequest = (_origin: string, requestId: RequestId) => {
     const permissions = confirmPermissionRequest(PermissionStatus.GRANTED)
     communicator?.send(permissions, requestId as string)
diff --git a/src/components/tx-flow/index.tsx b/src/components/tx-flow/index.tsx
index c4d44d5ba9..03bb39ae77 100644
--- a/src/components/tx-flow/index.tsx
+++ b/src/components/tx-flow/index.tsx
@@ -1,33 +1,41 @@
-import {
-  createContext,
-  type Dispatch,
-  type ReactElement,
-  type ReactNode,
-  type SetStateAction,
-  useState,
-  useEffect,
-  useCallback,
-} from 'react'
+import { createContext, type ReactElement, type ReactNode, useState, useEffect, useCallback } from 'react'
 import NewModalDialog from '@/components/common/NewModalDialog'
 import { useRouter } from 'next/router'
 
 const noop = () => {}
 
-export const TxModalContext = createContext<{
+type TxModalContextType = {
   txFlow: ReactNode | undefined
-  setTxFlow: Dispatch<SetStateAction<ReactNode | undefined>>
-}>({
+  setTxFlow: (txFlow: TxModalContextType['txFlow'], onClose?: () => void) => void
+  onClose: () => void
+}
+
+export const TxModalContext = createContext<TxModalContextType>({
   txFlow: undefined,
   setTxFlow: noop,
+  onClose: noop,
 })
 
 export const TxModalProvider = ({ children }: { children: ReactNode }): ReactElement => {
-  const [txFlow, setTxFlow] = useState<ReactNode | undefined>()
+  const [txFlow, setFlow] = useState<TxModalContextType['txFlow']>(undefined)
+  const [onClose, setOnClose] = useState<TxModalContextType['onClose']>(noop)
   const router = useRouter()
 
   const handleModalClose = useCallback(() => {
-    setTxFlow(undefined)
-  }, [setTxFlow])
+    setOnClose((prevOnClose) => {
+      prevOnClose?.()
+      return noop
+    })
+    setFlow(undefined)
+  }, [setFlow, setOnClose])
+
+  const setTxFlow = useCallback(
+    (txFlow: TxModalContextType['txFlow'], onClose?: () => void) => {
+      setFlow(txFlow)
+      setOnClose(() => onClose ?? noop)
+    },
+    [setFlow, setOnClose],
+  )
 
   // Close the modal if user navigates
   useEffect(() => {
@@ -36,7 +44,7 @@ export const TxModalProvider = ({ children }: { children: ReactNode }): ReactEle
   }, [router, handleModalClose])
 
   return (
-    <TxModalContext.Provider value={{ txFlow, setTxFlow }}>
+    <TxModalContext.Provider value={{ txFlow, setTxFlow, onClose }}>
       {children}
 
       <NewModalDialog open={!!txFlow} onClose={handleModalClose}>