diff --git a/src/Counter.sol b/src/Counter.sol.old similarity index 100% rename from src/Counter.sol rename to src/Counter.sol.old diff --git a/src/MyToken.sol b/src/MyToken.sol new file mode 100644 index 0000000..6929423 --- /dev/null +++ b/src/MyToken.sol @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +contract MyToken { + string public name = "MyToken"; + string public symbol = "MTK"; + uint8 public decimals = 18; + uint256 public totalSupply; + uint256 public tokenPrice = 0.01 ether; // Price per token in ETH + uint256 public sellFee = 1; // 1% fee on selling + + mapping(address => uint256) public balanceOf; + mapping(address => mapping(address => uint256)) public allowance; + + address public owner; + + event Transfer(address indexed from, address indexed to, uint256 value); + event Approval(address indexed owner, address indexed spender, uint256 value); + event Bought(address indexed buyer, uint256 amount, uint256 cost); + event Sold(address indexed seller, uint256 amount, uint256 revenue); + event PriceUpdated(uint256 newPrice); + event Mint(address indexed to, uint256 amount); + event Burn(address indexed from, uint256 amount); + + modifier onlyOwner() { + require(msg.sender == owner, "Only the owner can perform this action"); + _; + } + + constructor(uint256 _initialSupply) { + owner = msg.sender; + totalSupply = _initialSupply * (10 ** uint256(decimals)); + balanceOf[msg.sender] = totalSupply; + } + + function transfer(address _to, uint256 _value) public returns (bool success) { + _transfer(msg.sender, _to, _value); + return true; + } + + function approve(address _spender, uint256 _value) public returns (bool success) { + allowance[msg.sender][_spender] = _value; + emit Approval(msg.sender, _spender, _value); + return true; + } + + function transferFrom(address _from, address _to, uint256 _value) public returns (bool success) { + require(allowance[_from][msg.sender] >= _value, "Allowance exceeded"); + allowance[_from][msg.sender] -= _value; + _transfer(_from, _to, _value); + return true; + } + + function _transfer(address _from, address _to, uint256 _value) internal { + require(balanceOf[_from] >= _value, "Insufficient balance"); + balanceOf[_from] -= _value; + balanceOf[_to] += _value; + emit Transfer(_from, _to, _value); + } + + function mint(address _to, uint256 _amount) internal { + balanceOf[_to] += _amount; + totalSupply += _amount; + emit Mint(_to, _amount); + emit Transfer(address(0), _to, _amount); + } + + function burn(address _from, uint256 _amount) internal { + require(balanceOf[_from] >= _amount, "Insufficient balance to burn"); + balanceOf[_from] -= _amount; + totalSupply -= _amount; + emit Burn(_from, _amount); + emit Transfer(_from, address(0), _amount); + } + + function buyTokens() public payable returns (bool success) { + uint256 amountToBuy = msg.value / tokenPrice; + require(amountToBuy > 0, "You need to send some ETH"); + mint(msg.sender, amountToBuy); + emit Bought(msg.sender, amountToBuy, msg.value); + return true; + } + + function sellTokens(uint256 amountToSell) public returns (bool success) { + require(balanceOf[msg.sender] >= amountToSell, "Insufficient token balance"); + uint256 ethAmount = amountToSell * tokenPrice; + uint256 fee = ethAmount * sellFee / 100; + uint256 revenue = ethAmount - fee; + require(address(this).balance >= revenue, "Contract has insufficient ETH balance"); + burn(msg.sender, amountToSell); + payable(msg.sender).transfer(revenue); + emit Sold(msg.sender, amountToSell, revenue); + return true; + } +function sellTokens_reentrancy(uint256 amountToSell) public returns (bool success) { + require(balanceOf[msg.sender] >= amountToSell, "Insufficient token balance"); + uint256 ethAmount = amountToSell * tokenPrice; + uint256 fee = ethAmount * sellFee / 100; + uint256 revenue = ethAmount - fee; + require(address(this).balance >= revenue, "Contract has insufficient ETH balance"); + + // Vulnerable point: External call before state update + (bool sent, ) = msg.sender.call{value: revenue}(""); + require(sent, "Failed to send Ether"); + + // State update happens after the external call + burn(msg.sender, amountToSell); + + emit Sold(msg.sender, amountToSell, revenue); + return true; +} + + + function setPrice(uint256 newPrice) public onlyOwner { + tokenPrice = newPrice; + emit PriceUpdated(newPrice); + } + + function withdraw() public onlyOwner { + payable(owner).transfer(address(this).balance); + } + + receive() external payable { + buyTokens(); + } + +} diff --git a/test/Counter.t.sol b/test/Counter.t.sol.old similarity index 100% rename from test/Counter.t.sol rename to test/Counter.t.sol.old diff --git a/test/MyToken.t.sol b/test/MyToken.t.sol new file mode 100644 index 0000000..c4e1daf --- /dev/null +++ b/test/MyToken.t.sol @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Test, console} from "forge-std/Test.sol"; +import {MyToken} from "../src/MyToken.sol"; + +contract MyTokenTest is Test { + MyToken public myToken; + address public owner; + address public user; + + function setUp() public { + owner = address(0x100); + user = address(0x101); + + vm.deal(owner, 100 ether); + vm.deal(user, 100 ether); + + vm.startPrank(owner); + myToken = new MyToken(1000); + vm.stopPrank(); + } + + function test_buyAndTransferTokens() public { + vm.startPrank(user); // User operations under here + + uint ethAmount = 0.01 ether; + // uint ethAmount = 0.001 ether; // This will cause myToken.buyTokens() to revert + + uint userBalanceBeforeBuy = myToken.balanceOf(user); + uint totalSupplyBeforeBuy = myToken.totalSupply(); + console.log("User balance before buy:", userBalanceBeforeBuy); + console.log("Total supply before buy:", totalSupplyBeforeBuy); + + uint expectedBalance = userBalanceBeforeBuy + ethAmount / myToken.tokenPrice(); + + myToken.buyTokens{value: ethAmount}(); + + uint userBalanceAfterBuy = myToken.balanceOf(user); + uint totalSupplyAfterBuy = myToken.totalSupply(); + console.log("User balance after buy:", userBalanceAfterBuy); + console.log("Total supply after buy:", totalSupplyAfterBuy); + + assert(userBalanceAfterBuy > userBalanceBeforeBuy); + assertEq(userBalanceAfterBuy, expectedBalance); + assertEq(totalSupplyAfterBuy, totalSupplyBeforeBuy + ethAmount / myToken.tokenPrice()); + + console.log("===================="); + + uint ownerBalanceBeforeTransfer = myToken.balanceOf(owner); + uint userBalanceBeforeTransfer = myToken.balanceOf(user); + console.log("Owner balance before transfer:", ownerBalanceBeforeTransfer); + console.log("User balance before transfer:", userBalanceBeforeTransfer); + + myToken.transfer(owner, userBalanceBeforeTransfer); + + uint ownerBalanceAfterTransfer = myToken.balanceOf(owner); + uint userBalanceAfterTransfer = myToken.balanceOf(user); + console.log("Owner balance after transfer:", ownerBalanceAfterTransfer); + console.log("User balance after transfer:", userBalanceAfterTransfer); + + assert(ownerBalanceAfterTransfer > ownerBalanceBeforeTransfer); + assertEq(ownerBalanceAfterTransfer, ownerBalanceBeforeTransfer + userBalanceBeforeTransfer); + + vm.stopPrank(); + } +}