Feature: Safenet balance in the assets tab (#4345)

* show safenet balances in the assets tab

* **feat(hooks/loadables): add support for SafeNet balances**

- Remove unused `useLazyGetSafeNetBalanceQuery`
- Fetch SafeNet balances alongside Counterfactual Gateway balances
- Merge and return combined balance data when SafeNet is supported

* mock safenet config in tests

src/components/tx-flow/flows/EnableSafenet/ReviewEnableSafenet.tsx

@@ -19,7 +19,7 @@ export const ReviewEnableSafenet = ({ params }: { params: EnableSafenetFlowProps
   useEffect(() => {
     async function getTxs(): Promise<SafeTransaction> {
       if (params.tokensForPresetAllowances.length === 0) {
-        return await createEnableGuardTx(params.guardAddress)
+        return createEnableGuardTx(params.guardAddress)
       }
 
       const txs: MetaTransactionData[] = [(await createEnableGuardTx(params.guardAddress)).data]

src/hooks/__tests__/useLoadBalances.test.ts

@@ -1,4 +1,5 @@
 import * as store from '@/store'
+import * as safeNetStore from '@/store/safenet'
 import { defaultSafeInfo } from '@/store/safeInfoSlice'
 import { act, renderHook, waitFor } from '@/tests/test-utils'
 import { toBeHex } from 'ethers'
@@ -97,6 +98,35 @@ describe('useLoadBalances', () => {
     jest.clearAllMocks()
     localStorage.clear()
     jest.spyOn(useChainId, 'useChainId').mockReturnValue('5')
+    jest.spyOn(safeNetStore, 'useGetSafeNetConfigQuery').mockReturnValue({
+      data: {
+        chains: {
+          sources: [11155111],
+          destinations: [84532, 11155420],
+        },
+        guards: {
+          '84532': '0x865544E0599589BA604b0449858695937d571382',
+          '11155111': '0x865544E0599589BA604b0449858695937d571382',
+          '11155420': '0x865544E0599589BA604b0449858695937d571382',
+        },
+        tokens: {
+          USDC: {
+            '84532': '0x036CbD53842c5426634e7929541eC2318f3dCF7e',
+            '11155111': '0x1c7D4B196Cb0C7B01d743Fbc6116a902379C7238',
+            '11155420': '0x5fd84259d66Cd46123540766Be93DFE6D43130D7',
+          },
+          DAI: {
+            '84532': '0xE6F6e27c0BF1a4841E3F09d03D7D31Da8eAd0a27',
+            '11155111': '0xB4F1737Af37711e9A5890D9510c9bB60e170CB0D',
+            '11155420': '0x0091f4e75a03C11cB9be8E3717219005eb780D89',
+          },
+        },
+      },
+      isLoading: false,
+      error: undefined,
+      isSuccess: true,
+      refetch: jest.fn(),
+    })
   })
 
   test('without selected Safe', async () => {

src/hooks/loadables/useLoadBalances.ts

@@ -11,6 +11,9 @@ import { FEATURES, hasFeature } from '@/utils/chains'
 import { POLLING_INTERVAL } from '@/config/constants'
 import useIntervalCounter from '../useIntervalCounter'
 import useSafeInfo from '../useSafeInfo'
+import { useGetSafeNetConfigQuery } from '@/store/safenet'
+import { SafenetChainType, convertSafeNetBalanceToSafeClientGatewayBalance, isSupportedChain } from '@/utils/safenet'
+import { getSafeNetBalances } from '@/store/safenet'
 
 export const useTokenListSetting = (): boolean | undefined => {
   const chain = useCurrentChain()
@@ -24,33 +27,76 @@ export const useTokenListSetting = (): boolean | undefined => {
   return isTrustedTokenList
 }
 
+const mergeBalances = (balance1: SafeBalanceResponse, balance2: SafeBalanceResponse): SafeBalanceResponse => {
+  return {
+    fiatTotal: balance1.fiatTotal + balance2.fiatTotal,
+    items: [...balance1.items, ...balance2.items],
+  }
+}
+
 export const useLoadBalances = (): AsyncResult<SafeBalanceResponse> => {
   const [pollCount, resetPolling] = useIntervalCounter(POLLING_INTERVAL)
+  const {
+    data: safeNetConfig,
+    isSuccess: isSafeNetConfigSuccess,
+    isLoading: isSafeNetConfigLoading,
+  } = useGetSafeNetConfigQuery()
   const currency = useAppSelector(selectCurrency)
   const isTrustedTokenList = useTokenListSetting()
   const { safe, safeAddress } = useSafeInfo()
   const web3 = useWeb3()
   const chain = useCurrentChain()
   const chainId = safe.chainId
+  const chainSupportedBySafeNet =
+    isSafeNetConfigSuccess && isSupportedChain(Number(chainId), safeNetConfig, SafenetChainType.DESTINATION)
 
   // Re-fetch assets when the entire SafeInfo updates
   const [data, error, loading] = useAsync<SafeBalanceResponse | undefined>(
     () => {
-      if (!chainId || !safeAddress || isTrustedTokenList === undefined) return
+      if (!chainId || !safeAddress || isTrustedTokenList === undefined || isSafeNetConfigLoading) return
 
       if (!safe.deployed) {
         return getCounterfactualBalance(safeAddress, web3, chain)
       }
 
-      return getBalances(chainId, safeAddress, currency, {
-        trusted: isTrustedTokenList,
-      })
+      const balanceQueries = [
+        getBalances(chainId, safeAddress, currency, {
+          trusted: isTrustedTokenList,
+        }),
+      ]
+
+      if (isSafeNetConfigSuccess && chainSupportedBySafeNet) {
+        balanceQueries.push(
+          getSafeNetBalances(chainId, safeAddress)
+            .then((safenetBalances) =>
+              convertSafeNetBalanceToSafeClientGatewayBalance(safenetBalances, safeNetConfig, Number(chainId)),
+            )
+            .catch(() => ({
+              fiatTotal: '0',
+              items: [],
+            })),
+        )
+      }
+
+      return Promise.all(balanceQueries).then(([cgw, sn]) => (sn ? mergeBalances(cgw, sn) : cgw))
     },
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [safeAddress, chainId, currency, isTrustedTokenList, pollCount, safe.deployed, web3, chain],
+    [
+      safeAddress,
+      chainId,
+      currency,
+      isTrustedTokenList,
+      pollCount,
+      safe.deployed,
+      web3,
+      chain,
+      safeNetConfig,
+      isSafeNetConfigSuccess,
+      isSafeNetConfigLoading,
+      chainSupportedBySafeNet,
+    ],
     false, // don't clear data between polls
   )
-
   // Reset the counter when safe address/chainId changes
   useEffect(() => {
     resetPolling()
@@ -63,7 +109,7 @@ export const useLoadBalances = (): AsyncResult<SafeBalanceResponse> => {
     }
   }, [error])
 
-  return [data, error, loading]
+  return [data, error, loading || isSafeNetConfigLoading]
 }
 
 export default useLoadBalances

src/pages/settings/safenet.tsx

@@ -10,26 +10,23 @@ import { sameAddress } from '@/utils/addresses'
 import { useContext, useEffect, useMemo } from 'react'
 import { TxModalContext } from '@/components/tx-flow'
 import { EnableSafenetFlow } from '@/components/tx-flow/flows/EnableSafenet'
-import type { SafenetConfigEntity } from '@/store/safenet'
+import type { SafeNetConfigEntity } from '@/store/safenet'
 import {
-  useGetSafenetConfigQuery,
   useLazyGetSafeNetOffchainStatusQuery,
   useRegisterSafeNetMutation,
+  useGetSafeNetConfigQuery,
 } from '@/store/safenet'
 import type { ExtendedSafeInfo } from '@/store/safeInfoSlice'
 import { SAFE_FEATURES } from '@safe-global/protocol-kit/dist/src/utils'
 import { hasSafeFeature } from '@/utils/safe-versions'
+import { SafenetChainType, isSupportedChain } from '@/utils/safenet'
 
-const isSupportedChain = (chainId: number, safenetConfig: SafenetConfigEntity) => {
-  return safenetConfig.chains.sources.includes(chainId)
-}
-
-const getSafenetTokensByChain = (chainId: number, safenetConfig: SafenetConfigEntity): string[] => {
-  const tokenSymbols = Object.keys(safenetConfig.tokens)
+const getSafeNetTokensByChain = (chainId: number, safeNetConfig: SafeNetConfigEntity): string[] => {
+  const tokenSymbols = Object.keys(safeNetConfig.tokens)
 
   const tokens: string[] = []
   for (const symbol of tokenSymbols) {
-    const tokenAddress = safenetConfig.tokens[symbol][chainId]
+    const tokenAddress = safeNetConfig.tokens[symbol][chainId]
     if (tokenAddress) {
       tokens.push(tokenAddress)
     }
@@ -38,11 +35,11 @@ const getSafenetTokensByChain = (chainId: number, safenetConfig: SafenetConfigEn
   return tokens
 }
 
-const SafeNetContent = ({ safenetConfig, safe }: { safenetConfig: SafenetConfigEntity; safe: ExtendedSafeInfo }) => {
+const SafeNetContent = ({ safeNetConfig, safe }: { safeNetConfig: SafeNetConfigEntity; safe: ExtendedSafeInfo }) => {
   const isVersionWithGuards = hasSafeFeature(SAFE_FEATURES.SAFE_TX_GUARDS, safe.version)
-  const safenetGuardAddress = safenetConfig.guards[safe.chainId]
-  const isSafeNetGuardEnabled = isVersionWithGuards && sameAddress(safe.guard?.value, safenetGuardAddress)
-  const chainSupported = isSupportedChain(Number(safe.chainId), safenetConfig)
+  const safeNetGuardAddress = safeNetConfig.guards[safe.chainId]
+  const isSafeNetGuardEnabled = isVersionWithGuards && sameAddress(safe.guard?.value, safeNetGuardAddress)
+  const chainSupported = isSupportedChain(Number(safe.chainId), safeNetConfig, SafenetChainType.SOURCE)
   const { setTxFlow } = useContext(TxModalContext)
 
   // Lazy query because running it on unsupported chain throws an error
@@ -59,16 +56,16 @@ const SafeNetContent = ({ safenetConfig, safe }: { safenetConfig: SafenetConfigE
   // @ts-expect-error bad types. We don't want 404 to be an error - it just means that the safe is not registered
   const offchainLookupError = safeNetOffchainStatusError?.status === 404 ? null : safeNetOffchainStatusError
   const registeredOffchainStatus =
-    !offchainLookupError && sameAddress(safeNetOffchainStatus?.guard, safenetGuardAddress)
+    !offchainLookupError && sameAddress(safeNetOffchainStatus?.guard, safeNetGuardAddress)
 
-  const safenetStatusQueryWorked =
+  const safeNetStatusQueryWorked =
     safeNetOffchainStatusStatus === QueryStatus.fulfilled || safeNetOffchainStatusStatus === QueryStatus.rejected
-  const needsRegistration = safenetStatusQueryWorked && isSafeNetGuardEnabled && !registeredOffchainStatus
+  const needsRegistration = safeNetStatusQueryWorked && isSafeNetGuardEnabled && !registeredOffchainStatus
   const [registerSafeNet, { error: registerSafeNetError }] = useRegisterSafeNetMutation()
   const error = offchainLookupError || registerSafeNetError
-  const safenetAssets = useMemo(
-    () => getSafenetTokensByChain(Number(safe.chainId), safenetConfig),
-    [safe.chainId, safenetConfig],
+  const safeNetAssets = useMemo(
+    () => getSafeNetTokensByChain(Number(safe.chainId), safeNetConfig),
+    [safe.chainId, safeNetConfig],
   )
 
   if (error) {
@@ -92,7 +89,7 @@ const SafeNetContent = ({ safenetConfig, safe }: { safenetConfig: SafenetConfigE
       return (
         <Typography>
           SafeNet is not supported on this chain. List of supported chains ids:{' '}
-          {safenetConfig.chains.sources.join(', ')}
+          {safeNetConfig.chains.sources.join(', ')}
         </Typography>
       )
     case !isVersionWithGuards:
@@ -107,7 +104,7 @@ const SafeNetContent = ({ safenetConfig, safe }: { safenetConfig: SafenetConfigE
             variant="contained"
             onClick={() =>
               setTxFlow(
-                <EnableSafenetFlow guardAddress={safenetGuardAddress} tokensForPresetAllowances={safenetAssets} />,
+                <EnableSafenetFlow guardAddress={safeNetGuardAddress} tokensForPresetAllowances={safeNetAssets} />,
               )
             }
             sx={{ mt: 2 }}
@@ -125,22 +122,22 @@ const SafeNetContent = ({ safenetConfig, safe }: { safenetConfig: SafenetConfigE
 
 const SafeNetPage: NextPage = () => {
   const { safe, safeLoaded } = useSafeInfo()
-  const { data: safenetConfig, isLoading: safenetConfigLoading, error: safenetConfigError } = useGetSafenetConfigQuery()
+  const { data: safeNetConfig, isLoading: safeNetConfigLoading, error: safeNetConfigError } = useGetSafeNetConfigQuery()
 
-  if (!safeLoaded || safenetConfigLoading) {
+  if (!safeLoaded || safeNetConfigLoading) {
     return <CircularProgress />
   }
 
-  if (safenetConfigError) {
+  if (safeNetConfigError) {
     return <Typography>Error loading SafeNet config</Typography>
   }
 
-  if (!safenetConfig) {
+  if (!safeNetConfig) {
     // Should never happen, making TS happy
     return <Typography>No SafeNet config found</Typography>
   }
 
-  const safenetContent = <SafeNetContent safenetConfig={safenetConfig} safe={safe} />
+  const safeNetContent = <SafeNetContent safeNetConfig={safeNetConfig} safe={safe} />
 
   return (
     <>
@@ -174,7 +171,7 @@ const SafeNetPage: NextPage = () => {
             </Grid>
 
             <Grid item xs>
-              {safenetContent}
+              {safeNetContent}
             </Grid>
           </Grid>
         </Paper>

src/store/safenet.ts

@@ -2,13 +2,13 @@ import { createApi, fetchBaseQuery } from '@reduxjs/toolkit/query/react'
 
 import { SAFENET_API_URL } from '@/config/constants'
 
-export type SafenetSafeEntity = {
+export type SafeNetSafeEntity = {
   safe: string
   chainId: number
   guard: string
 }
 
-export type SafenetConfigEntity = {
+export type SafeNetConfigEntity = {
   chains: {
     sources: number[]
     destinations: number[]
@@ -17,15 +17,26 @@ export type SafenetConfigEntity = {
   tokens: Record<string, Record<string, string>>
 }
 
+export type SafeNetBalanceEntity = {
+  [tokenSymbol: string]: string
+}
+
+export const getSafeNetBalances = async (chainId: string, safeAddress: string): Promise<SafeNetBalanceEntity> => {
+  const response = await fetch(`${SAFENET_API_URL}/safenet/balances/${chainId}/${safeAddress}`)
+  const data = await response.json()
+  return data
+}
+
 export const safenetApi = createApi({
   reducerPath: 'safenetApi',
   baseQuery: fetchBaseQuery({ baseUrl: `${SAFENET_API_URL}/safenet` }),
-  tagTypes: ['SafeNetOffchainStatus'],
+  tagTypes: ['SafeNetConfig', 'SafeNetOffchainStatus', 'SafeNetBalance'],
   endpoints: (builder) => ({
-    getSafenetConfig: builder.query<SafenetConfigEntity, void>({
+    getSafeNetConfig: builder.query<SafeNetConfigEntity, void>({
       query: () => `/config/`,
+      providesTags: ['SafeNetConfig'],
     }),
-    getSafeNetOffchainStatus: builder.query<SafenetSafeEntity, { chainId: string; safeAddress: string }>({
+    getSafeNetOffchainStatus: builder.query<SafeNetSafeEntity, { chainId: string; safeAddress: string }>({
       query: ({ chainId, safeAddress }) => `/safes/${chainId}/${safeAddress}`,
       providesTags: (_, __, arg) => [{ type: 'SafeNetOffchainStatus', id: arg.safeAddress }],
     }),
@@ -40,7 +51,16 @@ export const safenetApi = createApi({
       }),
       invalidatesTags: (_, __, arg) => [{ type: 'SafeNetOffchainStatus', id: arg.safeAddress }],
     }),
+    getSafeNetBalance: builder.query<SafeNetBalanceEntity, { chainId: string; safeAddress: string }>({
+      query: ({ chainId, safeAddress }) => `/balances/${chainId}/${safeAddress}`,
+      providesTags: (_, __, arg) => [{ type: 'SafeNetBalance', id: arg.safeAddress }],
+    }),
   }),
 })
 
-export const { useLazyGetSafeNetOffchainStatusQuery, useRegisterSafeNetMutation, useGetSafenetConfigQuery } = safenetApi
+export const {
+  useLazyGetSafeNetOffchainStatusQuery,
+  useRegisterSafeNetMutation,
+  useGetSafeNetConfigQuery,
+  useLazyGetSafeNetBalanceQuery,
+} = safenetApi

src/utils/safenet.ts

@@ -0,0 +1,49 @@
+import type { SafeNetBalanceEntity, SafeNetConfigEntity } from '@/store/safenet'
+import { TokenType, type SafeBalanceResponse } from '@safe-global/safe-gateway-typescript-sdk'
+
+const enum SafenetChainType {
+  SOURCE = 'source',
+  DESTINATION = 'destination',
+}
+
+const isSupportedChain = (chainId: number, safeNetConfig: SafeNetConfigEntity, chainType: SafenetChainType) => {
+  return chainType === SafenetChainType.SOURCE
+    ? safeNetConfig.chains.sources.includes(chainId)
+    : safeNetConfig.chains.destinations.includes(chainId)
+}
+
+const convertSafeNetBalanceToSafeClientGatewayBalance = (
+  safeNetBalance: SafeNetBalanceEntity,
+  safeNetConfig: SafeNetConfigEntity,
+  chainId: number,
+): SafeBalanceResponse => {
+  const balances: SafeBalanceResponse = {
+    fiatTotal: safeNetBalance.fiatTotal,
+    items: [],
+  }
+
+  for (const [tokenName, balance] of Object.entries(safeNetBalance)) {
+    const tokenAddress = safeNetConfig.tokens[tokenName][chainId]
+    if (!tokenAddress) {
+      continue
+    }
+
+    balances.items.push({
+      tokenInfo: {
+        type: TokenType.ERC20,
+        address: tokenAddress,
+        decimals: 18,
+        symbol: tokenName,
+        name: `${tokenName} (SafeNet)`,
+        logoUri: '',
+      },
+      balance,
+      fiatBalance: '0',
+      fiatConversion: '0',
+    })
+  }
+
+  return balances
+}
+
+export { isSupportedChain, SafenetChainType, convertSafeNetBalanceToSafeClientGatewayBalance }