fix: potential time-of-check vs time-of-use vulnerability with maxUin…

src/actions/market/marketLeverageBorrowAction/index.ts

@@ -2,8 +2,7 @@ import { DEFAULT_SLIPPAGE_TOLERANCE, MarketId, MathLib } from "@morpho-org/blue-
 import { fetchMarket } from "@morpho-org/blue-sdk-viem";
 import { BundlerAction, BundlerCall, encodeBundle } from "@morpho-org/bundler-sdk-viem";
 import { populateSubBundle } from "@morpho-org/bundler-sdk-viem";
-import { Address, Client, erc20Abi, maxUint256 } from "viem";
-import { readContract } from "viem/actions";
+import { Address, Client, maxUint256 } from "viem";
 
 import { getParaswapExactBuyTxPayload } from "@/actions/data/paraswap/getParaswapExactBuy";
 import { getIsContract } from "@/actions/data/rpc/getIsContract";
@@ -26,7 +25,7 @@ interface MarketLeveragedBorrowActionParameters {
 
   accountAddress: Address;
 
-  initialCollateralAmount: bigint; // a.k.a margin, maxUint256 for entire wallet balance
+  initialCollateralAmount: bigint; // a.k.a margin
   leverageFactor: number; // (1, (1 + S) / (1 + S - LLTV_WITH_MARGIN))]
   maxSlippageTolerance: number; // (0,1)
 }
@@ -50,25 +49,23 @@ export async function marketLeveragedBorrowAction({
   leverageFactor,
   maxSlippageTolerance,
 }: MarketLeveragedBorrowActionParameters): Promise<MarketLeveragedBorrowAction> {
-  // Input validation performed within computeLeverageValues
+  // Disallow maxUint256, we require exact collateral amount
+  if (initialCollateralAmount >= maxUint256) {
+    return {
+      status: "error",
+      message: "Initial collateral amount cannot be greater than or equal to max uint256",
+    };
+  }
+
+  // Other input validation performed within computeLeverageValues
 
   let market = await fetchMarket(marketId, publicClient);
   const { collateralToken: collateralTokenAddress, loanToken: loanTokenAddress } = market.params;
 
-  const accountCollateralBalance = await readContract(publicClient, {
-    address: collateralTokenAddress,
-    abi: erc20Abi,
-    functionName: "balanceOf",
-    args: [accountAddress],
-  });
-
-  const initialCollateralAmountInternal =
-    initialCollateralAmount == maxUint256 ? accountCollateralBalance : initialCollateralAmount;
-
   let collateralAmount: bigint;
   let loanAmount: bigint;
   try {
-    const values = computeLeverageValues(initialCollateralAmountInternal, leverageFactor, maxSlippageTolerance, market);
+    const values = computeLeverageValues(initialCollateralAmount, leverageFactor, maxSlippageTolerance, market);
     collateralAmount = values.collateralAmount;
     loanAmount = values.loanAmount;
   } catch (e) {
@@ -96,13 +93,13 @@ export async function marketLeveragedBorrowAction({
   const positionLoanBefore = market.toBorrowAssets(accountPosition.borrowShares);
   const positionLtvBefore = market.getLtv(accountPosition) ?? 0n;
 
-  const requiredSwapCollateralAmount = collateralAmount - initialCollateralAmountInternal;
+  const requiredSwapCollateralAmount = collateralAmount - initialCollateralAmount;
 
   try {
     const inputSubbundle = inputTransferSubbundle({
       accountAddress,
       tokenAddress: collateralTokenAddress,
-      amount: initialCollateralAmount, // Handles maxUint256
+      amount: initialCollateralAmount,
       recipientAddress: GENERAL_ADAPTER_1_ADDRESS,
       config: {
         accountSupportsSignatures: !isContract,

src/actions/market/marketRepayWithCollateralAction.ts

@@ -65,6 +65,8 @@ export async function marketRepayWithCollateralAction({
   const market = simulationState.getMarket(marketId);
   const accountPosition = simulationState.getPosition(accountAddress, marketId);
 
+  const accountPositionInitialBorrowShares = accountPosition.borrowShares;
+
   const positionCollateralBefore = accountPosition.collateral;
   const positionLoanBefore = market.toBorrowAssets(accountPosition.borrowShares);
   const positionLtvBefore = market.getLtv(accountPosition) ?? 0n;
@@ -133,7 +135,7 @@ export async function marketRepayWithCollateralAction({
           id: marketId,
           onBehalf: accountAddress,
           receiver: PARASWAP_ADAPTER_ADDRESS,
-          assets: closingPosition ? accountPosition.collateral : maxCollateralSwapAmount, // Full collateral withdraw if closing position - bundler SDK has issues with maxUint256, but this works the same
+          assets: closingPosition ? accountPosition.collateral : maxCollateralSwapAmount,
         },
       },
       simulationState
@@ -152,7 +154,7 @@ export async function marketRepayWithCollateralAction({
           CHAIN_ID,
           market.params,
           closingPosition ? 0n : loanRepayAmount, // Assets
-          closingPosition ? maxUint256 : 0n, // Shares - full repay using shares when closing position
+          closingPosition ? accountPositionInitialBorrowShares : 0n, // Shares - full repay using shares when closing position
           maxSharePriceE27,
           accountAddress,
           // Repay loan callbacks, all actions will happen before Morpho pulls the required loan assets from GA1 (i.e must have required loan assets in GA1 at the end of these callbacks)

src/actions/subbundles/inputTransferSubbundle.ts

@@ -34,6 +34,13 @@ export function inputTransferSubbundle({
   simulationState,
 }: InputTransferSubbundleParameters): Subbundle & { erc20Amount: bigint; nativeAmount: bigint } {
   const isMaxTransfer = amount == maxUint256;
+
+  // Sanity check for now to disable maxUint256 input transfer.
+  // No actions should use anyways. Migration does but is disabled
+  if(isMaxTransfer) {
+    throw new Error("Max uint256 transfer is not supported");
+  }
+
   const isWrappedNative = isAddressEqual(tokenAddress, WRAPPED_NATIVE_ADDRESS);
 
   const accountErc20Holding = simulationState.getHolding(accountAddress, tokenAddress);

src/actions/vault/vaultSupplyAction.ts

@@ -16,7 +16,7 @@ interface VaultSupplyActionParameters {
   publicClient: Client;
   vaultAddress: Address;
   accountAddress: Address;
-  supplyAmount: bigint; // Max uint256 for entire account balanace
+  supplyAmount: bigint;
   allowWrappingNativeAssets: boolean; // Ignored if the vault asset is not wrapped native
 }
 
@@ -38,6 +38,14 @@ export async function vaultSupplyBundle({
     };
   }
 
+  // Disallow maxUint256 - require exact supply amount
+  if (supplyAmount >= maxUint256) {
+    return {
+      status: "error",
+      message: "Supply amount cannot be greater than or equal to max uint256",
+    };
+  }
+
   const [intitialSimulationState, accountIsContract] = await Promise.all([
     getSimulationState({
       actionType: "vault",
@@ -50,16 +58,14 @@ export async function vaultSupplyBundle({
 
   const vault = intitialSimulationState.getVault(vaultAddress);
 
-  const isMaxSupply = supplyAmount === maxUint256;
-
   try {
     const intermediateSimulationState = produceImmutable(intitialSimulationState, () => {});
 
     // Prepare input transfer, and modify simulation state accordingly
     const inputSubbundle = inputTransferSubbundle({
       accountAddress,
       tokenAddress: vault.underlying,
-      amount: supplyAmount, // Handles maxUint256
+      amount: supplyAmount,
       recipientAddress: GENERAL_ADAPTER_1_ADDRESS,
       config: {
         accountSupportsSignatures: !accountIsContract,
@@ -69,16 +75,14 @@ export async function vaultSupplyBundle({
       simulationState: intermediateSimulationState,
     });
 
-    const ga1AssetBalance = intermediateSimulationState.getHolding(GENERAL_ADAPTER_1_ADDRESS, vault.asset).balance;
-
     const vaultSupplySubbundle = subbundleFromInputOps({
       inputOps: [
         {
           type: "MetaMorpho_Deposit",
           sender: accountAddress,
           address: vaultAddress,
           args: {
-            assets: isMaxSupply ? ga1AssetBalance : supplyAmount,
+            assets: supplyAmount,
             owner: accountAddress,
             slippage: DEFAULT_SLIPPAGE_TOLERANCE,
           },

src/app/migrate/page.tsx

@@ -1,12 +1,6 @@
 import { Metadata } from "next";
 import Image from "next/image";
 
-import PositionMigrator from "@/components/PositionMigrator";
-import ProtocolMigratorTableWrapper from "@/components/ProtocolMigrator/ProtocolMigratorTableWrapper";
-import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
-import { MarketSummary, getMarketSummaries } from "@/data/whisk/getMarketSummaries";
-import { getVaultSummaries } from "@/data/whisk/getVaultSummaries";
-
 export const metadata: Metadata = {
   title: "Compound Blue | Migrate",
 };
@@ -24,7 +18,9 @@ export default function MigratePage() {
         </div>
       </section>
 
-      <Tabs defaultValue="protocol" className="flex flex-col gap-8">
+      <div>Migration is currently unavailable.</div>
+
+      {/* <Tabs defaultValue="protocol" className="flex flex-col gap-8">
         <TabsList className="w-[254px]">
           <TabsTrigger value="protocol">Protocol</TabsTrigger>
           <TabsTrigger value="position">Position</TabsTrigger>
@@ -39,19 +35,19 @@ export default function MigratePage() {
             <PositionMigratorWrapper />
           </div>
         </TabsContent>
-      </Tabs>
+      </Tabs> */}
     </>
   );
 }
 
-async function PositionMigratorWrapper() {
-  const [vaultSummaries, marketSummaries] = await Promise.all([getVaultSummaries(), getMarketSummaries()]);
+// async function PositionMigratorWrapper() {
+//   const [vaultSummaries, marketSummaries] = await Promise.all([getVaultSummaries(), getMarketSummaries()]);
 
-  if (!vaultSummaries || !marketSummaries) {
-    return null;
-  }
-  return <PositionMigrator vaultSummaries={vaultSummaries} marketSummaries={marketSummaries as MarketSummary[]} />;
-}
+//   if (!vaultSummaries || !marketSummaries) {
+//     return null;
+//   }
+//   return <PositionMigrator vaultSummaries={vaultSummaries} marketSummaries={marketSummaries as MarketSummary[]} />;
+// }
 
 export const dynamic = "force-static";
 export const revalidate = 300; // 5 min

src/components/Header/Nav.tsx

@@ -4,14 +4,13 @@ import { usePathname } from "next/navigation";
 import { ReactNode } from "react";
 
 import { Button } from "../ui/button";
-import ArrowFatLineDown from "../ui/icons/ArrowFatLineDown";
 import BarChart from "../ui/icons/BarChart";
 import CirclePlus from "../ui/icons/CirclePlus";
 
 const NAV_ITEMS: { href: string; name: string; icon?: ReactNode; isNew?: boolean }[] = [
   { href: "/", name: "Earn", icon: <BarChart className="fill-content-secondary" /> },
   { href: "/borrow", name: "Borrow", icon: <CirclePlus className="fill-content-secondary" /> },
-  { href: "/migrate", name: "Migrate", icon: <ArrowFatLineDown className="fill-content-secondary" /> },
+  // { href: "/migrate", name: "Migrate", icon: <ArrowFatLineDown className="fill-content-secondary" /> },
 ];
 
 export default function Nav() {

src/components/MarketActions/Borrow/MarketLeverageBorrow.tsx

@@ -5,7 +5,7 @@ import { useConnectModal } from "@rainbow-me/rainbowkit";
 import { ArrowRight, Info } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { useForm } from "react-hook-form";
-import { getAddress, maxUint256, parseUnits } from "viem";
+import { getAddress, parseUnits } from "viem";
 import { useAccount } from "wagmi";
 import { usePublicClient } from "wagmi";
 import { z } from "zod";
@@ -79,7 +79,6 @@ export default function MarketLeverageBorrow({
             const rawVal = parseUnits(val, market.collateralAsset.decimals);
             return rawVal <= BigInt(rawCollateralTokenBalance.balance);
           }, "Amount exceeds wallet balance."),
-        isMaxCollateral: z.boolean(),
         multiplier: z.coerce.number().min(MIN_MULTIPLIER),
         maxSlippageTolerance: z.coerce.number().min(0.2).max(MAX_SLIPPAGE_TOLERANCE_PCT),
       })
@@ -103,19 +102,13 @@ export default function MarketLeverageBorrow({
     resolver: zodResolver(formSchema),
     defaultValues: {
       initialCollateralAmount: "",
-      isMaxCollateral: false,
       multiplier: 2,
       maxSlippageTolerance: 0.5,
     },
   });
 
   const onSubmit = useCallback(
-    async ({
-      initialCollateralAmount,
-      isMaxCollateral,
-      multiplier,
-      maxSlippageTolerance,
-    }: z.infer<typeof formSchema>) => {
+    async ({ initialCollateralAmount, multiplier, maxSlippageTolerance }: z.infer<typeof formSchema>) => {
       if (!address) {
         openConnectModal?.();
         return;
@@ -128,10 +121,7 @@ export default function MarketLeverageBorrow({
 
       setSimulatingBundle(true);
 
-      const rawInitialCollateralAmount = isMaxCollateral
-        ? maxUint256
-        : parseUnits(initialCollateralAmount, market.collateralAsset.decimals);
-
+      const rawInitialCollateralAmount = parseUnits(initialCollateralAmount, market.collateralAsset.decimals);
       const leverageFactor = multiplier / (1 + maxSlippageTolerance / 100) + 1;
 
       const action = await marketLeveragedBorrowAction({
@@ -194,9 +184,6 @@ export default function MarketLeverageBorrow({
                 actionName="Add"
                 asset={market.collateralAsset}
                 rawAvailableBalance={rawCollateralTokenBalance ? BigInt(rawCollateralTokenBalance.balance) : undefined}
-                setIsMax={(isMax) => {
-                  form.setValue("isMaxCollateral", isMax);
-                }}
               />
               <div className="h-[1px] w-full bg-border-primary" />
 

src/components/VaultActions/VaultSupply.tsx

@@ -4,7 +4,7 @@ import { useConnectModal } from "@rainbow-me/rainbowkit";
 import { Info } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { useForm } from "react-hook-form";
-import { getAddress, isAddressEqual, maxUint256, parseUnits } from "viem";
+import { getAddress, isAddressEqual, parseUnits } from "viem";
 import { useAccount, useBalance, usePublicClient } from "wagmi";
 import { z } from "zod";
 
@@ -74,7 +74,6 @@ export default function VaultSupply({
           .nonempty("Amount is required.")
           .refine((val) => !isNaN(parseFloat(val)), "Amount must be a valid number.")
           .refine((val) => parseUnits(val, vault.asset.decimals) > 0n, "Amount must be greater than zero."), // This also catches the case where val is lower than token precision, but we prevent this in ActionFlowSummaryAssetItem
-        isMaxSupply: z.boolean(),
         allowWrappingNativeAssets: z.boolean(),
       })
       .superRefine((data, ctx) => {
@@ -96,7 +95,6 @@ export default function VaultSupply({
     resolver: zodResolver(formSchema),
     defaultValues: {
       supplyAmount: "",
-      isMaxSupply: false,
       allowWrappingNativeAssets: isWrappedNative,
     },
   });
@@ -115,9 +113,8 @@ export default function VaultSupply({
 
       setSimulatingBundle(true);
 
-      // Uint256 max if the user wants to supply their entire balance
-      const { supplyAmount, isMaxSupply, allowWrappingNativeAssets } = values;
-      const rawSupplyAmount = isMaxSupply ? maxUint256 : parseUnits(supplyAmount, vault.asset.decimals);
+      const { supplyAmount, allowWrappingNativeAssets } = values;
+      const rawSupplyAmount = parseUnits(supplyAmount, vault.asset.decimals);
 
       const preparedAction = await vaultSupplyBundle({
         publicClient,
@@ -192,9 +189,6 @@ export default function VaultSupply({
                 actionName="Supply"
                 asset={vault.asset}
                 rawAvailableBalance={computeRawAvailableBalance(allowWrappingNativeAssets)}
-                setIsMax={(isMax) => {
-                  form.setValue("isMaxSupply", isMax);
-                }}
               />
               {isWrappedNative && (
                 <div className="flex rounded-[8px] bg-background-inverse p-3">

test/config.ts

@@ -7,7 +7,6 @@ export const test = createViemTest(polygon, {
   forkUrl: process.env.NEXT_PUBLIC_RPC_URL_1!,
   forkBlockNumber: 71095170,
   // forkBlockNumber: 71054290,
-  // hardfork: "Latest",
 });
 
 export const polygonClient = createPublicClient({
@@ -19,7 +18,6 @@ export const currentBlock = await polygonClient.getBlock();
 export const currentBlockTest = createViemTest(polygon, {
   forkUrl: process.env.NEXT_PUBLIC_RPC_URL_1!,
   forkBlockNumber: currentBlock.number,
-  hardfork: "Latest",
 });
 
 // Used for tests that need a specific block number besides default from test (ex when have quotes)

test/helpers/morpho.ts

@@ -43,7 +43,24 @@ export async function getMorphoVaultPosition(
 
   const userShareBalance = await getErc20BalanceOf(client, vaultAddress, accountAddress);
   const userAssetBalance = vault.toAssets(userShareBalance);
-  return userAssetBalance;
+  return { userAssetBalance, userShareBalance };
+}
+
+export async function getMorphoVaultSharesToAssets(
+  client: AnvilTestClient,
+  vaultAddress: Address,
+  accountAddress: Address,
+  shares: bigint
+) {
+  const simulationState = await getSimulationState({
+    publicClient: client,
+    actionType: "vault",
+    accountAddress,
+    vaultAddress,
+  });
+
+  const vault = simulationState.getVault(vaultAddress);
+  return vault.toAssets(shares);
 }
 
 // Supply loan assets directly to a market from account for onBehalf