make ComponentToken inherit from ReentrancyGuardUpgradeable (#105)

plumenetwork · Dec 12, 2024 · 25c3a6f · 25c3a6f
1 parent 8a65a37
commit 25c3a6f
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 17 deletions.
diff --git a/nest/src/AggregateToken.sol b/nest/src/AggregateToken.sol
@@ -197,7 +197,10 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      * @param componentToken ComponentToken to approve
      * @param amount Amount of `asset` to approve
      */
-    function approveComponentToken(IComponentToken componentToken, uint256 amount) external onlyRole(ADMIN_ROLE) {
+    function approveComponentToken(
+        IComponentToken componentToken,
+        uint256 amount
+    ) external nonReentrant onlyRole(ADMIN_ROLE) {
         IERC20(componentToken.asset()).approve(address(componentToken), amount);
     }
 
@@ -208,7 +211,7 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      */
     function addComponentToken(
         IComponentToken componentToken
-    ) external onlyRole(ADMIN_ROLE) {
+    ) external nonReentrant onlyRole(ADMIN_ROLE) {
         AggregateTokenStorage storage $ = _getAggregateTokenStorage();
         if ($.componentTokenMap[componentToken]) {
             revert ComponentTokenAlreadyListed(componentToken);
@@ -225,7 +228,10 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      * @param componentToken ComponentToken to buy
      * @param assets Amount of `asset` to pay to receive the ComponentToken
      */
-    function buyComponentToken(IComponentToken componentToken, uint256 assets) public onlyRole(ADMIN_ROLE) {
+    function buyComponentToken(
+        IComponentToken componentToken,
+        uint256 assets
+    ) public nonReentrant onlyRole(ADMIN_ROLE) {
         AggregateTokenStorage storage $ = _getAggregateTokenStorage();
 
         if (!$.componentTokenMap[componentToken]) {
@@ -248,7 +254,7 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
     function sellComponentToken(
         IComponentToken componentToken,
         uint256 componentTokenAmount
-    ) public onlyRole(ADMIN_ROLE) {
+    ) public nonReentrant onlyRole(ADMIN_ROLE) {
         uint256 assets = componentToken.redeem(componentTokenAmount, address(this), address(this));
         emit ComponentTokenSold(msg.sender, componentToken, componentTokenAmount, assets);
     }
@@ -260,7 +266,10 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      * @param componentToken ComponentToken to buy
      * @param assets Amount of `asset` to pay to receive the ComponentToken
      */
-    function requestBuyComponentToken(IComponentToken componentToken, uint256 assets) public onlyRole(ADMIN_ROLE) {
+    function requestBuyComponentToken(
+        IComponentToken componentToken,
+        uint256 assets
+    ) public nonReentrant onlyRole(ADMIN_ROLE) {
         uint256 requestId = componentToken.requestDeposit(assets, address(this), address(this));
         emit ComponentTokenBuyRequested(msg.sender, componentToken, assets, requestId);
     }
@@ -275,7 +284,7 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
     function requestSellComponentToken(
         IComponentToken componentToken,
         uint256 componentTokenAmount
-    ) public onlyRole(ADMIN_ROLE) {
+    ) public nonReentrant onlyRole(ADMIN_ROLE) {
         uint256 requestId = componentToken.requestRedeem(componentTokenAmount, address(this), address(this));
         emit ComponentTokenSellRequested(msg.sender, componentToken, componentTokenAmount, requestId);
     }
@@ -289,7 +298,7 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      */
     function setAskPrice(
         uint256 askPrice
-    ) external onlyRole(PRICE_UPDATER_ROLE) {
+    ) external nonReentrant onlyRole(PRICE_UPDATER_ROLE) {
         _getAggregateTokenStorage().askPrice = askPrice;
     }
 
@@ -300,7 +309,7 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      */
     function setBidPrice(
         uint256 bidPrice
-    ) external onlyRole(PRICE_UPDATER_ROLE) {
+    ) external nonReentrant onlyRole(PRICE_UPDATER_ROLE) {
         _getAggregateTokenStorage().bidPrice = bidPrice;
     }
 
@@ -321,7 +330,7 @@ contract AggregateToken is ComponentToken, IAggregateToken, ERC1155Holder {
      * @notice Unpause the AggregateToken contract for deposits
      * @dev Only the owner can unpause the AggregateToken contract for deposits
      */
-    function unpause() external onlyRole(ADMIN_ROLE) {
+    function unpause() external nonReentrant onlyRole(ADMIN_ROLE) {
         AggregateTokenStorage storage $ = _getAggregateTokenStorage();
         if (!$.paused) {
             revert NotPaused();

diff --git a/nest/src/ComponentToken.sol b/nest/src/ComponentToken.sol
@@ -5,6 +5,8 @@ import { AccessControlUpgradeable } from "@openzeppelin/contracts-upgradeable/ac
 import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import { ERC4626Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/ERC4626Upgradeable.sol";
+
+import { ReentrancyGuardUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
 import { IERC4626 } from "@openzeppelin/contracts/interfaces/IERC4626.sol";
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { IERC20Metadata } from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
@@ -26,6 +28,7 @@ abstract contract ComponentToken is
     ERC4626Upgradeable,
     AccessControlUpgradeable,
     UUPSUpgradeable,
+    ReentrancyGuardUpgradeable,
     ERC165,
     IERC7540
 {
@@ -155,6 +158,7 @@ abstract contract ComponentToken is
         __ERC4626_init(asset_);
         __AccessControl_init();
         __UUPSUpgradeable_init();
+        __ReentrancyGuard_init();
 
         _grantRole(DEFAULT_ADMIN_ROLE, owner);
         _grantRole(ADMIN_ROLE, owner);
@@ -236,7 +240,7 @@ abstract contract ComponentToken is
         uint256 assets,
         address controller,
         address owner
-    ) public virtual returns (uint256 requestId) {
+    ) public virtual nonReentrant returns (uint256 requestId) {
         if (assets == 0) {
             revert ZeroAmount();
         }
@@ -264,7 +268,7 @@ abstract contract ComponentToken is
      * @param shares Amount of shares to receive in exchange
      * @param controller Controller of the request
      */
-    function _notifyDeposit(uint256 assets, uint256 shares, address controller) internal virtual {
+    function _notifyDeposit(uint256 assets, uint256 shares, address controller) internal virtual nonReentrant {
         if (assets == 0) {
             revert ZeroAmount();
         }
@@ -285,7 +289,11 @@ abstract contract ComponentToken is
     }
 
     /// @inheritdoc IComponentToken
-    function deposit(uint256 assets, address receiver, address controller) public virtual returns (uint256 shares) {
+    function deposit(
+        uint256 assets,
+        address receiver,
+        address controller
+    ) public virtual nonReentrant returns (uint256 shares) {
         if (assets == 0) {
             revert ZeroAmount();
         }
@@ -314,7 +322,11 @@ abstract contract ComponentToken is
     }
 
     /// @inheritdoc IERC7540
-    function mint(uint256 shares, address receiver, address controller) public virtual returns (uint256 assets) {
+    function mint(
+        uint256 shares,
+        address receiver,
+        address controller
+    ) public virtual nonReentrant returns (uint256 assets) {
         if (shares == 0) {
             revert ZeroAmount();
         }
@@ -347,7 +359,7 @@ abstract contract ComponentToken is
         uint256 shares,
         address controller,
         address owner
-    ) public virtual returns (uint256 requestId) {
+    ) public virtual nonReentrant returns (uint256 requestId) {
         if (shares == 0) {
             revert ZeroAmount();
         }
@@ -373,7 +385,7 @@ abstract contract ComponentToken is
      * @param shares Amount of shares that was redeemed by `requestRedeem`
      * @param controller Controller of the request
      */
-    function _notifyRedeem(uint256 assets, uint256 shares, address controller) internal virtual {
+    function _notifyRedeem(uint256 assets, uint256 shares, address controller) internal virtual nonReentrant {
         if (shares == 0) {
             revert ZeroAmount();
         }
@@ -398,7 +410,7 @@ abstract contract ComponentToken is
         uint256 shares,
         address receiver,
         address controller
-    ) public virtual override(ERC4626Upgradeable, IERC7540) returns (uint256 assets) {
+    ) public virtual override(ERC4626Upgradeable, IERC7540) nonReentrant returns (uint256 assets) {
         if (shares == 0) {
             revert ZeroAmount();
         }
@@ -431,7 +443,7 @@ abstract contract ComponentToken is
         uint256 assets,
         address receiver,
         address controller
-    ) public virtual override(ERC4626Upgradeable, IERC7540) returns (uint256 shares) {
+    ) public virtual override(ERC4626Upgradeable, IERC7540) nonReentrant returns (uint256 shares) {
         if (assets == 0) {
             revert ZeroAmount();
         }