build: limit modules

contracts/test/mocks/ERC4626/LossyStrategy.sol

@@ -73,6 +73,10 @@ contract ERC4626LossyStrategy is ERC4626BaseStrategyMock {
         return IERC20(asset()).balanceOf(address(this)) - lockedFunds;
     }
 
+    function maxRedeem(address) public view override returns (uint256) {
+        return convertToShares(IERC20(asset()).balanceOf(address(this)) - lockedFunds);
+    }
+
     function migrate(address _newStrategy) external override {
         IERC20(asset()).safeTransfer(
             _newStrategy,

contracts/test/mocks/periphery/LimitModule.vy

@@ -0,0 +1,55 @@
+# @version 0.3.7
+
+interface IVault:
+    def totalAssets() -> uint256: view
+
+enforce_whitelist: public(bool)
+
+whitelist: public(HashMap[address, bool])
+
+default_deposit_limt: public(uint256)
+
+default_withdraw_limit: public(uint256)
+
+@external
+def __init__(
+    default_deposit_limt: uint256,
+    default_withdraw_limit: uint256,
+    enforce_whitelist: bool
+):
+    self.default_deposit_limt = default_deposit_limt
+    self.default_withdraw_limit = default_withdraw_limit
+    self.enforce_whitelist = enforce_whitelist
+
+@view
+@external
+def available_deposit_limit(receiver: address) -> uint256:
+    if self.enforce_whitelist:
+        if not self.whitelist[receiver]:
+            return 0
+
+    if self.default_deposit_limt == MAX_UINT256:
+        return MAX_UINT256
+
+    return self.default_deposit_limt - IVault(msg.sender).totalAssets()
+
+@view
+@external
+def available_withdraw_limit(owner: address, max_loss: uint256, strategies: DynArray[address, 10]) -> uint256:
+    return self.default_withdraw_limit
+
+@external
+def set_whitelist(list: address):
+    self.whitelist[list] = True
+
+@external
+def set_default_deposit_limit(limit: uint256):
+    self.default_deposit_limt = limit
+
+@external
+def set_default_withdraw_limit(limit: uint256):
+    self.default_withdraw_limit = limit
+
+@external
+def set_enforce_whitelist(enforce: bool):
+    self.enforce_whitelist = enforce

tests/conftest.py

@@ -215,6 +215,7 @@ def create_vault(
             | ROLES.DEBT_MANAGER
             | ROLES.MAX_DEBT_MANAGER
             | ROLES.DEPOSIT_LIMIT_MANAGER
+            | ROLES.WITHDRAW_LIMIT_MANAGER
             | ROLES.MINIMUM_IDLE_MANAGER
             | ROLES.PROFIT_UNLOCK_MANAGER
             | ROLES.DEBT_PURCHASER
@@ -353,6 +354,19 @@ def deploy_faulty_accountant(vault):
     yield deploy_faulty_accountant
 
 
+@pytest.fixture(scope="session")
+def deploy_limit_module(project, gov):
+    def deploy_limit_module(
+        deposit_limit=MAX_INT, withdraw_limit=MAX_INT, whitelist=False
+    ):
+        limit_module = gov.deploy(
+            project.LimitModule, deposit_limit, withdraw_limit, whitelist
+        )
+        return limit_module
+
+    yield deploy_limit_module
+
+
 @pytest.fixture(scope="session")
 def mint_and_deposit_into_strategy(gov, asset):
     def mint_and_deposit_into_strategy(

tests/unit/vault/test_emergency_shutdown.py

@@ -1,6 +1,6 @@
 import ape
 import pytest
-from utils.constants import ROLES
+from utils.constants import ROLES, ZERO_ADDRESS
 
 
 @pytest.fixture(autouse=True)
@@ -10,7 +10,8 @@ def set_role(vault, gov):
         ROLES.EMERGENCY_MANAGER
         | ROLES.ADD_STRATEGY_MANAGER
         | ROLES.DEBT_MANAGER
-        | ROLES.MAX_DEBT_MANAGER,
+        | ROLES.MAX_DEBT_MANAGER
+        | ROLES.DEPOSIT_LIMIT_MANAGER,
         sender=gov,
     )
 
@@ -34,7 +35,7 @@ def test_shutdown__increase_deposit_limit__reverts(
     mint_and_deposit_into_vault(vault, gov)
     vault.shutdown_vault(sender=gov)
 
-    assert vault.availableDepositLimit() == 0
+    assert vault.maxDeposit(gov) == 0
 
     vault.set_role(gov, ROLES.DEPOSIT_LIMIT_MANAGER, sender=gov)
 
@@ -45,7 +46,50 @@ def test_shutdown__increase_deposit_limit__reverts(
     with ape.reverts():
         vault.set_deposit_limit(limit, sender=gov)
 
-    assert vault.availableDepositLimit() == 0
+    assert vault.maxDeposit(gov) == 0
+
+
+def test_shutdown__set_deposit_limit_module__reverts(
+    vault, gov, asset, mint_and_deposit_into_vault, deploy_limit_module
+):
+    mint_and_deposit_into_vault(vault, gov)
+    vault.shutdown_vault(sender=gov)
+
+    assert vault.maxDeposit(gov) == 0
+
+    vault.set_role(gov, ROLES.DEPOSIT_LIMIT_MANAGER, sender=gov)
+
+    assert ROLES.DEPOSIT_LIMIT_MANAGER in ROLES(vault.roles(gov))
+
+    limit_module = deploy_limit_module()
+
+    with ape.reverts():
+        vault.set_deposit_limit_module(limit_module, sender=gov)
+
+    assert vault.maxDeposit(gov) == 0
+
+
+def test_shutdown__deposit_limit_module_is_removed(
+    create_vault, gov, asset, mint_and_deposit_into_vault, deploy_limit_module
+):
+    vault = create_vault(asset)
+
+    mint_and_deposit_into_vault(vault, gov)
+
+    limit_module = deploy_limit_module()
+    vault.set_deposit_limit_module(limit_module, sender=gov)
+
+    assert vault.maxDeposit(gov) > 0
+
+    tx = vault.shutdown_vault(sender=gov)
+
+    event = list(tx.decode_logs(vault.UpdateDepositLimitModule))
+
+    assert len(event) == 1
+    assert event[0].deposit_limit_module == ZERO_ADDRESS
+
+    assert vault.deposit_limit_module() == ZERO_ADDRESS
+    assert vault.maxDeposit(gov) == 0
 
 
 def test_shutdown_cant_deposit_can_withdraw(
@@ -54,7 +98,7 @@ def test_shutdown_cant_deposit_can_withdraw(
     mint_and_deposit_into_vault(vault, gov)
     vault.shutdown_vault(sender=gov)
 
-    assert vault.availableDepositLimit() == 0
+    assert vault.maxDeposit(gov) == 0
     vault_balance_before = asset.balanceOf(vault)
 
     with ape.reverts():
@@ -77,7 +121,8 @@ def test_strategy_return_funds(
     assert asset.balanceOf(strategy) == vault_balance
     assert asset.balanceOf(vault) == 0
     vault.shutdown_vault(sender=gov)
-    assert vault.availableDepositLimit() == 0
+
+    assert vault.maxDeposit(gov) == 0
 
     vault.update_debt(strategy.address, 0, sender=gov)
     assert asset.balanceOf(strategy) == 0