batch payment

ethstorage · Aug 16, 2024 · 57854d8 · 57854d8
1 parent ffbcb04
commit 57854d8
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 18 deletions.
diff --git a/contracts/DecentralizedKV.sol b/contracts/DecentralizedKV.sol
@@ -103,9 +103,10 @@ contract DecentralizedKV is OwnableUpgradeable {
     function _paymentIn(uint256 _x, uint256 _fromTs, uint256 _toTs) internal view returns (uint256) {
         return _paymentInInterval(_x, _fromTs - START_TIME, _toTs - START_TIME);
     }
-
+    
     /// @notice Evaluate payment given the timestamp.
     function _upfrontPayment(uint256 _ts) internal view returns (uint256) {
+        require(_ts >= START_TIME, "DecentralizedKV: invalid timestamp to evaluate payment");
         return _paymentInf(STORAGE_COST, _ts - START_TIME);
     }
 
@@ -114,6 +115,11 @@ contract DecentralizedKV is OwnableUpgradeable {
         return _upfrontPayment(block.timestamp);
     }
 
+    /// @notice Upfront payment for the next batch insertion
+    function upfrontPaymentInBatch(uint256 _batchSize) public view virtual returns (uint256) {
+        return upfrontPayment() * _batchSize;
+    }
+
     /// @notice Checks before appending the key-value.
     function _prepareAppend(uint256 _batchSize) internal virtual {
         require(msg.value >= upfrontPayment() * _batchSize, "DecentralizedKV: not enough batch payment");

diff --git a/contracts/StorageContract.sol b/contracts/StorageContract.sol
@@ -150,26 +150,15 @@ abstract contract StorageContract is DecentralizedKV {
 
     /// @notice Checks the payment using the last mine time.
     function _prepareAppendWithTimestamp(uint256 _timestamp, uint256 _batchSize) internal {
-        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS;
-        uint256 totalEntries = kvEntryCount + _batchSize; // include the batch to be put
-        uint256 shardIdNew = totalEntries >> SHARD_ENTRY_BITS; // shard id after the batch
-        uint256 totalPayment = 0;
+        uint256 totalPayment = upfrontPaymentInBatch(_batchSize);
+        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
 
-        if (shardIdNew > shardId) {
-            uint256 kvCountNew = totalEntries % (1 << SHARD_ENTRY_BITS);
-            totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;
-            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);
+        uint256 shardIdNew = (kvEntryCount + _batchSize) >> SHARD_ENTRY_BITS; // shard id after the batch
+        if (shardIdNew > (kvEntryCount >> SHARD_ENTRY_BITS)) {
             // Open a new shard and mark the shard is ready to mine.
             // (TODO): Setup shard difficulty as current difficulty / factor?
-            if (shardIdNew != 0) {
-                // shard0 is already opened in constructor
-                infos[shardIdNew].lastMineTime = _timestamp;
-            }
-        } else {
-            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
+            infos[shardIdNew].lastMineTime = _timestamp;
         }
-
-        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
     }
 
     /// @notice Upfront payment for the next insertion
@@ -186,6 +175,21 @@ abstract contract StorageContract is DecentralizedKV {
         }
     }
 
+    /// @notice Upfront payment for a batch insertion
+    function upfrontPaymentInBatch(uint256 _batchSize) public view virtual override returns (uint256) {
+        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS;
+        uint256 totalEntries = kvEntryCount + _batchSize; // include the batch to be put
+        uint256 shardIdNew = totalEntries >> SHARD_ENTRY_BITS; // shard id after the batch
+        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);
+        } else {
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
+        }
+        return totalPayment;
+    }
+
     /// @inheritdoc DecentralizedKV
     function _prepareAppend(uint256 _batchSize) internal virtual override {
         return _prepareAppendWithTimestamp(block.timestamp, _batchSize);

diff --git a/contracts/test/EthStorageContractTest.t.sol b/contracts/test/EthStorageContractTest.t.sol
@@ -75,7 +75,7 @@ contract EthStorageContractTest is Test {
         storageContract.putBlobs{value: insufficientCost}(keys, blobIdxs, lengths);
 
         // Enough storage cost
-        uint256 sufficientCost = 2 * storageContract.upfrontPayment();
+        uint256 sufficientCost = storageContract.upfrontPaymentInBatch(2);
         storageContract.putBlobs{value: sufficientCost}(keys, blobIdxs, lengths);
 
         assertEq(storageContract.kvEntryCount(), 2);
@@ -99,4 +99,28 @@ contract EthStorageContractTest is Test {
         assertEq(storageContract.size(bytes32(uint256(1))), 20);
         assertEq(storageContract.size(bytes32(uint256(2))), 30);
     }
+
+    function testPutBlobsXshard() public {
+        uint256 size = 5;
+        bytes32[] memory keys = new bytes32[](size);
+        uint256[] memory blobIdxs = new uint256[](size);
+        uint256[] memory lengths = new uint256[](size);
+        for (uint256 i = 0; i < size; i++) {
+            keys[i] = bytes32(uint256(i));
+            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);
+
+        // Enough storage cost
+        uint256 sufficientCost = storageContract.upfrontPaymentInBatch(size);
+        storageContract.putBlobs{value: sufficientCost}(keys, blobIdxs, lengths);
+
+        assertEq(storageContract.kvEntryCount(), size);
+    }
 }