From cfb90210dad3dd6ad61fb9aece5ffc0a3e957aed Mon Sep 17 00:00:00 2001 From: dsudit01 <79417718+dsudit01@users.noreply.github.com> Date: Fri, 5 Apr 2024 15:45:47 -0400 Subject: [PATCH] check that the lz endpoint is the only address calling lz Compose, fix tests to mock this --- contracts/asd/asdRouter.sol | 6 +- contracts/test-contracts/MockLZEndpoint.sol | 11 ++++ test/ASDRouter.js | 67 ++++++++++++--------- 3 files changed, 55 insertions(+), 29 deletions(-) create mode 100644 contracts/test-contracts/MockLZEndpoint.sol diff --git a/contracts/asd/asdRouter.sol b/contracts/asd/asdRouter.sol index 5aaedfd..6591d12 100644 --- a/contracts/asd/asdRouter.sol +++ b/contracts/asd/asdRouter.sol @@ -22,6 +22,7 @@ contract ASDRouter is IOAppComposer { // canto chain params address public cNoteAddress; uint32 public cantoLzEID; + address public cantoLzEndpoint; // asdUSDC contract for swapping to $NOTE address public asdUSDC; @@ -41,13 +42,14 @@ contract ASDRouter is IOAppComposer { event ASDSent(bytes32 indexed _guid, address _to, address _asdAddress, uint _amount, uint32 _dstEid, bool _lzSend); - constructor(address _cNoteAddress, uint32 _cantoLzEID, address _crocSwapAddress, address _crocImpactAddress, address _asdUSDCAddress) { + constructor(address _cNoteAddress, uint32 _cantoLzEID, address _crocSwapAddress, address _crocImpactAddress, address _asdUSDCAddress, address _cantoLzEndpoint) { cNoteAddress = _cNoteAddress; cantoLzEID = _cantoLzEID; crocSwapAddress = _crocSwapAddress; crocImpactAddress = _crocImpactAddress; asdUSDC = _asdUSDCAddress; ASDUSDC(_asdUSDCAddress).approve(crocSwapAddress, type(uint).max); + cantoLzEndpoint = _cantoLzEndpoint; } function _getNoteAddress() internal returns (address) { @@ -64,6 +66,8 @@ contract ASDRouter is IOAppComposer { * @dev Cannot revert anywhere, must send the tokens to the intended receiver if something fails (token's will be lost otherwise) */ function lzCompose(address _from, bytes32 _guid, bytes calldata _message, address _executor, bytes calldata _extraData) external payable { + // only time this function will revert if the caller is incorrect (only lz endpoint can call this) + require(msg.sender == cantoLzEndpoint, "ASDROUTER: only lz endpoint"); /* log event */ emit LZReceived(_guid, _from, _message, _executor, _extraData, msg.value); diff --git a/contracts/test-contracts/MockLZEndpoint.sol b/contracts/test-contracts/MockLZEndpoint.sol new file mode 100644 index 0000000..a6dc2c7 --- /dev/null +++ b/contracts/test-contracts/MockLZEndpoint.sol @@ -0,0 +1,11 @@ +pragma solidity ^0.8.20; + +interface ILayerZeroComposer { + function lzCompose(address _from, bytes32 _guid, bytes calldata _message, address _executor, bytes calldata _extraData) external payable; +} + +contract MockLZEndpoint { + function lzCompose(address _from, address _to, bytes32 _guid, uint16 _index, bytes calldata _message, bytes calldata _extraData) external payable { + ILayerZeroComposer(_to).lzCompose{value: msg.value}(_from, _guid, _message, msg.sender, _extraData); + } +} diff --git a/test/ASDRouter.js b/test/ASDRouter.js index 2a36337..d6607ab 100644 --- a/test/ASDRouter.js +++ b/test/ASDRouter.js @@ -1,7 +1,6 @@ const { expect } = require("chai"); const helpers = require("@nomicfoundation/hardhat-toolbox/network-helpers"); const { ethers } = require("hardhat"); -const LZ_ENDPOINTS = require("../constants/lzEndpoints.json"); const generatedComposeMsg = (from, amount, payload) => ethers.solidityPacked( @@ -23,9 +22,8 @@ const errorMessages = { }; describe("ASDRouter", function () { - const cantoLzEndpoint = LZ_ENDPOINTS["canto-testnet"]; + const cantoLzEndpointId = 1; - const executorAddress = "0xc0ffee254729296a45a3885639AC7E10F9d54979"; // random address const refundAddress = "0x9C29A5579EdfaA8F08dE82E805ea5744D9c50F9D"; // random address // testing contracts @@ -36,6 +34,7 @@ describe("ASDRouter", function () { let CrocSwap; let CrocImpact; let ASDUSDC; + let LzEndpoint; // test amounts const amountUSDCSent = ethers.parseEther("100"); @@ -50,13 +49,16 @@ describe("ASDRouter", function () { CrocImpact = await ethers.deployContract("MockCrocImpact", [CrocSwap.target]); // ASDUSDC contract ASDUSDC = await ethers.deployContract("ASDUSDC"); + // mock lz endpoint to call lzCompose from + LzEndpoint = await ethers.deployContract("MockLZEndpoint", []); // Router ASDRouter = await ethers.deployContract("ASDRouter", [ Note.target, - cantoLzEndpoint.id, + cantoLzEndpointId, CrocSwap.target, CrocImpact.target, ASDUSDC.target, + LzEndpoint.target, ]); // transfer USDC to router as if it was already sent @@ -75,7 +77,9 @@ describe("ASDRouter", function () { it("lzCompose: invalid payload", async () => { // call lzCompose with invalid payload - await expect(ASDRouter.lzCompose(USDCOFT.target, guid, generatedComposeMsg(refundAddress, amountUSDCSent, "0x"), executorAddress, "0x")) + await expect( + LzEndpoint.lzCompose(USDCOFT.target, ASDRouter.target, guid, 0, generatedComposeMsg(refundAddress, amountUSDCSent, "0x"), "0x") + ) .to.emit(ASDRouter, "TokenRefund") .withArgs(guid, USDCOFT.target, refundAddress, amountUSDCSent, 0, errorMessages.invalidPayload); @@ -87,7 +91,7 @@ describe("ASDRouter", function () { const gas = ethers.parseEther("1"); // call lzCompose with invalid payload await expect( - ASDRouter.lzCompose(USDCOFT.target, guid, generatedComposeMsg(refundAddress, amountUSDCSent, "0x"), executorAddress, "0x", { + LzEndpoint.lzCompose(USDCOFT.target, ASDRouter.target, guid, 0, generatedComposeMsg(refundAddress, amountUSDCSent, "0x"), "0x", { value: gas, }) ) @@ -101,15 +105,16 @@ describe("ASDRouter", function () { it("lzCompose: not whitelisted", async () => { // call lzCompose with un-whitelisted token await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, - generatedRouterPayload(cantoLzEndpoint.id, refundAddress, TESTASD.target, TESTASD.target, "0", refundAddress, "0") + generatedRouterPayload(cantoLzEndpointId, refundAddress, TESTASD.target, TESTASD.target, "0", refundAddress, "0") ), - executorAddress, "0x" ) ) @@ -123,15 +128,16 @@ describe("ASDRouter", function () { const gas = ethers.parseEther("1"); // call lzCompose with un-whitelisted token await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, - generatedRouterPayload(cantoLzEndpoint.id, refundAddress, TESTASD.target, TESTASD.target, "0", refundAddress, "0") + generatedRouterPayload(cantoLzEndpointId, refundAddress, TESTASD.target, TESTASD.target, "0", refundAddress, "0") ), - executorAddress, "0x", { value: gas } ) @@ -148,14 +154,16 @@ describe("ASDRouter", function () { await ASDUSDC.updateWhitelist(USDCOFT.target, true); // call lzCompose with minASD too high await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, generatedRouterPayload( - cantoLzEndpoint.id, + cantoLzEndpointId, refundAddress, TESTASD.target, TESTASD.target, @@ -164,7 +172,6 @@ describe("ASDRouter", function () { "0" ) ), - executorAddress, "0x" ) ) @@ -180,14 +187,16 @@ describe("ASDRouter", function () { const gas = ethers.parseEther("1"); // call lzCompose with minASD too high await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, generatedRouterPayload( - cantoLzEndpoint.id, + cantoLzEndpointId, refundAddress, TESTASD.target, TESTASD.target, @@ -196,7 +205,6 @@ describe("ASDRouter", function () { "0" ) ), - executorAddress, "0x", { value: gas } ) @@ -213,14 +221,16 @@ describe("ASDRouter", function () { await ASDUSDC.updateWhitelist(USDCOFT.target, true); // call lzCompose with msg.value less than send fee (dst not canto) await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, generatedRouterPayload( - cantoLzEndpoint.id + 1, + cantoLzEndpointId + 1, refundAddress, TESTASD.target, TESTASD.target, @@ -229,7 +239,6 @@ describe("ASDRouter", function () { ethers.parseEther("10").toString() // test fee ) ), - executorAddress, "0x" ) ) @@ -245,14 +254,16 @@ describe("ASDRouter", function () { const gas = ethers.parseEther("1"); // call lzCompose with msg.value less than send fee (dst not canto) await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, generatedRouterPayload( - cantoLzEndpoint.id + 1, + cantoLzEndpointId + 1, refundAddress, TESTASD.target, TESTASD.target, @@ -261,7 +272,6 @@ describe("ASDRouter", function () { ethers.parseEther("10").toString() // test fee ) ), - executorAddress, "0x", { value: gas } ) @@ -278,20 +288,21 @@ describe("ASDRouter", function () { await ASDUSDC.updateWhitelist(USDCOFT.target, true); // call lzCompose with valid payload await expect( - ASDRouter.lzCompose( + LzEndpoint.lzCompose( USDCOFT.target, + ASDRouter.target, guid, + 0, generatedComposeMsg( refundAddress, amountUSDCSent, - generatedRouterPayload(cantoLzEndpoint.id, refundAddress, TESTASD.target, TESTASD.target, "0", refundAddress, "0") + generatedRouterPayload(cantoLzEndpointId, refundAddress, TESTASD.target, TESTASD.target, "0", refundAddress, "0") ), - executorAddress, "0x" ) ) .to.emit(ASDRouter, "ASDSent") - .withArgs(guid, refundAddress, TESTASD.target, amountUSDCSent, cantoLzEndpoint.id, false); + .withArgs(guid, refundAddress, TESTASD.target, amountUSDCSent, cantoLzEndpointId, false); // expect ASD to be sent to canto expect(await TESTASD.balanceOf(refundAddress)).to.equal(amountUSDCSent); });