From 786d655ce24d8f259014a121dc940a2472af4bab Mon Sep 17 00:00:00 2001
From: Daniel Lima <daniel.lima@orium.network>
Date: Tue, 23 Apr 2024 12:43:19 -0300
Subject: [PATCH] ON-814: PR fixes

---
 contracts/NftRentalMarketplace.sol | 32 ++++++++++++++++++++++++++++++
 test/NftRentalMarketplace.test.ts  | 20 ++++++++++++++++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/contracts/NftRentalMarketplace.sol b/contracts/NftRentalMarketplace.sol
index 0d595b8..ebb175c 100644
--- a/contracts/NftRentalMarketplace.sol
+++ b/contracts/NftRentalMarketplace.sol
@@ -4,6 +4,8 @@ pragma solidity 0.8.9;
 
 import { IERC721 } from '@openzeppelin/contracts/token/ERC721/IERC721.sol';
 import { IERC20 } from '@openzeppelin/contracts/token/ERC20/IERC20.sol';
+import { IERC7432VaultExtension } from './interfaces/IERC7432VaultExtension.sol';
+import { ERC165Checker } from '@openzeppelin/contracts/utils/introspection/ERC165Checker.sol';
 import { IOriumMarketplaceRoyalties } from './interfaces/IOriumMarketplaceRoyalties.sol';
 import { OwnableUpgradeable } from '@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol';
 import { Initializable } from '@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol';
@@ -16,6 +18,8 @@ import { LibNftRentalMarketplace, RentalOffer, Rental } from './libraries/LibNft
  * @author Orium Network Team - developers@orium.network
  */
 contract NftRentalMarketplace is Initializable, OwnableUpgradeable, PausableUpgradeable {
+    using ERC165Checker for address;
+
     /** ######### Global Variables ########### **/
 
     /// @dev oriumMarketplaceRoyalties stores the collection royalties and fees
@@ -190,6 +194,34 @@ contract NftRentalMarketplace is Initializable, OwnableUpgradeable, PausableUpgr
      */
 
     function cancelRentalOffer(RentalOffer calldata _offer) external whenNotPaused {
+        _cancelRentalOffer(_offer);
+    }
+
+    /**
+     * @notice Cancels a rental offer and withdraws the NFT.
+     * @dev Can only be called by the lender, and only withdraws the NFT if the rental has expired.
+     * @param _offer The rental offer struct. It should be the same as the one used to create the offer.
+     */
+
+    function cancelRentalOfferAndWithdraw(RentalOffer calldata _offer) external whenNotPaused {
+        _cancelRentalOffer(_offer);
+
+        address _rolesRegistry = IOriumMarketplaceRoyalties(oriumMarketplaceRoyalties).nftRolesRegistryOf(
+            _offer.tokenAddress
+        );
+        require(
+            _rolesRegistry.supportsERC165InterfaceUnchecked(type(IERC7432VaultExtension).interfaceId),
+            'NftRentalMarketplace: roles registry does not support IERC7432VaultExtension'
+        );
+        IERC7432VaultExtension(_rolesRegistry).withdraw(_offer.tokenAddress, _offer.tokenId);
+    }
+
+    /**
+     * @notice Cancels a rental offer.
+     * @dev Internal function to cancel a rental offer.
+     * @param _offer The rental offer struct. It should be the same as the one used to create the offer.
+     */
+    function _cancelRentalOffer(RentalOffer calldata _offer) internal {
         bytes32 _offerHash = LibNftRentalMarketplace.hashRentalOffer(_offer);
         LibNftRentalMarketplace.validateCancelRentalOfferParams(
             isCreated[_offerHash],
diff --git a/test/NftRentalMarketplace.test.ts b/test/NftRentalMarketplace.test.ts
index 94850d1..1dc8079 100644
--- a/test/NftRentalMarketplace.test.ts
+++ b/test/NftRentalMarketplace.test.ts
@@ -482,7 +482,7 @@ describe('NftRentalMarketplace', () => {
           })
 
           describe('Cancel Rental Offer', async () => {
-            it('Should cancel a rental offer and releaseTokens from rolesRegistry', async () => {
+            it('Should cancel a rental offer', async () => {
               await expect(marketplace.connect(lender).cancelRentalOffer(rentalOffer))
                 .to.emit(marketplace, 'RentalOfferCancelled')
                 .withArgs(rentalOffer.lender, rentalOffer.nonce)
@@ -568,6 +568,24 @@ describe('NftRentalMarketplace', () => {
                 )
               })
             })
+
+            describe('Cancel Rental Offer', async function () {
+              it('Should cancel a rental offer and withdraw from rolesRegistry, if rental is not active', async () => {
+                await expect(marketplace.connect(lender).cancelRentalOfferAndWithdraw(rentalOffer))
+                  .to.emit(marketplace, 'RentalOfferCancelled')
+                  .withArgs(rentalOffer.lender, rentalOffer.nonce)
+                  .to.emit(rolesRegistry, 'Withdraw')
+                  .withArgs(rentalOffer.lender, rentalOffer.tokenAddress, rentalOffer.tokenId)
+              })
+              it('Should NOT cancel a rental offer and withdraw from rolesRegistry, if rolesRegistry does not support IERC7432VaultExtension', async () => {
+                await marketplaceRoyalties
+                  .connect(operator)
+                  .setRolesRegistry(await mockERC721.getAddress(), AddressZero)
+                await expect(marketplace.connect(lender).cancelRentalOfferAndWithdraw(rentalOffer)).to.be.revertedWith(
+                  'NftRentalMarketplace: roles registry does not support IERC7432VaultExtension',
+                )
+              })
+            })
           })
         })
       })