From 5860b655f45398d21400c87265dcd3dc2ab2caf4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?lz=2Esir=CE=94rthurmoney=28=29?= <kyle@layerzerolabs.org>
Date: Tue, 7 Nov 2023 13:20:27 -0800
Subject: [PATCH] refactor

---
 .../token/oft/v2/fee/NativeOFTWithFee.sol     | 25 +++++-----
 test/oft/v2/NativeOFTWithFee.test.js          | 47 +++++++++++++++++++
 2 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/contracts/token/oft/v2/fee/NativeOFTWithFee.sol b/contracts/token/oft/v2/fee/NativeOFTWithFee.sol
index f498fd04..41f4c0e2 100644
--- a/contracts/token/oft/v2/fee/NativeOFTWithFee.sol
+++ b/contracts/token/oft/v2/fee/NativeOFTWithFee.sol
@@ -65,18 +65,19 @@ contract NativeOFTWithFee is OFTWithFee, ReentrancyGuard {
 
     function _debitFromNative(address _from, uint _amount, uint16 _dstChainId) internal returns (uint messageFee, uint amount) {
         uint fee = quoteOFTFee(_dstChainId, _amount);
-
-        // subtract fee from _amount
-        _amount -= fee;
-
-        // pay fee and update newMsgValue
-        uint newMsgValue;
-        if(balanceOf(_from) > fee) {
-            _transferFrom(_from, feeOwner, fee);
-            newMsgValue = msg.value;
-        } else {
-            _mint(feeOwner, fee);
-            newMsgValue = msg.value - fee;
+        uint newMsgValue = msg.value;
+
+        if(fee > 0) {
+            // subtract fee from _amount
+            _amount -= fee;
+
+            // pay fee and update newMsgValue
+            if(balanceOf(_from) >= fee) {
+                _transferFrom(_from, feeOwner, fee);
+            } else {
+                _mint(feeOwner, fee);
+                newMsgValue = msg.value - fee;
+            }
         }
 
         (amount,) = _removeDust(_amount);
diff --git a/test/oft/v2/NativeOFTWithFee.test.js b/test/oft/v2/NativeOFTWithFee.test.js
index 107ff4b8..49e117f5 100644
--- a/test/oft/v2/NativeOFTWithFee.test.js
+++ b/test/oft/v2/NativeOFTWithFee.test.js
@@ -162,6 +162,53 @@ describe("NativeOFTWithFee: ", function () {
         expect(await nativeOFTWithFee.balanceOf(nativeOFTWithFee.address)).to.be.equal(totalAmount.sub(fee))
     })
 
+    it("sendFrom() w/ fee change - tokens from main to other chain without taking dust", async function () {
+        expect(await ethers.provider.getBalance(localEndpoint.address)).to.be.equal(ethers.utils.parseEther("0"))
+
+        // set default fee to 50%
+        await nativeOFTWithFee.setDefaultFeeBp(1)
+        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 leftOverAmount = ethers.utils.parseEther("0")
+        let totalAmount = ethers.utils.parseEther("8.123456789")
+
+        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
+
+
+        let ld2sdRate = 10 ** (18 - sharedDecimals)
+        let dust = totalAmount.sub(fee).mod(ld2sdRate)
+        let totalMintAmount = (totalAmount.sub(fee)).sub(dust)
+
+        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)
+            totalMintAmount, // 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(localEndpoint.address)).to.be.equal(nativeFee) // collects
+        expect(await nativeOFTWithFee.balanceOf(owner.address)).to.be.equal(leftOverAmount)
+        expect(await nativeOFTWithFee.balanceOf(bob.address)).to.be.equal(fee)
+        expect(await nativeOFTWithFee.balanceOf(nativeOFTWithFee.address)).to.be.equal(totalMintAmount)
+        expect(await remoteOFTWithFee.balanceOf(alice.address)).to.be.equal(totalMintAmount)
+    })
+
     it("sendFrom() w/ fee change - deposit before send", async function () {
         expect(await ethers.provider.getBalance(localEndpoint.address)).to.be.equal(ethers.utils.parseEther("0"))