Track prepaid amount and withdraw treasury

contracts/StorageContract.sol

@@ -81,6 +81,9 @@ abstract contract StorageContract is DecentralizedKV {
     /// @notice
     uint256 public prepaidLastMineTime;
 
+    /// @notice Fund tracker for prepaid
+    uint256 public totalPrepaidAmount;
+
     // TODO: Reserve extra slots (to a total of 50?) in the storage layout for future upgrades
 
     /// @notice Emitted when a block is mined.
@@ -143,7 +146,9 @@ abstract contract StorageContract is DecentralizedKV {
     }
 
     /// @notice People can sent ETH to the contract.
-    function sendValue() public payable {}
+    function sendValue() public payable {
+        totalPrepaidAmount += msg.value;
+    }
 
     /// @notice Upfront payment for the next insertion
     function upfrontPayment() public view virtual override returns (uint256) {
@@ -230,17 +235,19 @@ abstract contract StorageContract is DecentralizedKV {
     function _rewardMiner(uint256 _shardId, address _miner, uint256 _minedTs, uint256 _diff) internal {
         // Mining is successful.
         // Send reward to coinbase and miner.
-        (bool updatePrepaidTime, uint256 treasuryReward, uint256 minerReward) = _miningReward(_shardId, _minedTs);
+        (bool updatePrepaidTime, uint256 prepaidAmountSaved, uint256 treasuryReward, uint256 minerReward) =
+            _miningReward(_shardId, _minedTs);
         if (updatePrepaidTime) {
             prepaidLastMineTime = _minedTs;
         }
-
+        if (prepaidAmountSaved > 0) {
+            totalPrepaidAmount += prepaidAmountSaved;
+        }
+        totalPrepaidAmount += treasuryReward;
         // Update mining info.
         MiningLib.update(infos[_shardId], _minedTs, _diff);
 
-        require(treasuryReward + minerReward <= address(this).balance, "StorageContract: not enough balance");
-        // TODO: avoid reentrancy attack
-        payable(treasury).transfer(treasuryReward);
+        require(minerReward <= address(this).balance, "StorageContract: not enough balance");
         payable(_miner).transfer(minerReward);
         emit MinedBlock(_shardId, _diff, infos[_shardId].blockMined, _minedTs, _miner, minerReward);
     }
@@ -249,13 +256,19 @@ abstract contract StorageContract is DecentralizedKV {
     /// @param _shardId  The shard id.
     /// @param _minedTs  The mined timestamp.
     /// @return updatePrepaidTime Whether to update the prepaid time.
+    /// @return prepaidAmountSaved The capped part of prepaid amount.
     /// @return treasuryReward    The treasury reward.
     /// @return minerReward       The miner reward.
-    function _miningReward(uint256 _shardId, uint256 _minedTs) internal view returns (bool, uint256, uint256) {
+    function _miningReward(uint256 _shardId, uint256 _minedTs)
+        internal
+        view
+        returns (bool, uint256, uint256, uint256)
+    {
         MiningLib.MiningInfo storage info = infos[_shardId];
         uint256 lastShardIdx = kvEntryCount > 0 ? (kvEntryCount - 1) >> SHARD_ENTRY_BITS : 0;
-        uint256 reward = 0;
         bool updatePrepaidTime = false;
+        uint256 prepaidAmountSaved = 0;
+        uint256 reward = 0;
         if (_shardId < lastShardIdx) {
             reward = _paymentIn(STORAGE_COST << SHARD_ENTRY_BITS, info.lastMineTime, _minedTs);
         } else if (_shardId == lastShardIdx) {
@@ -265,6 +278,8 @@ abstract contract StorageContract is DecentralizedKV {
                 uint256 prepaidAmountCap =
                     STORAGE_COST * ((1 << SHARD_ENTRY_BITS) - kvEntryCount % (1 << SHARD_ENTRY_BITS));
                 if (prepaidAmountCap > prepaidAmount) {
+                    prepaidAmountSaved = _paymentIn(prepaidAmountCap, prepaidLastMineTime, _minedTs)
+                        - _paymentIn(prepaidAmount, prepaidLastMineTime, _minedTs);
                     prepaidAmountCap = prepaidAmount;
                 }
                 reward += _paymentIn(prepaidAmountCap, prepaidLastMineTime, _minedTs);
@@ -274,7 +289,7 @@ abstract contract StorageContract is DecentralizedKV {
 
         uint256 treasuryReward = (reward * TREASURY_SHARE) / 10000;
         uint256 minerReward = reward - treasuryReward;
-        return (updatePrepaidTime, treasuryReward, minerReward);
+        return (updatePrepaidTime, prepaidAmountSaved, treasuryReward, minerReward);
     }
 
     /// @notice Get the mining reward.
@@ -283,7 +298,7 @@ abstract contract StorageContract is DecentralizedKV {
     /// @return The mining reward.
     function miningReward(uint256 _shardId, uint256 _blockNum) public view returns (uint256) {
         uint256 minedTs = _getMinedTs(_blockNum);
-        (,, uint256 minerReward) = _miningReward(_shardId, minedTs);
+        (,,, uint256 minerReward) = _miningReward(_shardId, minedTs);
         return minerReward;
     }
 
@@ -374,6 +389,14 @@ abstract contract StorageContract is DecentralizedKV {
         _rewardMiner(_shardId, _miner, mineTs, diff);
     }
 
+    /// @notice Withdraw treasury fund
+    function withdraw(uint256 _amount) public {
+        require(totalPrepaidAmount >= prepaidAmount + _amount, "StorageContract: not enough prepaid amount");
+        totalPrepaidAmount -= _amount;
+        require(address(this).balance >= _amount, "StorageContract: not enough balance");
+        payable(treasury).transfer(_amount);
+    }
+
     /// @notice Get the current block number
     function _blockNumber() internal view virtual returns (uint256) {
         return block.number;

contracts/test/StorageContractTest.t.sol

@@ -22,22 +22,82 @@ contract StorageContractTest is Test {
 
     function testMiningReward() public {
         // no key-value stored on EthStorage, only use prepaid amount as the reward
-        (,, uint256 reward) = storageContract.miningRewards(0, 1);
+        (,,, uint256 reward) = storageContract.miningRewards(0, 1);
         assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT, 0, 1));
 
         // 1 key-value stored on EthStorage
         storageContract.setKvEntryCount(1);
-        (,, reward) = storageContract.miningRewards(0, 1);
+        (,,, reward) = storageContract.miningRewards(0, 1);
         assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT + STORAGE_COST * 1, 0, 1));
 
         // 2 key-value stored on EthStorage
         storageContract.setKvEntryCount(2);
-        (,, reward) = storageContract.miningRewards(0, 1);
+        (,,, reward) = storageContract.miningRewards(0, 1);
         assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT + STORAGE_COST * 2, 0, 1));
 
         // 3 key-value stored on EthStorage, but the reward is capped with 4 * STORAGE_COST
         storageContract.setKvEntryCount(3);
-        (,, reward) = storageContract.miningRewards(0, 1);
+        (,,, reward) = storageContract.miningRewards(0, 1);
         assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT + STORAGE_COST * 2, 0, 1));
     }
+
+    function setUp1() public {
+        storageContract = new TestStorageContract(
+            StorageContract.Config(MAX_KV_SIZE, SHARD_SIZE_BITS, 2, 0, 0, 100), 0, STORAGE_COST, 0
+        );
+        storageContract.initialize(0, PREPAID_AMOUNT, 0, vm.addr(1), address(0x1));
+    }
+
+    function testWithdraw() public {
+        setUp1();
+
+        uint256 valueToSent = 3000;
+        uint256 withdrawAmount = 800;
+
+        storageContract.sendValue{value: valueToSent}();
+        assertEq(storageContract.totalPrepaidAmount(), valueToSent);
+
+        storageContract.withdraw(withdrawAmount);
+        assertEq(storageContract.totalPrepaidAmount(), valueToSent - withdrawAmount);
+        assertEq(storageContract.treasury().balance, withdrawAmount);
+    }
+
+    function testWithdrawRewardMiner() public {
+        setUp1();
+
+        uint256 valueToSent = 5000;
+        uint256 withdrawAmount = 800;
+        uint256 mineTs = 10000;
+        uint40 kvEntryCount = 1;
+        uint256 shardEntry = 1 << (SHARD_SIZE_BITS - MAX_KV_SIZE);
+
+        storageContract.sendValue{value: valueToSent}();
+        storageContract.setKvEntryCount(kvEntryCount);
+        uint256 reward = storageContract.paymentIn(STORAGE_COST * kvEntryCount, 0, mineTs);
+        uint256 prepaidReward = storageContract.paymentIn(PREPAID_AMOUNT, 0, mineTs);
+        reward += prepaidReward;
+        uint256 treasureReward = (reward * storageContract.treasuryShare()) / 10000;
+        uint256 minerReward = reward - treasureReward;
+        uint256 prepaidAmountCap = STORAGE_COST * (shardEntry - kvEntryCount);
+        uint256 prepaidAmountSaved = storageContract.paymentIn(prepaidAmountCap, 0, mineTs) - prepaidReward;
+
+        storageContract.rewardMiner(0, vm.addr(2), mineTs, 1);
+        uint256 totalPrepaid = valueToSent + treasureReward + prepaidAmountSaved;
+        assertEq(storageContract.totalPrepaidAmount(), totalPrepaid);
+
+        storageContract.withdraw(withdrawAmount);
+        assertEq(storageContract.totalPrepaidAmount(), totalPrepaid - withdrawAmount);
+        assertEq(storageContract.treasury().balance, withdrawAmount);
+        assertEq(address(storageContract).balance, valueToSent - minerReward - withdrawAmount);
+    }
+
+    function testWithdrawInsufficientFunds() public {
+        uint256 valueToSent = 3000;
+        uint256 withdrawAmount = 1500;
+
+        storageContract.sendValue{value: valueToSent}();
+
+        vm.expectRevert("StorageContract: not enough prepaid amount");
+        storageContract.withdraw(withdrawAmount);
+    }
 }

contracts/test/TestStorageContract.sol

@@ -38,7 +38,11 @@ contract TestStorageContract is StorageContract {
         return _paymentIn(_x, _fromTs, _toTs);
     }
 
-    function miningRewards(uint256 _shardId, uint256 _minedTs) public view returns (bool, uint256, uint256) {
+    function miningRewards(uint256 _shardId, uint256 _minedTs) public view returns (bool, uint256, uint256, uint256) {
         return _miningReward(_shardId, _minedTs);
     }
+
+    function rewardMiner(uint256 _shardId, address _miner, uint256 _minedTs, uint256 _diff) public {
+        return _rewardMiner(_shardId, _miner, _minedTs, _diff);
+    }
 }