diff --git a/lib/handlers/constants.ts b/lib/handlers/constants.ts index 50e21cdc..e748ee65 100644 --- a/lib/handlers/constants.ts +++ b/lib/handlers/constants.ts @@ -22,4 +22,6 @@ export const UR_FUNCTION_SIGNATURES: Record = { [UR_EXECUTE_SELECTOR]: "function execute(bytes commands, bytes[] inputs)", [UR_EXECUTE_WITH_DEADLINE_SELECTOR]: "function execute(bytes commands, bytes[] inputs, uint256 deadline)" }; -export const UR_EXECUTE_DEADLINE_BUFFER = 60; // Seconds to extend calldata deadline \ No newline at end of file +export const UR_EXECUTE_DEADLINE_BUFFER = 60; // Seconds to extend calldata deadline +export const UR_UNWRAP_WETH_PARAMETERS = ['address', 'uint256'] +export const UR_SWEEP_PARAMETERS = ['address', 'address', 'uint256'] \ No newline at end of file diff --git a/lib/util/UniversalRouterCalldata.ts b/lib/util/UniversalRouterCalldata.ts index 6b1d8386..e2f38928 100644 --- a/lib/util/UniversalRouterCalldata.ts +++ b/lib/util/UniversalRouterCalldata.ts @@ -1,7 +1,15 @@ import { CommandType } from "@uniswap/universal-router-sdk"; import { Logger } from '@aws-lambda-powertools/logger' import { defaultAbiCoder, Interface } from "ethers/lib/utils"; -import { UR_EXECUTE_DEADLINE_BUFFER, UR_EXECUTE_FUNCTION, UR_EXECUTE_SELECTOR, UR_EXECUTE_WITH_DEADLINE_SELECTOR, UR_FUNCTION_SIGNATURES } from "../handlers/constants"; +import { + UR_EXECUTE_DEADLINE_BUFFER, + UR_EXECUTE_FUNCTION, + UR_EXECUTE_SELECTOR, + UR_EXECUTE_WITH_DEADLINE_SELECTOR, + UR_FUNCTION_SIGNATURES, + UR_SWEEP_PARAMETERS, + UR_UNWRAP_WETH_PARAMETERS +} from "../handlers/constants"; export class UniversalRouterCalldata { private iface: Interface; @@ -68,12 +76,12 @@ export class UniversalRouterCalldata { const sweepInput = this.inputsArray[sweepIndex]; // Decode sweep parameters const [token, , amountMinimum] = defaultAbiCoder.decode( - ['address', 'address', 'uint256'], + UR_SWEEP_PARAMETERS, sweepInput ); // Encode the parameters with new recipient address const modifiedSweepInput = defaultAbiCoder.encode( - ['address', 'address', 'uint256'], + UR_SWEEP_PARAMETERS, [token, recipient, amountMinimum] ); this.inputsArray[sweepIndex] = modifiedSweepInput; @@ -81,6 +89,25 @@ export class UniversalRouterCalldata { return this; } + public modifyUnwrapRecipient(recipient: string): UniversalRouterCalldata { + const unwrapIndex = this.commandArray.findIndex(command => command == CommandType.UNWRAP_WETH); + if (unwrapIndex !== -1) { + const unwrapInput = this.inputsArray[unwrapIndex]; + // Decode unwrap parameters + const [, amountMin] = defaultAbiCoder.decode( + UR_UNWRAP_WETH_PARAMETERS, + unwrapInput + ); + // Encode the parameters with new recipient address + const modifiedUnwrapInput = defaultAbiCoder.encode( + UR_UNWRAP_WETH_PARAMETERS, + [recipient, amountMin] + ); + this.inputsArray[unwrapIndex] = modifiedUnwrapInput; + } + return this; + } + public encode(): string { try { let modifiedCalldata; @@ -111,6 +138,7 @@ export function artemisModifyCalldata(calldata: string, log: Logger, executeAddr return router .removePayPortionCommand() .modifySweepRecipient(executeAddress) + .modifyUnwrapRecipient(executeAddress) .encode(); } catch (e) { log.error('Error in artemisModifyCalldata', { diff --git a/test/unit/handlers/get-unimind/get-unimind.test.ts b/test/unit/handlers/get-unimind/get-unimind.test.ts index 06535a06..bb1977ff 100644 --- a/test/unit/handlers/get-unimind/get-unimind.test.ts +++ b/test/unit/handlers/get-unimind/get-unimind.test.ts @@ -540,4 +540,24 @@ describe('Correctly modify URA calldata for Artemis support', () => { // Check that the SWEEP recipient is the executor address expect(sweepInput?.params.find(param => param.name === 'recipient')?.value).toEqual(EXECUTOR_ADDRESS) }) -}) + + it('artemisModifyCalldata for UNWRAP_WETH', () => { + // Contains V3_SWAP_EXACT_IN and UNWRAP_WETH + const calldata = "0x3593564c000000000000000000000000000000000000000000000000000000000000006000000000000000000000000000000000000000000000000000000000000000a00000000000000000000000000000000000000000000000000000000067e5b6550000000000000000000000000000000000000000000000000000000000000002000c000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000160000000000000000000000000000000000000000000000000000000000000010000000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000034de435f13194096b7ba79b000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000a00000000000000000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000002bc2b2ea7f6218cc37debbafe71361c088329ae09000271042000000000000000000000000000000000000060000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000040000000000000000000000000123ab21e11111a12abcd123f36fee12d21f42c7b00000000000000000000000000000000000000000000000005cff61d130d3651" + const decoded = CommandParser.parseCalldata(calldata) + const swapCommand = decoded.commands[0] + const unwrapCommand = decoded.commands[1] + + const modifiedCalldata = artemisModifyCalldata(calldata, mockLog, EXECUTOR_ADDRESS) + const modifiedDecoded = CommandParser.parseCalldata(modifiedCalldata) + const modifiedSwapCommand = modifiedDecoded.commands[0] + const modifiedUnwrapCommand = modifiedDecoded.commands[1] + expect(modifiedSwapCommand).toEqual(swapCommand) + + // Confirm the old UNWRAP_WETH command is not to the executor address + expect(unwrapCommand?.params.find(param => param.name === 'recipient')?.value).not.toEqual(EXECUTOR_ADDRESS) + // Check that the new UNWRAP_WETH recipient is the executor address + expect(modifiedUnwrapCommand).toBeDefined() + expect(modifiedUnwrapCommand?.params.find(param => param.name === 'recipient')?.value).toEqual(EXECUTOR_ADDRESS) + }) +}) \ No newline at end of file