From 592d3552cb43649f732b37c983ec0edfcb36fdbf Mon Sep 17 00:00:00 2001
From: tommyrharper <thomasrobertharper@gmail.com>
Date: Thu, 11 Jul 2024 18:49:01 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=B7=F0=9F=8F=BB=20Handle=20insufficien?=
 =?UTF-8?q?t=20margin?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/MarginPaymaster.sol                       | 31 +++++++++++++++----
 src/interfaces/external/IPerpsMarketProxy.sol |  4 +++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/MarginPaymaster.sol b/src/MarginPaymaster.sol
index 1654af7..d493df8 100644
--- a/src/MarginPaymaster.sol
+++ b/src/MarginPaymaster.sol
@@ -117,7 +117,6 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
                                 POST OP
     //////////////////////////////////////////////////////////////*/
 
-    // TODO: handle user not having enough funds
     function postOp(
         PostOpMode,
         bytes calldata context,
@@ -144,9 +143,12 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
             uint256 sUSDToWithdrawFromMargin = (costOfGasInUSDC -
                 availableUSDCInWallet) * 1e12;
             // TODO: handle users who don't have an snx account or margin or enough margin
-            withdrawFromMargin(sender, sUSDToWithdrawFromMargin);
+            uint256 withdrawn = withdrawFromMargin(
+                sender,
+                sUSDToWithdrawFromMargin
+            );
             // zap sUSD into USDC
-            _zapOut(sUSDToWithdrawFromMargin);
+            _zapOut(withdrawn);
         }
     }
 
@@ -212,16 +214,33 @@ contract MarginPaymaster is IPaymaster, Zap, Ownable {
         return uint128(snxV3AccountsModule.tokenOfOwnerByIndex(wallet, 0));
     }
 
+    /// @notice withdraws sUSD from margin account
+    /// @notice if insufficent margin, pulls out whatever is available
     function withdrawFromMargin(
         address sender,
         uint256 sUSDToWithdrawFromMargin
-    ) internal {
+    ) internal returns (uint256) {
+        uint128 accountId = getWalletAccountId(sender);
+        int256 withdrawableMargin = perpsMarketSNXV3.getWithdrawableMargin(
+            accountId
+        );
+
+        if (withdrawableMargin < 0) return 0;
+        uint256 withdrawableMarginUint = uint256(withdrawableMargin);
+
+        uint256 amountToPullFromMargin = sUSDToWithdrawFromMargin <
+            withdrawableMarginUint
+            ? sUSDToWithdrawFromMargin
+            : withdrawableMarginUint;
+
         // pull sUSD from margin
         perpsMarketSNXV3.modifyCollateral(
-            getWalletAccountId(sender),
+            accountId,
             sUSDId,
-            -int256(sUSDToWithdrawFromMargin)
+            -int256(amountToPullFromMargin)
         );
+
+        return amountToPullFromMargin;
     }
 
     function swapUSDCForWETH(
diff --git a/src/interfaces/external/IPerpsMarketProxy.sol b/src/interfaces/external/IPerpsMarketProxy.sol
index 2b1f4b4..916d872 100644
--- a/src/interfaces/external/IPerpsMarketProxy.sol
+++ b/src/interfaces/external/IPerpsMarketProxy.sol
@@ -23,6 +23,10 @@ interface IPerpsMarketProxy {
 
     function getAccountTokenAddress() external view returns (address accountNftToken);
 
+    function getWithdrawableMargin(
+        uint128 accountId
+    ) external view returns (int256 withdrawableMargin);
+
     /// @notice Returns the address that owns a given account, as recorded by the system.
     /// @param accountId The account id whose owner is being retrieved.
     /// @return owner The owner of the given account id.