feat: pass selector in transfers

morpho-org · Nov 22, 2023 · 464479d · 464479d
1 parent 560c176
commit 464479d
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 16 deletions.
diff --git a/src/Morpho.sol b/src/Morpho.sol
@@ -188,7 +188,7 @@ contract Morpho is IMorphoStaticTyping {
 
         if (data.length > 0) IMorphoSupplyCallback(msg.sender).onMorphoSupply(assets, data);
 
-        IERC20(marketParams.loanToken).safeTransferFrom(msg.sender, address(this), assets);
+        IERC20(marketParams.loanToken).safeTransferFrom(msg.sender, address(this), assets, Morpho.supply.selector);
 
         return (assets, shares);
     }
@@ -221,7 +221,7 @@ contract Morpho is IMorphoStaticTyping {
 
         emit EventsLib.Withdraw(id, msg.sender, onBehalf, receiver, assets, shares);
 
-        IERC20(marketParams.loanToken).safeTransfer(receiver, assets);
+        IERC20(marketParams.loanToken).safeTransfer(receiver, assets, Morpho.withdraw.selector);
 
         return (assets, shares);
     }
@@ -257,7 +257,7 @@ contract Morpho is IMorphoStaticTyping {
 
         emit EventsLib.Borrow(id, msg.sender, onBehalf, receiver, assets, shares);
 
-        IERC20(marketParams.loanToken).safeTransfer(receiver, assets);
+        IERC20(marketParams.loanToken).safeTransfer(receiver, assets, Morpho.borrow.selector);
 
         return (assets, shares);
     }
@@ -289,7 +289,7 @@ contract Morpho is IMorphoStaticTyping {
 
         if (data.length > 0) IMorphoRepayCallback(msg.sender).onMorphoRepay(assets, data);
 
-        IERC20(marketParams.loanToken).safeTransferFrom(msg.sender, address(this), assets);
+        IERC20(marketParams.loanToken).safeTransferFrom(msg.sender, address(this), assets, Morpho.repay.selector);
 
         return (assets, shares);
     }
@@ -313,7 +313,9 @@ contract Morpho is IMorphoStaticTyping {
 
         if (data.length > 0) IMorphoSupplyCollateralCallback(msg.sender).onMorphoSupplyCollateral(assets, data);
 
-        IERC20(marketParams.collateralToken).safeTransferFrom(msg.sender, address(this), assets);
+        IERC20(marketParams.collateralToken).safeTransferFrom(
+            msg.sender, address(this), assets, Morpho.supplyCollateral.selector
+        );
     }
 
     /// @inheritdoc IMorphoBase
@@ -335,7 +337,7 @@ contract Morpho is IMorphoStaticTyping {
 
         emit EventsLib.WithdrawCollateral(id, msg.sender, onBehalf, receiver, assets);
 
-        IERC20(marketParams.collateralToken).safeTransfer(receiver, assets);
+        IERC20(marketParams.collateralToken).safeTransfer(receiver, assets, Morpho.withdrawCollateral.selector);
     }
 
     /* LIQUIDATION */
@@ -398,14 +400,16 @@ contract Morpho is IMorphoStaticTyping {
             position[id][borrower].borrowShares = 0;
         }
 
-        IERC20(marketParams.collateralToken).safeTransfer(msg.sender, seizedAssets);
+        IERC20(marketParams.collateralToken).safeTransfer(msg.sender, seizedAssets, Morpho.liquidate.selector);
 
         // `repaidAssets` may be greater than `totalBorrowAssets` by 1.
         emit EventsLib.Liquidate(id, msg.sender, borrower, repaidAssets, repaidShares, seizedAssets, badDebtShares);
 
         if (data.length > 0) IMorphoLiquidateCallback(msg.sender).onMorphoLiquidate(repaidAssets, data);
 
-        IERC20(marketParams.loanToken).safeTransferFrom(msg.sender, address(this), repaidAssets);
+        IERC20(marketParams.loanToken).safeTransferFrom(
+            msg.sender, address(this), repaidAssets, Morpho.liquidate.selector
+        );
 
         return (seizedAssets, repaidAssets);
     }
@@ -414,13 +418,13 @@ contract Morpho is IMorphoStaticTyping {
 
     /// @inheritdoc IMorphoBase
     function flashLoan(address token, uint256 assets, bytes calldata data) external {
-        IERC20(token).safeTransfer(msg.sender, assets);
+        IERC20(token).safeTransfer(msg.sender, assets, Morpho.flashLoan.selector);
 
         emit EventsLib.FlashLoan(msg.sender, token, assets);
 
         IMorphoFlashLoanCallback(msg.sender).onMorphoFlashLoan(assets, data);
 
-        IERC20(token).safeTransferFrom(msg.sender, address(this), assets);
+        IERC20(token).safeTransferFrom(msg.sender, address(this), assets, Morpho.flashLoan.selector);
     }
 
     /* AUTHORIZATION */

diff --git a/src/libraries/SafeTransferLib.sol b/src/libraries/SafeTransferLib.sol
@@ -18,17 +18,17 @@ interface IERC20Internal {
 /// @dev It is the responsibility of the market creator to make sure that the address of the token has non-zero code.
 library SafeTransferLib {
     /// @dev Warning: It does not revert on `token` with no code.
-    function safeTransfer(IERC20 token, address to, uint256 value) internal {
+    function safeTransfer(IERC20 token, address to, uint256 value, bytes4 selector) internal {
         (bool success, bytes memory returndata) =
-            address(token).call(abi.encodeCall(IERC20Internal.transfer, (to, value)));
+            address(token).call(abi.encodeWithSelector(IERC20Internal.transfer.selector, to, value, selector));
         require(success, ErrorsLib.TRANSFER_REVERTED);
         require(returndata.length == 0 || abi.decode(returndata, (bool)), ErrorsLib.TRANSFER_RETURNED_FALSE);
     }
 
     /// @dev Warning: It does not revert on `token` with no code.
-    function safeTransferFrom(IERC20 token, address from, address to, uint256 value) internal {
+    function safeTransferFrom(IERC20 token, address from, address to, uint256 value, bytes4 selector) internal {
         (bool success, bytes memory returndata) =
-            address(token).call(abi.encodeCall(IERC20Internal.transferFrom, (from, to, value)));
+            address(token).call(abi.encodeWithSelector(IERC20Internal.transferFrom.selector, from, to, value, selector));
         require(success, ErrorsLib.TRANSFER_FROM_REVERTED);
         require(returndata.length == 0 || abi.decode(returndata, (bool)), ErrorsLib.TRANSFER_FROM_RETURNED_FALSE);
     }

diff --git a/test/forge/libraries/SafeTransferLibTest.sol b/test/forge/libraries/SafeTransferLibTest.sol
@@ -86,10 +86,10 @@ contract SafeTransferLibTest is Test {
     }
 
     function safeTransfer(address token, address to, uint256 amount) external {
-        IERC20(token).safeTransfer(to, amount);
+        IERC20(token).safeTransfer(to, amount, 0x00000000);
     }
 
     function safeTransferFrom(address token, address from, address to, uint256 amount) external {
-        IERC20(token).safeTransferFrom(from, to, amount);
+        IERC20(token).safeTransferFrom(from, to, amount, 0x00000000);
     }
 }