From fd94eb9a1886c26cbe2ab792e3a643ffdaa9bfdd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?lz=2Esir=CE=94rthurmoney=28=29?= <kyle@layerzerolabs.org>
Date: Mon, 6 Nov 2023 15:40:10 -0800
Subject: [PATCH] Adding in update to NativeOFTWithFee

---
 .../token/oft/v2/fee/NativeOFTWithFee.sol     | 15 +++++++
 test/oft/v2/NativeOFTWithFee.test.js          | 42 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/contracts/token/oft/v2/fee/NativeOFTWithFee.sol b/contracts/token/oft/v2/fee/NativeOFTWithFee.sol
index 0f5e7e8a..c38abf0c 100644
--- a/contracts/token/oft/v2/fee/NativeOFTWithFee.sol
+++ b/contracts/token/oft/v2/fee/NativeOFTWithFee.sol
@@ -25,6 +25,21 @@ contract NativeOFTWithFee is OFTWithFee, ReentrancyGuard {
         emit Withdrawal(msg.sender, _amount);
     }
 
+    /************************************************************************
+    * public functions
+    ************************************************************************/
+    function sendFrom(address _from, uint16 _dstChainId, bytes32 _toAddress, uint _amount, uint _minAmount, LzCallParams calldata _callParams) public payable virtual override {
+        _amount = _send(_from, _dstChainId, _toAddress, _amount, _callParams.refundAddress, _callParams.zroPaymentAddress, _callParams.adapterParams);
+        (_amount,) = _payOFTFee(address(this), _dstChainId, _amount);
+        require(_amount >= _minAmount, "BaseOFTWithFee: amount is less than minAmount");
+    }
+
+    function sendAndCall(address _from, uint16 _dstChainId, bytes32 _toAddress, uint _amount, uint _minAmount, bytes calldata _payload, uint64 _dstGasForCall, LzCallParams calldata _callParams) public payable virtual override {
+        _amount = _sendAndCall(_from, _dstChainId, _toAddress, _amount, _payload, _dstGasForCall, _callParams.refundAddress, _callParams.zroPaymentAddress, _callParams.adapterParams);
+        (_amount,) = _payOFTFee(address(this), _dstChainId, _amount);
+        require(_amount >= _minAmount, "BaseOFTWithFee: amount is less than minAmount");
+    }
+
     function _send(address _from, uint16 _dstChainId, bytes32 _toAddress, uint _amount, address payable _refundAddress, address _zroPaymentAddress, bytes memory _adapterParams) internal virtual override returns (uint amount) {
         _checkGasLimit(_dstChainId, PT_SEND, _adapterParams, NO_EXTRA_GAS);
 
diff --git a/test/oft/v2/NativeOFTWithFee.test.js b/test/oft/v2/NativeOFTWithFee.test.js
index 5a7b7fbd..156c8654 100644
--- a/test/oft/v2/NativeOFTWithFee.test.js
+++ b/test/oft/v2/NativeOFTWithFee.test.js
@@ -133,6 +133,48 @@ describe("NativeOFTWithFee: ", function () {
         expect(await remoteOFTWithFee.balanceOf(owner.address)).to.equal(0)
         expect(await ethers.provider.getBalance(nativeOFTWithFee.address)).to.equal(0)
 
+        let leftOverAmount = ethers.utils.parseEther("0")
+        let totalAmount = ethers.utils.parseEther("8")
+
+        expect(await ethers.provider.getBalance(localEndpoint.address)).to.be.equal(0)
+        expect(await nativeOFTWithFee.balanceOf(nativeOFTWithFee.address)).to.be.equal(0)
+        expect(await nativeOFTWithFee.balanceOf(owner.address)).to.be.equal(0)
+        expect(await remoteOFTWithFee.balanceOf(owner.address)).to.be.equal(0)
+
+        const aliceAddressBytes32 = ethers.utils.defaultAbiCoder.encode(["address"], [alice.address])
+        // estimate nativeFees
+        let fee = await nativeOFTWithFee.quoteOFTFee(remoteChainId, totalAmount)
+        let nativeFee = (await nativeOFTWithFee.estimateSendFee(remoteChainId, aliceAddressBytes32, totalAmount, false, defaultAdapterParams))
+            .nativeFee
+        await nativeOFTWithFee.sendFrom(
+            owner.address,
+            remoteChainId, // destination chainId
+            aliceAddressBytes32, // destination address to send tokens to
+            totalAmount, // quantity of tokens to send (in units of wei)
+            totalAmount.sub(fee), // quantity of tokens to send (in units of wei)
+            [owner.address, ethers.constants.AddressZero, defaultAdapterParams],
+            { value: nativeFee.add(totalAmount) } // pass a msg.value to pay the LayerZero message fee
+        )
+        expect(await ethers.provider.getBalance(nativeOFTWithFee.address)).to.be.equal(totalAmount)
+        expect(await ethers.provider.getBalance(localEndpoint.address)).to.be.equal(nativeFee) // collects
+        expect(await nativeOFTWithFee.balanceOf(owner.address)).to.be.equal(leftOverAmount)
+        expect(await nativeOFTWithFee.balanceOf(alice.address)).to.be.equal(leftOverAmount)
+        expect(await nativeOFTWithFee.balanceOf(bob.address)).to.be.equal(fee)
+        expect(await nativeOFTWithFee.balanceOf(nativeOFTWithFee.address)).to.be.equal(totalAmount.sub(fee))
+    })
+
+    it("sendFrom() w/ fee change - deposit before send", async function () {
+        expect(await ethers.provider.getBalance(localEndpoint.address)).to.be.equal(ethers.utils.parseEther("0"))
+
+        // set default fee to 50%
+        await nativeOFTWithFee.setDefaultFeeBp(5000)
+        await nativeOFTWithFee.setFeeOwner(bob.address)
+
+        // ensure they're both allocated initial amounts
+        expect(await nativeOFTWithFee.balanceOf(owner.address)).to.equal(0)
+        expect(await remoteOFTWithFee.balanceOf(owner.address)).to.equal(0)
+        expect(await ethers.provider.getBalance(nativeOFTWithFee.address)).to.equal(0)
+
         let depositAmount = ethers.utils.parseEther("7")
         await nativeOFTWithFee.deposit({ value: depositAmount })