new version&&fix potential transfer callback vulnerability

contracts/ETHPool.sol

@@ -0,0 +1,87 @@
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/utils/Address.sol";
+import "@openzeppelin/contracts/utils/Context.sol";
+import "@openzeppelin/contracts/utils/math/SafeMath.sol";
+import "@openzeppelin/contracts/access/Ownable.sol";
+import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
+
+
+contract ETHPool is Context,Ownable,ReentrancyGuard {
+	using SafeMath for uint256;
+	using Address for address payable;
+
+	event Deposit(address user,uint256 balance);
+	event Withdraw(address user,uint256 balance);
+
+	struct Account {
+		uint256 balance;
+		uint256 shares;
+		uint64 lastBalanceChange;
+	}
+
+	mapping (address=>Account) public depositAccount;
+	uint256 public totalBalance=0;
+
+	struct Reward {
+		uint256 settleBalance;
+		uint256 reward;
+	}
+
+	Reward[] public rewardHistory;
+	uint256 public rewardPool=0;
+
+	//update must before user's balance change.
+	function updateShares(Account storage account) private {
+		for (uint256 index = account.lastBalanceChange; index < rewardHistory.length; index++) {
+			//add shares for each time.
+			Reward storage history=rewardHistory[index];
+			account.shares=account.shares.add(history.reward.mul(account.balance).div(history.settleBalance));
+		}
+		account.lastBalanceChange=uint64(rewardHistory.length);
+	}
+
+	function deposit() public payable {
+		address user=address(_msgSender());
+		uint256 amount=msg.value;
+		require(amount>0);
+
+		Account storage account=depositAccount[user];
+		updateShares(account);
+		account.balance=account.balance.add(amount);
+		totalBalance=totalBalance.add(amount);
+		emit Deposit(user, account.balance.add(account.shares));
+	}
+
+	function withdraw(uint256 amount) public nonReentrant {
+		address payable user=payable(_msgSender());
+		Account storage account=depositAccount[user];
+		updateShares(account);
+		require(account.balance.add(account.shares)>=amount);
+
+		//withdraw from balance first
+		uint256 amountFromBalance=amount;
+		if(amount>account.balance) {
+			amountFromBalance=account.balance;
+			uint256 amountFromShares=amount.sub(account.balance);
+			account.shares=account.shares.sub(amountFromShares);
+		}
+		account.balance=account.balance.sub(amountFromBalance);
+		totalBalance=totalBalance.sub(amountFromBalance);
+		//avoid from potential vulnerability
+		user.sendValue(amount);
+		emit Withdraw(user, account.balance.add(account.shares));
+	}
+
+	function depositReward() public payable onlyOwner {
+		uint256 amount=msg.value;
+		rewardPool=rewardPool.add(amount);
+		rewardHistory.push(Reward({settleBalance:totalBalance,reward:amount}));
+	}
+
+	function balanceOf(address user) public returns(uint256) {
+		Account storage account=depositAccount[user];
+		updateShares(account);
+		return account.balance.add(account.shares);
+	}
+}

scripts/get_balance.js

@@ -0,0 +1,8 @@
+const Web3 = require("web3");
+
+(async ()=>
+{
+	const contract="0xFe834810F61ABf397B53eC21bCc2652f91d1F7fD";
+	let web3=new Web3(`https://ropsten.infura.io/v3/e61d7489169641adb4ba7a698ad81d51`);
+	console.log("pool has "+web3.utils.fromWei(await web3.eth.getBalance(contract))+" ethers");
+})();

test/ethpool.js

@@ -0,0 +1,88 @@
+const { assert } = require("chai");
+const Web3 = require("web3");
+const expect = require('chai').expect;
+const ETHPool = artifacts.require("ETHPool");
+
+
+contract("ethpool test", async (accounts) => {
+	let [owner,alice,bob,charlie] = accounts;
+	let [reward_amount,alice_amount,bob_amount,charlie_amount]=[1,5,10,2];
+	let instanceOfPool;
+	let web3;
+	let weis=function(num) {
+		return web3.utils.toWei(String(num), "ether");
+	}
+    before(async () => {
+		instanceOfPool = await ETHPool.deployed();
+		console.log("pool:"+instanceOfPool.address);
+		console.log("alice:"+alice+"\nbob:"+bob);
+		web3 = new Web3(Web3.givenProvider || 'ws://127.0.0.1:7545');
+    });
+
+	it("check contract and alice balance after alice deposit", async () => {
+		let result = await instanceOfPool.deposit({from:alice,value:weis(alice_amount)});
+		expect(result.receipt.status).to.equal(true);
+		expect(await web3.eth.getBalance(instanceOfPool.address)).to.equal(weis(alice_amount));
+		result=String(await instanceOfPool.balanceOf.call(alice));
+		expect(result).to.equal(weis(alice_amount));
+	});
+	it("check contract and bob balance after bob deposit", async () => {
+		let result = await instanceOfPool.deposit({from:bob,value:weis(bob_amount)});
+		expect(result.receipt.status).to.equal(true);
+		expect(await web3.eth.getBalance(instanceOfPool.address)).to.equal(weis(alice_amount+bob_amount));
+		result=String(await instanceOfPool.balanceOf.call(bob));
+		expect(result).to.equal(weis(bob_amount));
+	});
+	it("check contract and users balance after reward deposit", async () => {
+		let result = await instanceOfPool.depositReward({from:owner,value:weis(reward_amount)});
+		expect(result.receipt.status).to.equal(true);
+		expect(await web3.eth.getBalance(instanceOfPool.address)).to.equal(weis(reward_amount+alice_amount+bob_amount));
+		//range based check
+		result=Number(await instanceOfPool.balanceOf.call(alice));
+		expect(result).to.greaterThan(Number(weis(alice_amount+0.99*reward_amount/3)));
+		expect(result).to.lessThan(Number(weis(alice_amount+1.01*reward_amount/3)));
+		result=Number(await instanceOfPool.balanceOf.call(bob));
+		expect(result).to.greaterThan(Number(weis(bob_amount+1.99*reward_amount/3)));
+		expect(result).to.lessThan(Number(weis(bob_amount+2.01*reward_amount/3)));
+	});
+	it("check contract and charlie balance after charlie deposit", async () => {
+		let result = await instanceOfPool.deposit({from:charlie,value:weis(charlie_amount)});
+		expect(result.receipt.status).to.equal(true);
+		expect(await web3.eth.getBalance(instanceOfPool.address)).to.equal(weis(reward_amount+alice_amount+bob_amount+charlie_amount));
+		result=String(await instanceOfPool.balanceOf.call(charlie));
+		expect(result).to.equal(weis(charlie_amount));
+	});
+	it("check contract and alice balance after alice withdraw", async () => {
+		let contractBalance=Number(await web3.eth.getBalance(instanceOfPool.address));
+		let balance=Number(await instanceOfPool.balanceOf.call(alice));
+		let result = await instanceOfPool.withdraw(String(balance),{from:alice});
+		expect(result.receipt.status).to.equal(true);
+		//Don't know where 2000 weis went
+		result=Number(await web3.eth.getBalance(instanceOfPool.address));
+		expect(result).to.greaterThan((contractBalance-balance)*0.99);
+		expect(result).to.lessThan((contractBalance-balance)*1.01);
+		//neither
+		result=Number(await instanceOfPool.balanceOf.call(alice));
+		expect(result).to.greaterThan(0);
+		expect(result).to.lessThan(2000);
+	});
+
+	it("check insufficient deposit", async () => {
+		try {
+			let result = await instanceOfPool.deposit({from:alice,value:weis(1000)});
+			assert(false);
+		} catch {}
+	});
+	it("check non-owner depositReward", async () => {
+		try {
+			let result = await instanceOfPool.depositReward({from:alice,value:weis(alice_amount)});
+			assert(false);
+		} catch {}
+	});
+	it("check insufficient withdraw", async () => {
+		try {
+			let result = await instanceOfPool.withdraw(weis(100),{from:bob});
+			assert(false);
+		} catch {}
+	});
+});