enforces specified protocol amount

DePayFi · Sep 17, 2024 · 10dd5b5 · 10dd5b5
1 parent c4375b6
commit 10dd5b5
Show file tree

Hide file tree

Showing 3 changed files with 189 additions and 8 deletions.
diff --git a/contracts/DePayRouterV3.sol b/contracts/DePayRouterV3.sol
@@ -27,6 +27,7 @@ contract DePayRouterV3 is Ownable2Step {
   error PaymentToZeroAddressNotAllowed();
   error InsufficientBalanceInAfterPayment();
   error InsufficientBalanceOutAfterPayment();
+  error InsufficientProtocolAmount();
 
   /// @notice Address representing the NATIVE token (e.g. ETH, BNB, MATIC, etc.)
   address constant NATIVE = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
@@ -278,6 +279,20 @@ contract DePayRouterV3 is Ownable2Step {
         revert InsufficientBalanceOutAfterPayment();
       }
     }
+
+    // Ensure protocolAmount remained within router
+    if(payment.protocolAmount > 0) {
+
+      if(payment.tokenOutAddress == NATIVE) {
+        if((address(this).balance - payment.protocolAmount) < balanceOutBefore) {
+          revert InsufficientProtocolAmount();
+        }
+      } else {
+        if((IERC20(payment.tokenOutAddress).balanceOf(address(this)) - payment.protocolAmount) < balanceOutBefore) {
+          revert InsufficientProtocolAmount();
+        }
+      }
+    }
   }
 
   /// @dev Emits payment event.

diff --git a/test/_pay-with-native.ts b/test/_pay-with-native.ts
@@ -190,6 +190,56 @@ export default ({ blockchain })=>{
         expect(feeReceiverBalanceAfter).to.eq(feeReceiverBalanceBefore.add(feeAmount))
       })
 
+      it('pays payment receiver, fee receiver and protocol and emits Payment event to validate transfers easily', async ()=> {
+        const amountIn = 1000000000
+        const paymentAmount = 900000000
+        const feeAmount = 50000000
+        const protocolAmount = 50000000
+
+        const paymentReceiverBalanceBefore = await provider.getBalance(wallets[1].address)
+        const feeReceiverBalanceBefore = await provider.getBalance(wallets[2].address)
+        const routerBalanceBefore = await provider.getBalance(router.address)
+
+        await expect(
+          router.connect(wallets[0])[PAY]({
+            amountIn: amountIn,
+            paymentAmount: paymentAmount,
+            feeAmount: feeAmount,
+            protocolAmount: protocolAmount,
+            tokenInAddress: NATIVE,
+            exchangeAddress: ZERO,
+            tokenOutAddress: NATIVE,
+            paymentReceiverAddress: wallets[1].address,
+            feeReceiverAddress: wallets[2].address,
+            exchangeType: 0,
+            receiverType: 0,
+            exchangeCallData: ZERO,
+            receiverCallData: ZERO,
+            deadline,
+          }, { value: 1000000000 })
+        )
+        .to.emit(router, 'Payment').withArgs(
+          wallets[0].address, // from
+          wallets[1].address, // to
+          deadline, // deadline
+          amountIn,
+          paymentAmount,
+          feeAmount,
+          protocolAmount,
+          NATIVE,
+          NATIVE,
+          wallets[2].address
+        )
+
+        const paymentReceiverBalanceAfter = await provider.getBalance(wallets[1].address)
+        const feeReceiverBalanceAfter = await provider.getBalance(wallets[2].address)
+        const routerBalanceAfter = await provider.getBalance(router.address)
+
+        expect(paymentReceiverBalanceAfter).to.eq(paymentReceiverBalanceBefore.add(paymentAmount))
+        expect(feeReceiverBalanceAfter).to.eq(feeReceiverBalanceBefore.add(feeAmount))
+        expect(routerBalanceAfter).to.eq(routerBalanceBefore.add(protocolAmount))
+      })
+
       it('fails if balanceIn is less after payment', async()=>{
         await wallets[0].sendTransaction({ to: router.address, value: 1000000000 });
         await expect(
@@ -213,6 +263,40 @@ export default ({ blockchain })=>{
           'InsufficientBalanceInAfterPayment()'
         )
       })
+
+      it('fails if protocolAmount is less than specified', async()=>{
+
+        const amountIn = 1000000000
+        const paymentAmount = 900000000
+        const feeAmount = 50000000
+        const protocolAmount = 50000000
+
+        const paymentReceiverBalanceBefore = await provider.getBalance(wallets[1].address)
+        const feeReceiverBalanceBefore = await provider.getBalance(wallets[2].address)
+        const routerBalanceBefore = await provider.getBalance(router.address)
+
+        await expect(
+          router.connect(wallets[0])[PAY]({
+            amountIn: amountIn,
+            paymentAmount: paymentAmount,
+            feeAmount: feeAmount,
+            protocolAmount: 60000000,
+            tokenInAddress: NATIVE,
+            exchangeAddress: ZERO,
+            tokenOutAddress: NATIVE,
+            paymentReceiverAddress: wallets[1].address,
+            feeReceiverAddress: wallets[2].address,
+            exchangeType: 0,
+            receiverType: 0,
+            exchangeCallData: ZERO,
+            receiverCallData: ZERO,
+            deadline,
+          }, { value: amountIn })
+        ).to.be.revertedWith(
+          'InsufficientProtocolAmount()'
+        )
+
+      })
     })
   })
 }
diff --git a/test/_pay-with-token.ts b/test/_pay-with-token.ts
@@ -59,8 +59,8 @@ export default ({ blockchain, token, fromAccount, reversalReason })=>{
       })
 
       it('pays payment receiver', async ()=> {
-        const amountIn = 1000000000
-        const paymentAmount = 1000000000
+        const amountIn = 1000000
+        const paymentAmount = 1000000
 
         const paymentReceiverBalanceBefore = await tokenContract.balanceOf(wallets[1].address)
 
@@ -102,9 +102,9 @@ export default ({ blockchain, token, fromAccount, reversalReason })=>{
       })
 
       it('pays payment receiver and fee receiver', async ()=> {
-        const amountIn = 1000000000
-        const paymentAmount = 900000000
-        const feeAmount = 100000000
+        const amountIn = 100000
+        const paymentAmount = 90000
+        const feeAmount = 10000
 
         const paymentReceiverBalanceBefore = await tokenContract.balanceOf(wallets[1].address)
         const feeReceiverBalanceBefore = await tokenContract.balanceOf(wallets[2].address)
@@ -128,17 +128,70 @@ export default ({ blockchain, token, fromAccount, reversalReason })=>{
           deadline,
         })
 
+
         const paymentReceiverBalanceAfter = await tokenContract.balanceOf(wallets[1].address)
         const feeReceiverBalanceAfter = await tokenContract.balanceOf(wallets[2].address)
 
         expect(paymentReceiverBalanceAfter).to.eq(paymentReceiverBalanceBefore.add(paymentAmount))
         expect(feeReceiverBalanceAfter).to.eq(feeReceiverBalanceBefore.add(feeAmount))
       })
 
+      it('pays payment receiver, fee receiver and protocol', async ()=> {
+        const amountIn = 100000
+        const paymentAmount = 90000
+        const feeAmount = 5000
+        const protocolAmount = 5000
+
+        const paymentReceiverBalanceBefore = await tokenContract.balanceOf(wallets[1].address)
+        const feeReceiverBalanceBefore = await tokenContract.balanceOf(wallets[2].address)
+        const protocolBalanceBefore = await tokenContract.balanceOf(router.address)
+
+        await tokenContract.connect(fromAccount).approve(router.address, amountIn)
+
+        await expect(
+          router.connect(fromAccount)[PAY]({
+            amountIn: amountIn,
+            paymentAmount: paymentAmount,
+            feeAmount: feeAmount,
+            protocolAmount: protocolAmount,
+            tokenInAddress: TOKEN,
+            exchangeAddress: ZERO,
+            tokenOutAddress: TOKEN,
+            paymentReceiverAddress: wallets[1].address,
+            feeReceiverAddress: wallets[2].address,
+            exchangeType: 0,
+            receiverType: 0,
+            exchangeCallData: ZERO,
+            receiverCallData: ZERO,
+            deadline,
+          })
+        )
+        .to.emit(router, 'Payment').withArgs(
+          fromAccount._address, // from
+          wallets[1].address, // to
+          deadline, // deadline
+          amountIn,
+          paymentAmount,
+          feeAmount,
+          protocolAmount,
+          TOKEN,
+          TOKEN,
+          wallets[2].address
+        )
+
+        const paymentReceiverBalanceAfter = await tokenContract.balanceOf(wallets[1].address)
+        const feeReceiverBalanceAfter = await tokenContract.balanceOf(wallets[2].address)
+        const protocolBalanceAfter = await tokenContract.balanceOf(router.address)
+
+        expect(paymentReceiverBalanceAfter).to.eq(paymentReceiverBalanceBefore.add(paymentAmount))
+        expect(feeReceiverBalanceAfter).to.eq(feeReceiverBalanceBefore.add(feeAmount))
+        expect(protocolBalanceAfter).to.eq(protocolBalanceBefore.add(protocolAmount))
+      })
+
       it('fails if balanceIn is less after payment', async()=>{
-        const amountIn = 1000000000
-        const paymentAmount = 1000000000
-        const feeAmount = 100000000
+        const amountIn = 1000000
+        const paymentAmount = 1000000
+        const feeAmount = 100000
 
         const paymentReceiverBalanceBefore = await tokenContract.balanceOf(wallets[1].address)
         const feeReceiverBalanceBefore = await tokenContract.balanceOf(wallets[2].address)
@@ -168,6 +221,35 @@ export default ({ blockchain, token, fromAccount, reversalReason })=>{
           'InsufficientBalanceInAfterPayment()'
         )
       })
+
+      it('fails if protocolAmount is less than specified', async ()=> {
+        const amountIn = 100000
+        const paymentAmount = 90000
+        const feeAmount = 5000
+        const protocolAmount = 5000
+
+        await expect(
+          router.connect(fromAccount)[PAY]({
+            amountIn: amountIn,
+            paymentAmount: paymentAmount,
+            feeAmount: feeAmount,
+            protocolAmount: 6000,
+            tokenInAddress: TOKEN,
+            exchangeAddress: ZERO,
+            tokenOutAddress: TOKEN,
+            paymentReceiverAddress: wallets[1].address,
+            feeReceiverAddress: wallets[2].address,
+            exchangeType: 0,
+            receiverType: 0,
+            exchangeCallData: ZERO,
+            receiverCallData: ZERO,
+            deadline,
+          })
+        ).to.be.revertedWith(
+          'InsufficientProtocolAmount()'
+        )
+      })
+
     })
   })
 }