fix compute payment on putBlobs

ethstorage · Aug 19, 2024 · 6437963 · 6437963
1 parent 57854d8
commit 6437963
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/contracts/StorageContract.sol b/contracts/StorageContract.sol
@@ -150,15 +150,21 @@ abstract contract StorageContract is DecentralizedKV {
 
     /// @notice Checks the payment using the last mine time.
     function _prepareAppendWithTimestamp(uint256 _timestamp, uint256 _batchSize) internal {
-        uint256 totalPayment = upfrontPaymentInBatch(_batchSize);
-        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
-
-        uint256 shardIdNew = (kvEntryCount + _batchSize) >> SHARD_ENTRY_BITS; // shard id after the batch
-        if (shardIdNew > (kvEntryCount >> SHARD_ENTRY_BITS)) {
+        uint256 shardId = (kvEntryCount - _batchSize) >> SHARD_ENTRY_BITS;
+        uint256 shardIdNew = kvEntryCount >> SHARD_ENTRY_BITS; // shard id after the batch
+        uint256 totalPayment = 0;
+        if (shardIdNew > shardId) {
+            uint256 kvCountNew = kvEntryCount % (1 << SHARD_ENTRY_BITS);
+            totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);
             // Open a new shard and mark the shard is ready to mine.
             // (TODO): Setup shard difficulty as current difficulty / factor?
             infos[shardIdNew].lastMineTime = _timestamp;
+        } else {
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
         }
+
+        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
     }
 
     /// @notice Upfront payment for the next insertion
@@ -183,7 +189,8 @@ abstract contract StorageContract is DecentralizedKV {
         uint256 totalPayment = 0;
         if (shardIdNew > shardId) {
             uint256 kvCountNew = totalEntries % (1 << SHARD_ENTRY_BITS);
-            totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);
+            totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);
         } else {
             totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
         }

diff --git a/contracts/test/EthStorageContractTest.t.sol b/contracts/test/EthStorageContractTest.t.sol
@@ -110,13 +110,13 @@ contract EthStorageContractTest is Test {
             blobIdxs[i] = i;
             lengths[i] = 10 + i*10;
         }
-        vm.warp(vm.unixTime());
-
-        // uint256 insufficientCost = storageContract.upfrontPaymentInBatch(size - 1);
-        // // Expect the specific revert reason from _prepareBatchAppend due to insufficient msg.value
-        // vm.expectRevert("StorageContract: not enough batch payment");
-        // storageContract.putBlobs{value: insufficientCost}(keys, blobIdxs, lengths);
+        vm.warp(0);
 
+        uint256 insufficientCost = storageContract.upfrontPaymentInBatch(size) - 1;
+        // Expect the specific revert reason from _prepareBatchAppend due to insufficient msg.value
+        vm.expectRevert("StorageContract: not enough batch payment");
+        storageContract.putBlobs{value: insufficientCost}(keys, blobIdxs, lengths);
+
         // Enough storage cost
         uint256 sufficientCost = storageContract.upfrontPaymentInBatch(size);
         storageContract.putBlobs{value: sufficientCost}(keys, blobIdxs, lengths);