From 43b229a954bdb520ab8d9102b9c716aa3766d7ac Mon Sep 17 00:00:00 2001
From: hensha256 <henshawalice@gmail.com>
Date: Tue, 13 Feb 2024 16:25:56 -0300
Subject: [PATCH] Safe mode

---
 src/entities/protocols/uniswap.ts        |  6 +++++-
 test/forge/SwapERC20CallParameters.t.sol | 16 ++++++++++++++++
 test/forge/interop.json                  |  4 ++++
 test/uniswapTrades.test.ts               | 16 ++++++++++++++++
 4 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/src/entities/protocols/uniswap.ts b/src/entities/protocols/uniswap.ts
index f7262e4d..9c058420 100644
--- a/src/entities/protocols/uniswap.ts
+++ b/src/entities/protocols/uniswap.ts
@@ -18,7 +18,7 @@ import {
 import { Permit2Permit } from '../../utils/inputTokens'
 import { Currency, TradeType, CurrencyAmount, Percent } from '@uniswap/sdk-core'
 import { Command, RouterTradeType, TradeConfig } from '../Command'
-import { SENDER_AS_RECIPIENT, ROUTER_AS_RECIPIENT, CONTRACT_BALANCE } from '../../utils/constants'
+import { SENDER_AS_RECIPIENT, ROUTER_AS_RECIPIENT, CONTRACT_BALANCE, ETH_ADDRESS } from '../../utils/constants'
 import { encodeFeeBips } from '../../utils/numbers'
 import { BigNumber, BigNumberish } from 'ethers'
 
@@ -29,9 +29,11 @@ export type FlatFeeOptions = {
 
 // the existing router permit object doesn't include enough data for permit2
 // so we extend swap options with the permit2 permit
+// when safe mode is enabled, the SDK will add an extra ETH sweep for security
 export type SwapOptions = Omit<RouterSwapOptions, 'inputTokenPermit'> & {
   inputTokenPermit?: Permit2Permit
   flatFee?: FlatFeeOptions
+  safeMode?: boolean
 }
 
 const REFUND_ETH_PRICE_IMPACT_THRESHOLD = new Percent(50, 100)
@@ -152,6 +154,8 @@ export class UniswapTrade implements Command {
       // we need to send back the change to the user
       planner.addCommand(CommandType.UNWRAP_WETH, [this.options.recipient, 0])
     }
+
+    if (this.options.safeMode) planner.addCommand(CommandType.SWEEP, [ETH_ADDRESS, this.options.recipient, 0])
   }
 }
 
diff --git a/test/forge/SwapERC20CallParameters.t.sol b/test/forge/SwapERC20CallParameters.t.sol
index bf9ca41a..fbfa44eb 100644
--- a/test/forge/SwapERC20CallParameters.t.sol
+++ b/test/forge/SwapERC20CallParameters.t.sol
@@ -427,6 +427,22 @@ contract SwapERC20CallParametersTest is Test, Interop, DeployRouter {
         assertGt(DAI.balanceOf(RECIPIENT), 1000 * ONE_DAI);
     }
 
+    function testV3ExactInputNativeWithSafeMode() public {
+        MethodParameters memory params = readFixture(json, "._UNISWAP_V3_ETH_FOR_DAI_SAFE_MODE");
+
+        assertEq(from.balance, BALANCE);
+        assertEq(DAI.balanceOf(RECIPIENT), 0);
+
+        // intended call value is 1e18 but 5e18 are sent in
+        assertEq(params.value, 1e18);
+        (bool success,) = address(router).call{value: 5e18}(params.data);
+        require(success, "call failed");
+
+        // the final balance only decreased by 1e18 because safemode swept back the excess
+        assertLe(from.balance, BALANCE - params.value);
+        assertGt(DAI.balanceOf(RECIPIENT), 1000 * ONE_DAI);
+    }
+
     function testV3ExactOutputSingleNative() public {
         MethodParameters memory params = readFixture(json, "._UNISWAP_V3_ETH_FOR_1000_USDC");
 
diff --git a/test/forge/interop.json b/test/forge/interop.json
index 592fc43e..d4c1e555 100644
--- a/test/forge/interop.json
+++ b/test/forge/interop.json
@@ -242,5 +242,9 @@
   "_UNISWAP_V2_USCD_FOR_10_ETH_WITH_FLAT_FEE": {
     "calldata": "0x24856bc300000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000080000000000000000000000000000000000000000000000000000000000000000309050c0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000300000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000180000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000000002000000000000000000000000000000000000000000000000d02ab486cedc0000000000000000000000000000000000000000000000000000000000047ab1042900000000000000000000000000000000000000000000000000000000000000a000000000000000000000000000000000000000000000000000000000000000010000000000000000000000000000000000000000000000000000000000000002000000000000000000000000a0b86991c6218b36c1d19d4a2e9eb0ce3606eb48000000000000000000000000c02aaa39b223fe8d0a0e5c4f27ead9083c756cc20000000000000000000000000000000000000000000000000000000000000060000000000000000000000000c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2000000000000000000000000bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb0000000000000000000000000000000000000000000000004563918244f400000000000000000000000000000000000000000000000000000000000000000040000000000000000000000000aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa0000000000000000000000000000000000000000000000008ac7230489e80000",
     "value": "0"
+  },
+  "_UNISWAP_V3_ETH_FOR_DAI_SAFE_MODE": {
+    "calldata": "0x24856bc30000000000000000000000000000000000000000000000000000000000000040000000000000000000000000000000000000000000000000000000000000008000000000000000000000000000000000000000000000000000000000000000030b000400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000003000000000000000000000000000000000000000000000000000000000000006000000000000000000000000000000000000000000000000000000000000000c00000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000de0b6b3a76400000000000000000000000000000000000000000000000000000000000000000120000000000000000000000000aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa0000000000000000000000000000000000000000000000000de0b6b3a764000000000000000000000000000000000000000000000000003eb3459f0ce6ae000b00000000000000000000000000000000000000000000000000000000000000a000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000042c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2000bb8a0b86991c6218b36c1d19d4a2e9eb0ce3606eb480001f46b175474e89094c44da98b954eedeac495271d0f00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000000000000000000000000000000aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa0000000000000000000000000000000000000000000000000000000000000000",
+    "value": "1000000000000000000"
   }
 }
diff --git a/test/uniswapTrades.test.ts b/test/uniswapTrades.test.ts
index d4037234..2ac15bd6 100644
--- a/test/uniswapTrades.test.ts
+++ b/test/uniswapTrades.test.ts
@@ -425,6 +425,22 @@ describe('Uniswap', () => {
       expect(methodParameters.value).to.eq(methodParametersV2.value)
     })
 
+    it('encodes a single exactInput ETH->USDC->DAI swap in safemode, sends too much ETH', async () => {
+      const inputEther = utils.parseEther('1').toString()
+      const trade = await V3Trade.fromRoute(
+        new RouteV3([WETH_USDC_V3, USDC_DAI_V3], ETHER, DAI),
+        CurrencyAmount.fromRawAmount(ETHER, inputEther),
+        TradeType.EXACT_INPUT
+      )
+      const opts = swapOptions({ safeMode: true })
+      const methodParameters = SwapRouter.swapERC20CallParameters(buildTrade([trade]), opts)
+      let methodParametersV2 = SwapRouter.swapCallParameters(new UniswapTrade(buildTrade([trade]), opts))
+      registerFixture('_UNISWAP_V3_ETH_FOR_DAI_SAFE_MODE', methodParametersV2)
+      expect(hexToDecimalString(methodParameters.value)).to.eq(inputEther)
+      expect(methodParameters.calldata).to.eq(methodParametersV2.calldata)
+      expect(methodParameters.value).to.eq(methodParametersV2.value)
+    })
+
     it('encodes a single exactOutput ETH->USDC swap', async () => {
       const outputUSDC = utils.parseUnits('1000', 6).toString()
       const trade = await V3Trade.fromRoute(