send protocolFeeSub to protocolFeeReceiver

masa-finance · Jul 27, 2023 · 818bb81 · 818bb81
1 parent 36847bb
commit 818bb81
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/contracts/dex/PaymentGateway.sol b/contracts/dex/PaymentGateway.sol
@@ -310,18 +310,23 @@ abstract contract PaymentGateway is AccessControl {
         uint256 protocolFee
     ) internal paymentParamsAlreadySet(amount.add(protocolFee)) {
         if (amount == 0 && protocolFee == 0) return;
-        if (protocolFee > 0 && protocolFeeReceiver == address(0))
-            revert ProtocolFeeReceiverNotSet();
+
+        uint256 protocolFeeSub = _getProtocolFeeSub(amount);
+
+        if (
+            (protocolFee > 0 || protocolFeeSub > 0) &&
+            protocolFeeReceiver == address(0)
+        ) revert ProtocolFeeReceiverNotSet();
 
         if (!enabledPaymentMethod[paymentMethod])
             revert InvalidPaymentMethod(paymentMethod);
         if (paymentMethod == address(0)) {
             // ETH
             if (msg.value < amount.add(protocolFee))
                 revert InsufficientEthAmount(amount.add(protocolFee));
-            if (amount > 0) {
+            if (amount.sub(protocolFeeSub) > 0) {
                 (bool success, ) = payable(projectFeeReceiver).call{
-                    value: amount
+                    value: amount.sub(protocolFeeSub)
                 }("");
                 if (!success) revert TransferFailed();
             }
@@ -331,6 +336,12 @@ abstract contract PaymentGateway is AccessControl {
                 }("");
                 if (!success) revert TransferFailed();
             }
+            if (protocolFeeSub > 0) {
+                (bool success, ) = payable(protocolFeeReceiver).call{
+                    value: protocolFeeSub
+                }("");
+                if (!success) revert TransferFailed();
+            }
             if (msg.value > amount.add(protocolFee)) {
                 // return diff
                 uint256 refund = msg.value.sub(amount.add(protocolFee));
@@ -339,11 +350,11 @@ abstract contract PaymentGateway is AccessControl {
             }
         } else {
             // ERC20 token, including MASA and USDC
-            if (amount > 0) {
+            if (amount.sub(protocolFeeSub) > 0) {
                 IERC20(paymentMethod).safeTransferFrom(
                     msg.sender,
                     projectFeeReceiver,
-                    amount
+                    amount.sub(protocolFeeSub)
                 );
             }
             if (protocolFee > 0) {
@@ -353,6 +364,13 @@ abstract contract PaymentGateway is AccessControl {
                     protocolFee
                 );
             }
+            if (protocolFeeSub > 0) {
+                IERC20(paymentMethod).safeTransferFrom(
+                    msg.sender,
+                    protocolFeeReceiver,
+                    protocolFeeSub
+                );
+            }
         }
     }