flash loan

contracts/DODOVendorMachine/impl/DVM.sol

@@ -22,7 +22,6 @@ contract DVM is DVMTrader, DVMFunding {
         address lpFeeRateModel,
         address mtFeeRateModel,
         address tradePermissionManager,
-        address fundingPermissionManager,
         uint256 i,
         uint256 k
     ) external {
@@ -33,7 +32,6 @@ contract DVM is DVMTrader, DVMFunding {
         _LP_FEE_RATE_MODEL_ = IFeeRateModel(lpFeeRateModel);
         _MT_FEE_RATE_MODEL_ = IFeeRateModel(mtFeeRateModel);
         _TRADE_PERMISSION_ = IPermissionManager(tradePermissionManager);
-        _FUNDING_PERMISSION_ = IPermissionManager(fundingPermissionManager);
         _MAINTAINER_ = maintainer;
         _I_ = i;
         _K_ = k;

contracts/DODOVendorMachine/impl/DVMStorage.sol

@@ -22,16 +22,14 @@ contract DVMStorage is InitializableOwnable, ReentrancyGuard {
 
     // ============ Variables for Control ============
 
-    bool public _CLOSED_;
     uint256 public _GAS_PRICE_LIMIT_;
 
     // ============ Advanced Controls ============
 
-    bool public _BUYING_ALLOWED_;
+    bool public _BUYING_CLOSE_;
-    bool public _SELLING_ALLOWED_;
+    bool public _SELLING_CLOSE_;
 
     IPermissionManager public _TRADE_PERMISSION_;
-    IPermissionManager public _FUNDING_PERMISSION_;
 
     // ============ Core Address ============
 
@@ -51,8 +49,16 @@ contract DVMStorage is InitializableOwnable, ReentrancyGuard {
 
     // ============ Modifiers ============
 
-    modifier notClosed() {
+    modifier isBuyAllow(address trader) {
-        require(!_CLOSED_, "DODO_CLOSED");
+        require(!_BUYING_CLOSE_ && _TRADE_PERMISSION_.isAllowed(trader), "TRADER_BUY_NOT_ALLOWED");
+        _;
+    }
+
+    modifier isSellAllow(address trader) {
+        require(
+            !_SELLING_CLOSE_ && _TRADE_PERMISSION_.isAllowed(trader),
+            "TRADER_SELL_NOT_ALLOWED"
+        );
         _;
     }
 
@@ -83,11 +89,19 @@ contract DVMStorage is InitializableOwnable, ReentrancyGuard {
         _TRADE_PERMISSION_ = IPermissionManager(newTradePermissionManager);
     }
 
-    function setFundingPermissionManager(address newFundingPermissionManager) external onlyOwner {
-        _FUNDING_PERMISSION_ = IPermissionManager(newFundingPermissionManager);
-    }
-
     function setMaintainer(address newMaintainer) external onlyOwner {
         _MAINTAINER_ = newMaintainer;
     }
+
+    function setGasPriceLimit(uint256 newGasPriceLimit) external onlyOwner {
+        _GAS_PRICE_LIMIT_ = newGasPriceLimit;
+    }
+
+    function setBuy(bool open) external onlyOwner {
+        _BUYING_CLOSE_ = !open;
+    }
+
+    function setSell(bool open) external onlyOwner {
+        _SELLING_CLOSE_ = !open;
+    }
 }

contracts/DODOVendorMachine/impl/DVMTrader.sol

@@ -17,7 +17,7 @@ import {IExternalCall} from "../intf/IExternalCall.sol";
 contract DVMTrader is DVMStorage {
     using SafeMath for uint256;
 
-    function sellBase(address to) external returns (uint256 receiveQuoteAmount) {
+    function sellBase(address to) external isSellAllow(to) returns (uint256 receiveQuoteAmount) {
         uint256 baseInput = _VAULT_.getBaseInput();
         uint256 mtFee;
         (receiveQuoteAmount, mtFee) = querySellBase(to, baseInput);
@@ -29,7 +29,7 @@ contract DVMTrader is DVMStorage {
         return receiveQuoteAmount;
     }
 
-    function sellQuote(address to) external returns (uint256 receiveBaseAmount) {
+    function sellQuote(address to) external isBuyAllow(to) returns (uint256 receiveBaseAmount) {
         uint256 quoteInput = _VAULT_.getQuoteInput();
         uint256 mtFee;
         (receiveBaseAmount, mtFee) = querySellQuote(to, quoteInput);
@@ -48,23 +48,37 @@ contract DVMTrader is DVMStorage {
         address call,
         bytes calldata data
     ) external {
-        (uint256 baseReserve, uint256 quoteReserve) = _VAULT_.getVaultReserve();
-        uint256 B0 = calculateBase0(baseReserve, quoteReserve);
-
-        uint256 mtFeeRate = _MT_FEE_RATE_MODEL_.getFeeRate(assetTo, quoteAmount);
-        uint256 baseMtFee = DecimalMath.mulCeil(baseAmount, mtFeeRate);
-        uint256 quoteMtFee = DecimalMath.mulCeil(quoteAmount, mtFeeRate);
-
-        _VAULT_.transferBaseOut(_MAINTAINER_, baseMtFee);
-        _VAULT_.transferQuoteOut(_MAINTAINER_, quoteMtFee);
         _VAULT_.transferBaseOut(assetTo, baseAmount);
         _VAULT_.transferQuoteOut(assetTo, quoteAmount);
-
         IExternalCall(call).DVMCall(data);
 
+        (uint256 baseReserve, uint256 quoteReserve) = _VAULT_.getVaultReserve();
         (uint256 baseBalance, uint256 quoteBalance) = _VAULT_.getVaultBalance();
-        uint256 newB0 = calculateBase0(baseBalance, quoteBalance);
+
-        require(newB0 >= B0, "FLASH_LOAN_FAILED");
+        uint256 mtFeeRate = _MT_FEE_RATE_MODEL_.getFeeRate(assetTo);
+        uint256 lpFeeRate = _LP_FEE_RATE_MODEL_.getFeeRate(assetTo);
+        if (baseBalance < baseReserve) {
+            uint256 validBaseOut = DecimalMath.divCeil(
+                baseReserve - baseBalance,
+                DecimalMath.ONE.sub(mtFeeRate).sub(lpFeeRate)
+            );
+            baseBalance = baseReserve.sub(validBaseOut);
+            _VAULT_.transferBaseOut(_MAINTAINER_, DecimalMath.mulCeil(validBaseOut, mtFeeRate));
+        }
+        if (quoteBalance < quoteReserve) {
+            uint256 validQuoteOut = DecimalMath.divCeil(
+                quoteReserve - quoteBalance,
+                DecimalMath.ONE.sub(mtFeeRate).sub(lpFeeRate)
+            );
+            quoteBalance = quoteReserve.sub(validQuoteOut);
+            _VAULT_.transferQuoteOut(_MAINTAINER_, DecimalMath.mulCeil(validQuoteOut, mtFeeRate));
+        }
+
+        require(
+            calculateBase0(baseBalance, quoteBalance) >= calculateBase0(baseReserve, quoteReserve),
+            "FLASH_LOAN_FAILED"
+        );
+
         _VAULT_.sync();
     }
 
@@ -80,8 +94,8 @@ contract DVMTrader is DVMStorage {
         require(B0 >= B1, "DODO_BASE_BALANCE_NOT_ENOUGH");
         uint256 Q = DODOMath._GeneralIntegrate(B0, B1, baseReserve, _I_, _K_);
 
-        uint256 lpFeeRate = _LP_FEE_RATE_MODEL_.getFeeRate(trader, Q);
+        uint256 lpFeeRate = _LP_FEE_RATE_MODEL_.getFeeRate(trader);
-        uint256 mtFeeRate = _MT_FEE_RATE_MODEL_.getFeeRate(trader, Q);
+        uint256 mtFeeRate = _MT_FEE_RATE_MODEL_.getFeeRate(trader);
         mtFee = DecimalMath.mulCeil(Q, mtFeeRate);
         receiveQuoteAmount = Q.sub(mtFee).sub(DecimalMath.mulCeil(Q, lpFeeRate));
 
@@ -105,8 +119,8 @@ contract DVMTrader is DVMStorage {
             _K_
         );
         uint256 deltaBase = baseReserve.sub(newBaseReserve);
-        uint256 lpFeeRate = _LP_FEE_RATE_MODEL_.getFeeRate(trader, payQuoteAmount);
+        uint256 lpFeeRate = _LP_FEE_RATE_MODEL_.getFeeRate(trader);
-        uint256 mtFeeRate = _MT_FEE_RATE_MODEL_.getFeeRate(trader, payQuoteAmount);
+        uint256 mtFeeRate = _MT_FEE_RATE_MODEL_.getFeeRate(trader);
         mtFee = DecimalMath.mulCeil(deltaBase, mtFeeRate);
         receiveBaseAmount = deltaBase.sub(mtFee).sub(DecimalMath.mulCeil(deltaBase, lpFeeRate));
         return (receiveBaseAmount, mtFee);

contracts/DODOVendorMachine/impl/DVMVault.sol

@@ -110,11 +110,15 @@ contract DVMVault is InitializableOwnable {
     }
 
     function transferBaseOut(address to, uint256 amount) public onlyOwner {
-        IERC20(_BASE_TOKEN_).safeTransfer(to, amount);
+        if (amount > 0) {
+            IERC20(_BASE_TOKEN_).safeTransfer(to, amount);
+        }
     }
 
     function transferQuoteOut(address to, uint256 amount) public onlyOwner {
-        IERC20(_QUOTE_TOKEN_).safeTransfer(to, amount);
+        if (amount > 0) {
+            IERC20(_QUOTE_TOKEN_).safeTransfer(to, amount);
+        }
     }
 
     // Shares related

contracts/DODOVendorMachine/intf/IDVM.sol

@@ -16,7 +16,6 @@ interface IDVM {
         address lpFeeRateModel,
         address mtFeeRateModel,
         address tradePermissionManager,
-        address fundingPermissionManager,
         uint256 i,
         uint256 k
     ) external;

contracts/Factory/DVMFactory.sol

@@ -59,7 +59,6 @@ contract DVMFactory is Ownable {
             createConstFeeRateModel(msg.sender, lpFeeRate),
             createConstFeeRateModel(msg.sender, mtFeeRate),
             createPermissionManager(msg.sender),
-            createPermissionManager(msg.sender),
             i,
             k
         );

contracts/intf/IFeeRateModel.sol

@@ -9,5 +9,5 @@ pragma solidity 0.6.9;
 pragma experimental ABIEncoderV2;
 
 interface IFeeRateModel {
-    function getFeeRate(address trader, uint256 amount) external view returns (uint256);
+    function getFeeRate(address trader) external view returns (uint256);
 }

contracts/lib/ConstFeeRateModel.sol

@@ -12,12 +12,10 @@ import {IFeeRateModel} from "../intf/IFeeRateModel.sol";
 import {Ownable} from "../lib/Ownable.sol";
 import {InitializableOwnable} from "../lib/InitializableOwnable.sol";
 
-interface IConstFeeRateModel {
+interface IConstFeeRateModel is IFeeRateModel {
     function init(address owner, uint256 feeRate) external;
 
     function setFeeRate(uint256 newFeeRate) external;
-
-    function getFeeRate(address, uint256) external view returns (uint256);
 }
 
 contract ConstFeeRateModel is InitializableOwnable, IFeeRateModel {
@@ -32,7 +30,7 @@ contract ConstFeeRateModel is InitializableOwnable, IFeeRateModel {
         _FEE_RATE_ = newFeeRate;
     }
 
-    function getFeeRate(address, uint256) external override view returns (uint256) {
+    function getFeeRate(address) external override view returns (uint256) {
         return _FEE_RATE_;
     }
 }