build: pre and post hooks

yearn · Jan 20, 2024 · 327dcbd · 327dcbd
1 parent b1ed340
commit 327dcbd
Show file tree

Hide file tree

Showing 6 changed files with 499 additions and 7 deletions.
diff --git a/src/HealthCheck/BaseHealthCheck.sol → src/Bases/HealthCheck/BaseHealthCheck.sol b/src/HealthCheck/BaseHealthCheck.sol → src/Bases/HealthCheck/BaseHealthCheck.sol
diff --git a/src/HealthCheck/IBaseHealthCheck.sol → src/Bases/HealthCheck/IBaseHealthCheck.sol b/src/HealthCheck/IBaseHealthCheck.sol → src/Bases/HealthCheck/IBaseHealthCheck.sol
diff --git a/src/Bases/Hooks/BaseHooks.sol b/src/Bases/Hooks/BaseHooks.sol
@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: AGPL-3.0
+pragma solidity 0.8.18;
+
+import {BaseHealthCheck, ERC20} from "../HealthCheck/BaseHealthCheck.sol";
+
+abstract contract DepositHooks {
+    function _preDepositHook(
+        uint256 assets,
+        uint256 shares,
+        address receiver
+    ) internal virtual {}
+
+    function _postDepositHook(
+        uint256 assets,
+        uint256 shares,
+        address receiver
+    ) internal virtual {}
+}
+
+abstract contract WithdrawHooks {
+    function _preWithdrawHook(
+        uint256 assets,
+        uint256 shares,
+        address receiver,
+        address owner,
+        uint256 maxLoss
+    ) internal virtual {}
+
+    function _postWithdrawHook(
+        uint256 assets,
+        uint256 shares,
+        address receiver,
+        address owner,
+        uint256 maxLoss
+    ) internal virtual {}
+}
+
+abstract contract TransferHooks {
+    function _preTransferHook(
+        address from,
+        address to,
+        uint256 amount
+    ) internal virtual {}
+
+    function _postTransferHook(
+        address from,
+        address to,
+        uint256 amount,
+        bool success
+    ) internal virtual {}
+}
+
+abstract contract Hooks is DepositHooks, WithdrawHooks, TransferHooks {}
+
+/**
+ *   @title Base Hooks
+ *   @author Yearn.finance
+ *   @notice This contract can be inherited by any Yearn
+ *   strategy wishing to implement pre or post deposit, withdraw
+ *   or transfer hooks in their strategy.
+ */
+abstract contract BaseHooks is BaseHealthCheck, Hooks {
+    constructor(
+        address _asset,
+        string memory _name
+    ) BaseHealthCheck(_asset, _name) {}
+
+    // Deposit
+    function deposit(
+        uint256 assets,
+        address receiver
+    ) external virtual returns (uint256 shares) {
+        _preDepositHook(assets, shares, receiver);
+        shares = abi.decode(
+            _delegateCall(
+                abi.encodeCall(TokenizedStrategy.deposit, (assets, receiver))
+            ),
+            (uint256)
+        );
+        _postDepositHook(assets, shares, receiver);
+    }
+
+    // Mint
+    function mint(
+        uint256 shares,
+        address receiver
+    ) external virtual returns (uint256 assets) {
+        _preDepositHook(assets, shares, receiver);
+        assets = abi.decode(
+            _delegateCall(
+                abi.encodeCall(TokenizedStrategy.mint, (shares, receiver))
+            ),
+            (uint256)
+        );
+        _postDepositHook(assets, shares, receiver);
+    }
+
+    // Withdraw
+    function withdraw(
+        uint256 assets,
+        address receiver,
+        address owner
+    ) external virtual returns (uint256 shares) {
+        return withdraw(assets, receiver, owner, 0);
+    }
+
+    function withdraw(
+        uint256 assets,
+        address receiver,
+        address owner,
+        uint256 maxLoss
+    ) public virtual returns (uint256 shares) {
+        _preWithdrawHook(assets, shares, receiver, owner, maxLoss);
+        shares = abi.decode(
+            _delegateCall(
+                // Have to use encodeWithSignature due to overloading parameters.
+                abi.encodeWithSignature(
+                    "withdraw(uint256,address,address,uint256)",
+                    assets,
+                    receiver,
+                    owner,
+                    maxLoss
+                )
+            ),
+            (uint256)
+        );
+        _postWithdrawHook(assets, shares, receiver, owner, maxLoss);
+    }
+
+    // Redeem
+    function redeem(
+        uint256 shares,
+        address receiver,
+        address owner
+    ) external virtual returns (uint256) {
+        // We default to not limiting a potential loss.
+        return redeem(shares, receiver, owner, MAX_BPS);
+    }
+
+    function redeem(
+        uint256 shares,
+        address receiver,
+        address owner,
+        uint256 maxLoss
+    ) public returns (uint256 assets) {
+        _preWithdrawHook(assets, shares, receiver, owner, maxLoss);
+        assets = abi.decode(
+            _delegateCall(
+                // Have to use encodeWithSignature due to overloading parameters.
+                abi.encodeWithSignature(
+                    "redeem(uint256,address,address,uint256)",
+                    shares,
+                    receiver,
+                    owner,
+                    maxLoss
+                )
+            ),
+            (uint256)
+        );
+        _postWithdrawHook(assets, shares, receiver, owner, maxLoss);
+    }
+
+    // Transfer
+    function transferFrom(
+        address from,
+        address to,
+        uint256 amount
+    ) public virtual returns (bool success) {
+        _preTransferHook(from, to, amount);
+        success = abi.decode(
+            _delegateCall(
+                abi.encodeCall(
+                    TokenizedStrategy.transferFrom,
+                    (from, to, amount)
+                )
+            ),
+            (bool)
+        );
+        _postTransferHook(from, to, amount, success);
+    }
+
+    // Transfer from
+    function transfer(
+        address to,
+        uint256 amount
+    ) external virtual returns (bool success) {
+        _preTransferHook(msg.sender, to, amount);
+        success = abi.decode(
+            _delegateCall(
+                abi.encodeCall(TokenizedStrategy.transfer, (to, amount))
+            ),
+            (bool)
+        );
+        _postTransferHook(msg.sender, to, amount, success);
+    }
+}
diff --git a/src/test/BaseHook.t.sol b/src/test/BaseHook.t.sol
@@ -0,0 +1,197 @@
+// SPDX-License-Identifier: AGPL-3.0
+pragma solidity 0.8.18;
+
+import {Setup, IStrategy} from "./utils/Setup.sol";
+
+import {MockHooks, HookEvents} from "./mocks/MockHooks.sol";
+
+contract BaseHookTest is Setup, HookEvents {
+    function setUp() public override {
+        super.setUp();
+
+        mockStrategy = IStrategy(address(new MockHooks(address(asset))));
+
+        mockStrategy.setKeeper(keeper);
+        mockStrategy.setPerformanceFeeRecipient(performanceFeeRecipient);
+        mockStrategy.setPendingManagement(management);
+        // Accept management.
+        vm.prank(management);
+        mockStrategy.acceptManagement();
+    }
+
+    function test_depositHooks(uint256 _amount) public {
+        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
+
+        airdrop(asset, user, _amount);
+
+        vm.prank(user);
+        asset.approve(address(mockStrategy), _amount);
+
+        // Make sure we get both events with the correct amounts.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre deposit wont have a shares amount yet
+        emit PreDepositHook(_amount, 0, user);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostDepositHook(_amount, _amount, user);
+
+        vm.prank(user);
+        mockStrategy.deposit(_amount, user);
+
+        assertEq(mockStrategy.balanceOf(user), _amount);
+    }
+
+    function test_mintHooks(uint256 _amount) public {
+        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
+
+        airdrop(asset, user, _amount);
+
+        vm.prank(user);
+        asset.approve(address(mockStrategy), _amount);
+
+        // Make sure we get both events with the correct amounts.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre mint wont have a assets amount yet
+        emit PreDepositHook(0, _amount, user);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostDepositHook(_amount, _amount, user);
+
+        vm.prank(user);
+        mockStrategy.mint(_amount, user);
+
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+    }
+
+    function test_withdrawHooks(uint256 _amount) public {
+        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
+
+        mintAndDepositIntoStrategy(mockStrategy, user, _amount);
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+
+        // Make sure we get both events with the correct amounts.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre withdraw wont have a shares amount yet
+        emit PreWithdrawHook(_amount, 0, user, user, 0);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostWithdrawHook(_amount, _amount, user, user, 0);
+
+        vm.prank(user);
+        mockStrategy.withdraw(_amount, user, user);
+
+        checkStrategyTotals(mockStrategy, 0, 0, 0);
+
+        // Deposit back in
+        mintAndDepositIntoStrategy(mockStrategy, user, _amount);
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+
+        // Make sure works on both withdraw versions.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre withdraw wont have a shares amount yet
+        emit PreWithdrawHook(_amount, 0, user, user, 8);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostWithdrawHook(_amount, _amount, user, user, 8);
+
+        vm.prank(user);
+        mockStrategy.withdraw(_amount, user, user, 8);
+
+        checkStrategyTotals(mockStrategy, 0, 0, 0);
+    }
+
+    function test_redeemHooks(uint256 _amount) public {
+        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
+
+        mintAndDepositIntoStrategy(mockStrategy, user, _amount);
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+
+        // Make sure we get both events with the correct amounts.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre withdraw wont have a shares amount yet
+        emit PreWithdrawHook(0, _amount, user, user, MAX_BPS);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostWithdrawHook(_amount, _amount, user, user, MAX_BPS);
+
+        vm.prank(user);
+        mockStrategy.redeem(_amount, user, user);
+        checkStrategyTotals(mockStrategy, 0, 0, 0);
+
+        // Deposit back in
+        mintAndDepositIntoStrategy(mockStrategy, user, _amount);
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+
+        // Make sure works on both withdraw versions.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre withdraw wont have a shares amount yet
+        emit PreWithdrawHook(0, _amount, user, user, 8);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostWithdrawHook(_amount, _amount, user, user, 8);
+
+        vm.prank(user);
+        mockStrategy.redeem(_amount, user, user, 8);
+
+        checkStrategyTotals(mockStrategy, 0, 0, 0);
+    }
+
+    function test_transferHooks(uint256 _amount) public {
+        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
+
+        mintAndDepositIntoStrategy(mockStrategy, user, _amount);
+
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        assertEq(mockStrategy.balanceOf(management), 0);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+
+        // Make sure we get both events with the correct amounts.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre withdraw wont have a shares amount yet
+        emit PreTransferHook(user, management, _amount);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostTransferHook(user, management, _amount, true);
+
+        vm.prank(user);
+        mockStrategy.transfer(management, _amount);
+
+        assertEq(mockStrategy.balanceOf(user), 0);
+        assertEq(mockStrategy.balanceOf(management), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+    }
+
+    function test_transferFromHooks(uint256 _amount) public {
+        vm.assume(_amount >= minFuzzAmount && _amount <= maxFuzzAmount);
+
+        mintAndDepositIntoStrategy(mockStrategy, user, _amount);
+
+        assertEq(mockStrategy.balanceOf(user), _amount);
+        assertEq(mockStrategy.balanceOf(management), 0);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+
+        // Approve daddy to move funds
+        vm.prank(user);
+        mockStrategy.approve(daddy, _amount);
+
+        // Make sure we get both events with the correct amounts.
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        // Pre withdraw wont have a shares amount yet
+        emit PreTransferHook(user, management, _amount);
+
+        vm.expectEmit(true, true, true, true, address(mockStrategy));
+        emit PostTransferHook(user, management, _amount, true);
+
+        vm.prank(daddy);
+        mockStrategy.transferFrom(user, management, _amount);
+
+        assertEq(mockStrategy.balanceOf(user), 0);
+        assertEq(mockStrategy.balanceOf(management), _amount);
+        checkStrategyTotals(mockStrategy, _amount, 0, _amount);
+    }
+}