fix: adding reinitializer(2) to init the ReentrancyGuard_init

OriumNetwork · Aug 12, 2024 · 0f115d1 · 0f115d1
1 parent 938e11b
commit 0f115d1
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 36 deletions.
diff --git a/contracts/NftRentalMarketplace.sol b/contracts/NftRentalMarketplace.sol
@@ -93,13 +93,20 @@ contract NftRentalMarketplace is Initializable, OwnableUpgradeable, PausableUpgr
     function initialize(address _owner, address _oriumMarketplaceRoyalties) external initializer {
         __Pausable_init();
         __Ownable_init();
-        __ReentrancyGuard_init();
 
         oriumMarketplaceRoyalties = _oriumMarketplaceRoyalties;
 
         transferOwnership(_owner);
     }
 
+    /**
+     * @notice Initializes the reentrancy guard for the new version.
+     * @dev This function can only be called once, after the contract upgrade.
+     */
+    function initializeV2() external reinitializer(2) {
+        __ReentrancyGuard_init();
+    }
+
     /** ============================ Rental Functions  ================================== **/
 
     /** ######### Setters ########### **/
@@ -164,38 +171,7 @@ contract NftRentalMarketplace is Initializable, OwnableUpgradeable, PausableUpgr
             _expirationDate
         );
 
-        if (_offer.feeTokenAddress == address(0)) {
-            uint256 totalFeeAmount = _offer.feeAmountPerSecond * _duration;
-            require(msg.value >= totalFeeAmount, 'NftRentalMarketplace: Incorrect native token amount');
-
-            uint256 marketplaceFeeAmount = LibNftRentalMarketplace.getAmountFromPercentage(
-                totalFeeAmount,
-                IOriumMarketplaceRoyalties(oriumMarketplaceRoyalties).marketplaceFeeOf(_offer.tokenAddress)
-            );
-            IOriumMarketplaceRoyalties.RoyaltyInfo memory royaltyInfo = IOriumMarketplaceRoyalties(
-                oriumMarketplaceRoyalties
-            ).royaltyInfoOf(_offer.tokenAddress);
-
-            uint256 royaltyAmount = LibNftRentalMarketplace.getAmountFromPercentage(
-                totalFeeAmount,
-                royaltyInfo.royaltyPercentageInWei
-            );
-            uint256 lenderAmount = totalFeeAmount - marketplaceFeeAmount - royaltyAmount;
-
-            payable(owner()).transfer(marketplaceFeeAmount);
-            payable(royaltyInfo.treasury).transfer(royaltyAmount);
-            payable(_offer.lender).transfer(lenderAmount);
-        } else {
-            LibNftRentalMarketplace.transferFees(
-                _offer.feeTokenAddress,
-                owner(),
-                _offer.lender,
-                oriumMarketplaceRoyalties,
-                _offer.tokenAddress,
-                _offer.feeAmountPerSecond,
-                _duration
-            );
-        }
+        _transferFees(_offer.tokenAddress, _offer.feeTokenAddress, _offer.feeAmountPerSecond, _duration, _offer.lender);
 
         LibNftRentalMarketplace.grantRoles(
             oriumMarketplaceRoyalties,
@@ -332,6 +308,54 @@ contract NftRentalMarketplace is Initializable, OwnableUpgradeable, PausableUpgr
         emit RentalOfferCancelled(_offer.lender, _offer.nonce);
     }
 
+    /**
+     * @dev Transfers the fees to the marketplace, the creator and the lender.
+     * @param _feeTokenAddress The address of the ERC20 token for rental fees.
+     * @param _feeAmountPerSecond  The amount of fee per second.
+     * @param _duration The duration of the rental.
+     * @param _lenderAddress The address of the lender.
+     */
+    function _transferFees(
+        address _tokenAddress,
+        address _feeTokenAddress,
+        uint256 _feeAmountPerSecond,
+        uint64 _duration,
+        address _lenderAddress
+    ) internal {
+        if (_feeTokenAddress == address(0)) {
+            uint256 totalFeeAmount = _feeAmountPerSecond * _duration;
+            require(msg.value >= totalFeeAmount, 'NftRentalMarketplace: Incorrect native token amount');
+
+            uint256 marketplaceFeeAmount = LibNftRentalMarketplace.getAmountFromPercentage(
+                totalFeeAmount,
+                IOriumMarketplaceRoyalties(oriumMarketplaceRoyalties).marketplaceFeeOf(_tokenAddress)
+            );
+            IOriumMarketplaceRoyalties.RoyaltyInfo memory royaltyInfo = IOriumMarketplaceRoyalties(
+                oriumMarketplaceRoyalties
+            ).royaltyInfoOf(_tokenAddress);
+
+            uint256 royaltyAmount = LibNftRentalMarketplace.getAmountFromPercentage(
+                totalFeeAmount,
+                royaltyInfo.royaltyPercentageInWei
+            );
+            uint256 lenderAmount = totalFeeAmount - marketplaceFeeAmount - royaltyAmount;
+
+            payable(owner()).transfer(marketplaceFeeAmount);
+            payable(royaltyInfo.treasury).transfer(royaltyAmount);
+            payable(_lenderAddress).transfer(lenderAmount);
+        } else {
+            LibNftRentalMarketplace.transferFees(
+                _feeTokenAddress,
+                owner(),
+                _lenderAddress,
+                oriumMarketplaceRoyalties,
+                _tokenAddress,
+                _feeAmountPerSecond,
+                _duration
+            );
+        }
+    }
+
     /** ============================ Core Functions  ================================== **/
 
     /** ######### Setters ########### **/

diff --git a/contracts/OriumSftMarketplace.sol b/contracts/OriumSftMarketplace.sol
@@ -111,13 +111,20 @@ contract OriumSftMarketplace is Initializable, OwnableUpgradeable, PausableUpgra
     function initialize(address _owner, address _oriumMarketplaceRoyalties) public initializer {
         __Pausable_init();
         __Ownable_init();
-        __ReentrancyGuard_init();
 
         oriumMarketplaceRoyalties = _oriumMarketplaceRoyalties;
 
         transferOwnership(_owner);
     }
 
+    /**
+     * @notice Initializes the reentrancy guard for the new version.
+     * @dev This function can only be called once, after the contract upgrade.
+     */
+    function initializeV2() external reinitializer(2) {
+        __ReentrancyGuard_init();
+    }
+
     /** ============================ Rental Functions  ================================== **/
 
     /** ######### Setters ########### **/
@@ -445,7 +452,7 @@ contract OriumSftMarketplace is Initializable, OwnableUpgradeable, PausableUpgra
         uint256 _lenderAmount = _feeAmount - _royaltyAmount - _marketplaceFeeAmount;
 
         if (_feeTokenAddress == address(0)) {
-            require(msg.value >= _feeAmount, 'OriumSftMarketplace: Insufficient native token amount');
+            require(msg.value == _feeAmount, 'OriumSftMarketplace: Insufficient native token amount');
             payable(owner()).transfer(_marketplaceFeeAmount);
             payable(_royaltyInfo.treasury).transfer(_royaltyAmount);
             payable(_lenderAddress).transfer(_lenderAmount);

diff --git a/test/NftRentalMarketplace.test.ts b/test/NftRentalMarketplace.test.ts
@@ -1049,6 +1049,10 @@ describe('NftRentalMarketplace', () => {
             'Initializable: contract is already initialized',
           )
         })
+        it('Should NOT initialize the contract with initializeV2 if already initialized', async () => {
+          await marketplace.initializeV2()
+          await expect(marketplace.initializeV2()).to.be.revertedWith('Initializable: contract is already initialized')
+        })
       })
       describe('Pausable', async () => {
         describe('Pause', async () => {

diff --git a/test/OriumSftMarketplace.test.ts b/test/OriumSftMarketplace.test.ts
@@ -16,7 +16,6 @@ import {
   OriumSftMarketplace,
   SftRolesRegistrySingleRole,
   SftRolesRegistrySingleRoleLegacy,
-  ReentrancyAttack,
 } from '../typechain-types'
 
 describe('OriumSftMarketplace', () => {
@@ -1609,6 +1608,10 @@ describe('OriumSftMarketplace', () => {
             'Initializable: contract is already initialized',
           )
         })
+        it('Should NOT initialize the contract with initializeV2 if already initialized', async () => {
+          await marketplace.initializeV2()
+          await expect(marketplace.initializeV2()).to.be.revertedWith('Initializable: contract is already initialized')
+        })
       })
 
       describe('Pausable', async () => {