Safe Mode

src/entities/protocols/uniswap.ts

@@ -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 extra token sweeps for security
 export type SwapOptions = Omit<RouterSwapOptions, 'inputTokenPermit'> & {
   inputTokenPermit?: Permit2Permit
   flatFee?: FlatFeeOptions
+  tokensToSweep?: string[]
 }
 
 const REFUND_ETH_PRICE_IMPACT_THRESHOLD = new Percent(50, 100)
@@ -152,6 +154,14 @@ 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.tokensToSweep) addExtraTokenSweeps(planner, this.options.recipient, this.options.tokensToSweep)
+  }
+}
+
+function addExtraTokenSweeps(planner: RoutePlanner, recipient: string, tokensToSweep: string[]) {
+  for (const token of tokensToSweep) {
+    planner.addCommand(CommandType.SWEEP, [token, recipient, 0])
   }
 }
 
@@ -295,7 +305,7 @@ function addMixedSwap<TInput extends Currency, TOutput extends Currency>(
         // if not last section: send tokens directly to the first v2 pair of the next section
         // note: because of the partitioning function we can be sure that the next section is v2
         isLastSectionInRoute(i) ? tradeRecipient : (sections[i + 1][0] as Pair).liquidityToken.address,
-        i == 0 ? amountIn : CONTRACT_BALANCE, // amountIn
+        i === 0 ? amountIn : CONTRACT_BALANCE, // amountIn
         !isLastSectionInRoute(i) ? 0 : amountOut, // amountOut
         path, // path
         payerIsUser && i === 0, // payerIsUser