add contract and tests
This commit is contained in:
parent
e93c870a87
commit
f538e88d49
|
|
@ -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();
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
@ -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();
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue