add overrides for deposit related functions (#112)

plumenetwork · Dec 10, 2024 · d6ffde7 · d6ffde7
1 parent d9a09dd
commit d6ffde7
Showing 1 changed file with 54 additions and 1 deletion.
diff --git a/nest/src/AggregateToken.sol b/nest/src/AggregateToken.sol
@@ -174,13 +174,66 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
         uint256 assets,
         address receiver,
         address controller
-    ) public override(ComponentToken, IComponentToken) returns (uint256 shares) {
+    ) public override(ComponentToken, IComponentToken, ERC4626Upgradeable) returns (uint256 shares) {
         if (_getAggregateTokenStorage().paused) {
             revert DepositPaused();
         }
         return super.deposit(assets, receiver, controller);
     }
 
+    /**
+     * @inheritdoc ERC4626Upgradeable
+     * @dev Overridden to add pause check before deposit
+     * @param assets Amount of assets to deposit
+     * @param receiver Address that will receive the shares
+     * @return shares Amount of shares minted
+     */
+    function deposit(
+        uint256 assets,
+        address receiver
+    ) public virtual override(ERC4626Upgradeable) returns (uint256 shares) {
+        if (_getAggregateTokenStorage().paused) {
+            revert DepositPaused();
+        }
+        return super.deposit(assets, receiver);
+    }
+
+    /**
+     * @inheritdoc ComponentToken
+     * @dev Overridden to add pause check before minting
+     * @param shares Amount of shares to mint
+     * @param receiver Address that will receive the shares
+     * @param controller Address that controls the minting
+     * @return assets Amount of assets deposited
+     */
+    function mint(
+        uint256 shares,
+        address receiver,
+        address controller
+    ) public virtual override(ComponentToken) returns (uint256 assets) {
+        if (_getAggregateTokenStorage().paused) {
+            revert DepositPaused();
+        }
+        return super.mint(shares, receiver, controller);
+    }
+
+    /**
+     * @inheritdoc ERC4626Upgradeable
+     * @dev Overridden to add pause check before minting
+     * @param shares Amount of shares to mint
+     * @param receiver Address that will receive the shares
+     * @return assets Amount of assets deposited
+     */
+    function mint(
+        uint256 shares,
+        address receiver
+    ) public virtual override(ERC4626Upgradeable) returns (uint256 assets) {
+        if (_getAggregateTokenStorage().paused) {
+            revert DepositPaused();
+        }
+        return super.mint(shares, receiver);
+    }
+
     /// @inheritdoc IComponentToken
     function redeem(
         uint256 shares,