From 6900f8fe03f7fc394eb5b1f4e968cdeb5d06396a Mon Sep 17 00:00:00 2001
From: brusher_ru <brusher.designer@gmail.com>
Date: Mon, 16 Sep 2024 13:18:27 -0300
Subject: [PATCH] feat: check account balances before publishing tx

---
 src/api/schemas/strNumber.ts                |  2 +-
 src/components/sendTx/ConfirmationModal.tsx | 20 +++++++--
 src/components/sendTx/SendTxModal.tsx       | 45 +++++++++++++++++++++
 3 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/api/schemas/strNumber.ts b/src/api/schemas/strNumber.ts
index 2626b67..759a0a0 100644
--- a/src/api/schemas/strNumber.ts
+++ b/src/api/schemas/strNumber.ts
@@ -10,6 +10,6 @@ export const BigIntStringSchema = z.custom<string>((a) => {
 });
 
 export const BigIntMin = (min: bigint) =>
-  z.custom<string>((a) => BigInt(a) >= min, {
+  z.custom<string>((a) => a && BigInt(a) >= min, {
     message: `Should be greater than or equal to ${String(min)}`,
   });
diff --git a/src/components/sendTx/ConfirmationModal.tsx b/src/components/sendTx/ConfirmationModal.tsx
index b9dc7c9..e129cbd 100644
--- a/src/components/sendTx/ConfirmationModal.tsx
+++ b/src/components/sendTx/ConfirmationModal.tsx
@@ -77,6 +77,7 @@ type ConfirmationModalProps = ConfirmationData & {
   isOpen: boolean;
   estimatedGas: bigint | null;
   isLedgerRejected?: boolean;
+  insufficientFunds: string;
 };
 
 const renderTemplateSpecificFields = (form: FormValues) => {
@@ -222,6 +223,7 @@ function ConfirmationModal({
   signatures = undefined,
   required = 1,
   isMultiSig = false,
+  insufficientFunds,
   isOpen,
   onClose,
   onSubmit,
@@ -320,7 +322,13 @@ function ConfirmationModal({
     if (!isMultiSig) {
       return (
         <ButtonGroup isAttached>
-          <Button variant="whiteModal" onClick={submit} ml={2} mr="1px">
+          <Button
+            variant="whiteModal"
+            onClick={submit}
+            ml={2}
+            mr="1px"
+            isDisabled={insufficientFunds.length > 0}
+          >
             {hasSingleSig ? 'Publish' : 'Sign & Publish'}
           </Button>
           <Menu>
@@ -349,10 +357,11 @@ function ConfirmationModal({
     return (
       <ButtonGroup isAttached>
         <Button
-          colorScheme="blue"
+          variant="whiteModal"
           onClick={mayPublish ? submit : exportSigned}
           ml={2}
           mr="1px"
+          isDisabled={mayPublish ? insufficientFunds.length > 0 : false}
         >
           {/* eslint-disable-next-line no-nested-ternary */}
           {hasAllSignatures
@@ -365,7 +374,7 @@ function ConfirmationModal({
           <MenuButton
             as={IconButton}
             icon={<IconChevronDown />}
-            colorScheme="blue"
+            variant="whiteModal"
             minW={8}
           />
           <MenuList>
@@ -486,6 +495,11 @@ function ConfirmationModal({
               )}
             </>
           )}
+          {!!insufficientFunds && (
+            <Text mt={2} color="brand.red" fontSize="xs">
+              {insufficientFunds}
+            </Text>
+          )}
         </ModalBody>
         <ModalFooter
           justifyContent="space-between"
diff --git a/src/components/sendTx/SendTxModal.tsx b/src/components/sendTx/SendTxModal.tsx
index d31f5b7..ac9914f 100644
--- a/src/components/sendTx/SendTxModal.tsx
+++ b/src/components/sendTx/SendTxModal.tsx
@@ -71,6 +71,7 @@ import {
 } from '../../utils/account';
 import { fromBase64 } from '../../utils/base64';
 import { generateAddress, getWords } from '../../utils/bech32';
+import { DEFAULT_ACCOUNT_STATES } from '../../utils/constants';
 import { fromHexString, toHexString } from '../../utils/hexString';
 import { isForeignKey } from '../../utils/keys';
 import { formatSmidge } from '../../utils/smh';
@@ -92,11 +93,13 @@ import ConfirmationModal, { ConfirmationData } from './ConfirmationModal';
 import Drain from './Drain';
 import MultiSigSpawn from './MultiSigSpawn';
 import {
+  DrainPayload,
   DrainSchema,
   FormSchema,
   FormValues,
   MultiSigSpawnSchema,
   SingleSigSpawnSchema,
+  SpendPayload,
   SpendSchema,
   VaultSpawnSchema,
 } from './schemas';
@@ -194,6 +197,47 @@ function SendTxModal({ isOpen, onClose }: SendTxModalProps): JSX.Element {
     isVaultAccount
   );
 
+  const insufficientFunds = O.mapWithDefault(accountState, '', (acc) => {
+    if (!estimatedGas) {
+      return 'Waiting for gas estimation...';
+    }
+
+    const tx = getValues();
+    const fee =
+      estimatedGas *
+      BigInt(Number.isNaN(tx.gasPrice) || !tx.gasPrice ? 1 : tx.gasPrice);
+    if (BigInt(acc.projected.balance) < fee) {
+      return 'You have insufficient funds to pay for gas';
+    }
+    if (SpendSchema.safeParse(tx.payload).success) {
+      // If it is a Spend tx — check if the balance is enough to send the amount
+      const txPayload = tx.payload as SpendPayload;
+      const res =
+        BigInt(acc.projected.balance) >= BigInt(txPayload.Amount) + fee;
+      return res ? '' : 'You have insufficient funds to send this amount';
+    }
+    if (DrainSchema.safeParse(tx.payload).success) {
+      // If it is a Drain tx — check that Vault has enough balance to be drained
+      const txPayload = tx.payload as DrainPayload;
+      return pipe(
+        genesisID,
+        O.flatMap((genId) => getAccountData(genId, txPayload.Vault)),
+        O.mapWithDefault('', (vault) => {
+          if (vault.state === DEFAULT_ACCOUNT_STATES) {
+            // No account data found — do not block the transaction
+            return '';
+          }
+
+          const res =
+            BigInt(vault.state.projected.balance) >= BigInt(txPayload.Amount);
+          return res ? '' : 'The vault has insufficient funds to be drained';
+        })
+      );
+    }
+    // In any other case
+    return '';
+  });
+
   const methodOptions = useMemo(
     () =>
       [
@@ -1287,6 +1331,7 @@ function SendTxModal({ isOpen, onClose }: SendTxModalProps): JSX.Element {
           isMultiSig={txData.isMultiSig ?? false}
           required={txData.required}
           isLedgerRejected={isLedgerRejected}
+          insufficientFunds={insufficientFunds}
         />
       )}
     </>