transfer out tokens before transactions

thirdweb-dev · Aug 8, 2024 · 49618aa · 49618aa
1 parent 213bbc2
commit 49618aa
Showing 1 changed file with 13 additions and 9 deletions.
diff --git a/src/PayGatewayExtension.sol b/src/PayGatewayExtension.sol
@@ -206,6 +206,16 @@ contract PayGatewayExtension is Ownable, EIP712, ModularExtension, ReentrancyGua
 
     function initiateTokenPurchase(PayRequest calldata req, bytes calldata signature) external payable nonReentrant {
         bool isTokenNative = _isTokenNative(req.tokenAddress);
+
+        // transfer out any leftover balance
+        if (address(this).balance != msg.value) {
+            _withdrawTo(NATIVE_TOKEN_ADDRESS, address(this).balance - msg.value, payable(owner()));
+        }
+        if (!isTokenNative) {
+            uint256 balance = ERC20(req.tokenAddress).balanceOf(address(this));
+            _withdrawTo(req.tokenAddress, balance, payable(owner()));
+        }
+
         // verify amount
         if (req.tokenAmount == 0) {
             revert PayGatewayInvalidAmount(req.tokenAmount);
@@ -259,15 +269,9 @@ contract PayGatewayExtension is Ownable, EIP712, ModularExtension, ReentrancyGua
                     }
                     require(target.tokenIn != address(0), "Invalid token addresses");
 
-                    uint256 amount;
-                    if (sendValue != 0) {
-                        amount = sendValue;
-                        sendValue = 0;
-                    } else {
-                        amount = _isTokenNative(target.tokenIn)
-                            ? address(this).balance
-                            : ERC20(target.tokenIn).balanceOf(address(this));
-                    }
+                    uint256 amount = _isTokenNative(target.tokenIn)
+                        ? address(this).balance
+                        : ERC20(target.tokenIn).balanceOf(address(this));
 
                     if (target.amountInOffset != type(uint256).max) {
                         // Modify calldata to update amount to send