From 0c40e9e97923d5bad61aa812ba739c2fe4260cf8 Mon Sep 17 00:00:00 2001
From: royvardhan <royvardhan.official@gmail.com>
Date: Wed, 12 Feb 2025 22:30:13 +0530
Subject: [PATCH] fix: handle native token balance changes

---
 foundry/src/executors/UniswapV4Executor.sol | 24 +++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/foundry/src/executors/UniswapV4Executor.sol b/foundry/src/executors/UniswapV4Executor.sol
index 9cafb93..0b31939 100644
--- a/foundry/src/executors/UniswapV4Executor.sol
+++ b/foundry/src/executors/UniswapV4Executor.sol
@@ -39,13 +39,29 @@ contract UniswapV4Executor is IExecutor, V4Router {
             uint256 amount
         ) = _decodeData(data);
 
-        uint256 tokenOutBalanceBefore = IERC20(tokenOut).balanceOf(receiver);
-        uint256 tokenInBalanceBefore = IERC20(tokenIn).balanceOf(address(this));
+        uint256 tokenOutBalanceBefore;
+        uint256 tokenInBalanceBefore;
+
+        tokenOutBalanceBefore = tokenOut == address(0)
+            ? receiver.balance
+            : IERC20(tokenOut).balanceOf(receiver);
+
+        tokenInBalanceBefore = tokenIn == address(0)
+            ? address(this).balance
+            : IERC20(tokenIn).balanceOf(address(this));
 
         _executeActions(data);
 
-        uint256 tokenOutBalanceAfter = IERC20(tokenOut).balanceOf(receiver);
-        uint256 tokenInBalanceAfter = IERC20(tokenIn).balanceOf(address(this));
+        uint256 tokenOutBalanceAfter;
+        uint256 tokenInBalanceAfter;
+
+        tokenOutBalanceAfter = tokenOut == address(0)
+            ? receiver.balance
+            : IERC20(tokenOut).balanceOf(receiver);
+
+        tokenInBalanceAfter = tokenIn == address(0)
+            ? address(this).balance
+            : IERC20(tokenIn).balanceOf(address(this));
 
         if (isExactInput) {
             calculatedAmount = tokenOutBalanceAfter - tokenOutBalanceBefore;