From 452500b6ee0574a6f145a41a6c37ff7358209f15 Mon Sep 17 00:00:00 2001
From: ungaro <alp.guneysel@gmail.com>
Date: Mon, 9 Dec 2024 23:46:40 -0500
Subject: [PATCH] safetransfer changes

---
 nest/src/ComponentToken.sol | 52 +++++++++++++++++++------------------
 1 file changed, 27 insertions(+), 25 deletions(-)

diff --git a/nest/src/ComponentToken.sol b/nest/src/ComponentToken.sol
index ad76f9a..675d70c 100644
--- a/nest/src/ComponentToken.sol
+++ b/nest/src/ComponentToken.sol
@@ -257,6 +257,7 @@ abstract contract ComponentToken is
 
         SafeERC20.safeTransferFrom(IERC20(asset()), owner, address(this), assets);
         $.pendingDepositRequest[controller] += assets;
+
         emit DepositRequest(controller, owner, REQUEST_ID, owner, assets);
         return REQUEST_ID;
     }
@@ -332,21 +333,22 @@ abstract contract ComponentToken is
         }
 
         ComponentTokenStorage storage $ = _getComponentTokenStorage();
-        assets = convertToAssets(shares);
-
         if ($.asyncDeposit) {
-            if ($.claimableDepositRequest[controller] < assets) {
-                revert InsufficientRequestBalance(controller, assets, 1);
+            // Check shares directly instead of converting to assets
+            if ($.sharesDepositRequest[controller] < shares) {
+                revert InsufficientRequestBalance(controller, shares, 1);
             }
-            $.claimableDepositRequest[controller] -= assets;
-            $.sharesDepositRequest[controller] -= shares;
+            // Use the pre-calculated assets amount from when deposit was notified
+            assets = $.claimableDepositRequest[controller];
+            $.claimableDepositRequest[controller] = 0;
+            $.sharesDepositRequest[controller] = 0;
         } else {
-            SafeERC20.safeTransferFrom(IERC20(asset()), controller, address(this), assets);
+            assets = previewMint(shares);
+            _deposit(msg.sender, receiver, assets, shares);
         }
-
         _mint(receiver, shares);
-
-        emit Deposit(controller, receiver, assets, shares);
+        emit Deposit(msg.sender, receiver, assets, shares);
+        return assets;
     }
 
     /// @inheritdoc IComponentToken
@@ -437,7 +439,7 @@ abstract contract ComponentToken is
         address receiver,
         address controller
     ) public virtual override(ERC4626Upgradeable, IERC7540) nonReentrant returns (uint256 shares) {
-        if (assets == 0) {
+        if (shares == 0) {
             revert ZeroAmount();
         }
         if (msg.sender != controller) {
@@ -445,21 +447,21 @@ abstract contract ComponentToken is
         }
 
         ComponentTokenStorage storage $ = _getComponentTokenStorage();
-        shares = convertToShares(assets);
-
         if ($.asyncRedeem) {
-            if ($.claimableRedeemRequest[controller] < shares) {
-                revert InsufficientRequestBalance(controller, shares, 3);
+            // Use the pre-calculated assets amount from when redeem was notified
+            if ($.assetsRedeemRequest[controller] < assets) {
+                revert InsufficientRequestBalance(controller, assets, 3);
             }
-            $.claimableRedeemRequest[controller] -= shares;
-            $.assetsRedeemRequest[controller] -= assets;
+            shares = $.claimableRedeemRequest[controller];
+            $.claimableRedeemRequest[controller] = 0;
+            $.assetsRedeemRequest[controller] = 0;
         } else {
-            _burn(controller, shares);
+            shares = previewWithdraw(assets);
+            _withdraw(msg.sender, receiver, msg.sender, assets, shares);
         }
-
-        SafeERC20.safeTransfer(IERC20(asset()), receiver, assets);
-
-        emit Withdraw(controller, receiver, controller, assets, shares);
+        _burn(msg.sender, shares);
+        emit Withdraw(msg.sender, receiver, msg.sender, assets, shares);
+        return shares;
     }
 
     // Getter View Functions
@@ -518,7 +520,7 @@ abstract contract ComponentToken is
         if (_getComponentTokenStorage().asyncDeposit) {
             revert Unimplemented();
         }
-        assets = super.previewDeposit(shares);
+        assets = convertToAssets(shares);
     }
 
     /**
@@ -545,7 +547,7 @@ abstract contract ComponentToken is
         if (_getComponentTokenStorage().asyncRedeem) {
             revert Unimplemented();
         }
-        shares = super.previewWithdraw(assets);
+        shares = convertToShares(assets);
     }
 
     /// @inheritdoc IERC7540
@@ -553,4 +555,4 @@ abstract contract ComponentToken is
         revert Unimplemented();
     }
 
-}
\ No newline at end of file
+}