Skip to content
This repository has been archived by the owner on Apr 25, 2024. It is now read-only.

Safe Mode #160

Closed
wants to merge 8 commits into from
Closed
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 54 additions & 6 deletions src/entities/protocols/uniswap.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand All @@ -27,11 +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
hensha256 marked this conversation as resolved.
Show resolved Hide resolved
export type SwapOptions = Omit<RouterSwapOptions, 'inputTokenPermit'> & {
inputTokenPermit?: Permit2Permit
flatFee?: FlatFeeOptions
safetyMode?: SafetyModes
}

const REFUND_ETH_PRICE_IMPACT_THRESHOLD = new Percent(50, 100)
Expand All @@ -53,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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe we rename to shouldSweepETH and shouldSweepIntermediateTokens for readability, know it's a bool flag

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
Expand Down Expand Up @@ -147,14 +167,42 @@ export class UniswapTrade implements Command {
}
}

if (inputIsNative && (this.trade.tradeType === TradeType.EXACT_OUTPUT || riskOfPartialFill(this.trade))) {
// for exactOutput swaps that take native currency as input
// we need to send back the change to the user
planner.addCommand(CommandType.UNWRAP_WETH, [this.options.recipient, 0])
if (inputIsNative) {
if (this.trade.tradeType === TradeType.EXACT_OUTPUT || riskOfPartialFill(this.trade)) {
// for exactOutput swaps that take native currency as input
// we need to send back the change to the user
planner.addCommand(CommandType.UNWRAP_WETH, [this.options.recipient, 0])
}
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
function addV2Swap<TInput extends Currency, TOutput extends Currency>(
planner: RoutePlanner,
Expand Down Expand Up @@ -295,7 +343,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
Expand Down
Loading