Adding multiple modes

src/entities/protocols/uniswap.ts

@@ -27,13 +27,26 @@ export type FlatFeeOptions = {
   recipient: string
 }
 
+// SafetyModes can be used to protect against any tokens being left in the router. With a correct integration to the SDK, no
+// safety modes should be needed.
+// Normal: The normal router SDK encodings are used. This is the default.
+// InputOutputSweeps: This adds an extra ETH sweep for ETH-input trades incase extra ETH was accidentally sent to the transaction
+// IntermediateSweeps: This sweeps all intermediate tokens on a route, and protects against trading on very low liquidity pools
+// AllSweeps: Both InputOutputSweeps and IntermediateSweeps are applied.
+export enum SafetyModes {
+  Normal = 'Normal',
+  InputOutputSweeps = 'InputOutputSweeps',
+  IntermediateSweeps = 'IntermediateSweeps',
+  AllSweeps = 'AllSweeps',
+}
+
 // 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
-  safeMode?: boolean
+  safetyMode?: SafetyModes
 }
 
 const REFUND_ETH_PRICE_IMPACT_THRESHOLD = new Percent(50, 100)
@@ -55,6 +68,11 @@ export class UniswapTrade implements Command {
   encode(planner: RoutePlanner, _config: TradeConfig): void {
     let payerIsUser = true
 
+    let safetyMode = this.options.safetyMode ?? SafetyModes.Normal
+    let inputETHSweep = safetyMode == SafetyModes.InputOutputSweeps || safetyMode == SafetyModes.AllSweeps
+    let intermediateSweeps = safetyMode == SafetyModes.IntermediateSweeps || safetyMode == SafetyModes.AllSweeps
+    let tokensToSweep = new Set<string>()
+
     // If the input currency is the native currency, we need to wrap it with the router as the recipient
     if (this.trade.inputAmount.currency.isNative) {
       // TODO: optimize if only one v2 pool we can directly send this to the pool
@@ -155,11 +173,34 @@ 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) {
+      if (inputETHSweep) {
         planner.addCommand(CommandType.SWEEP, [ETH_ADDRESS, this.options.recipient, 0])
       }
+      if (intermediateSweeps) {
+        addIntermediateSweeps(planner, this.trade, this.options.recipient, tokensToSweep)
+      }
+    }
+  }
+}
+
+function addIntermediateSweeps(
+  planner: RoutePlanner,
+  trade: RouterTrade<Currency, Currency, TradeType>,
+  recipient: string,
+  tokensToSweep: Set<string>
+) {
+  for (const swap of trade.swaps) {
+    for (const token of swap.route.path) {
+      tokensToSweep.add(token.address)
     }
   }
+
+  tokensToSweep.delete(trade.inputAmount.currency.wrapped.address)
+  tokensToSweep.delete(trade.outputAmount.currency.wrapped.address)
+
+  for (const token of tokensToSweep) {
+    planner.addCommand(CommandType.SWEEP, [token, recipient, 0])
+  }
 }
 
 // encode a uniswap v2 swap