Merge pull request #108 from ethstorage/reentrancy

Fix re-entrance

contracts/StorageContract.sol

@@ -1,14 +1,16 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.0;
 
+// TODO: upgrade OpenZeppelin to next release and import "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol"
+import "./libraries/ReentrancyGuardTransient.sol";
 import "./DecentralizedKV.sol";
 import "./libraries/MiningLib.sol";
 import "./libraries/RandaoLib.sol";
 
 /// @custom:upgradeable
 /// @title StorageContract
 /// @notice EthStorage L1 Contract with Decentralized KV Interface and Proof of Storage Verification
-abstract contract StorageContract is DecentralizedKV {
+abstract contract StorageContract is DecentralizedKV, ReentrancyGuardTransient {
     /// @notice Represents the configuration of the storage contract.
     /// @custom:field maxKvSizeBits  Maximum size of a single key-value pair.
     /// @custom:field shardSizeBits  Storage shard size.
@@ -78,7 +80,7 @@ abstract contract StorageContract is DecentralizedKV {
     /// @notice Treasury address
     address public treasury;
 
-    /// @notice
+    /// @notice Prepaid timestamp of last mined
     uint256 public prepaidLastMineTime;
 
     // TODO: Reserve extra slots (to a total of 50?) in the storage layout for future upgrades
@@ -239,7 +241,10 @@ abstract contract StorageContract is DecentralizedKV {
         MiningLib.update(infos[_shardId], _minedTs, _diff);
 
         require(treasuryReward + minerReward <= address(this).balance, "StorageContract: not enough balance");
-        // TODO: avoid reentrancy attack
+        // Actually `transfer` is limited by the amount of gas allocated, which is not sufficient to enable reentrancy attacks.
+        // However, this behavior may restrict the extensibility of scenarios where the receiver is a contract that requires
+        // additional gas for its fallback functions of proper operations.
+        // Therefore, we use `ReentrancyGuard` in case `call` replaces `transfer` in the future.
         payable(treasury).transfer(treasuryReward);
         payable(_miner).transfer(minerReward);
         emit MinedBlock(_shardId, _diff, infos[_shardId].blockMined, _minedTs, _miner, minerReward);
@@ -307,7 +312,7 @@ abstract contract StorageContract is DecentralizedKV {
         bytes calldata _randaoProof,
         bytes[] calldata _inclusiveProofs,
         bytes[] calldata _decodeProof
-    ) public virtual {
+    ) public virtual nonReentrant {
         _mine(
             _blockNum, _shardId, _miner, _nonce, _encodedSamples, _masks, _randaoProof, _inclusiveProofs, _decodeProof
         );

contracts/libraries/ReentrancyGuardTransient.sol

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {StorageSlot} from "./StorageSlot.sol";
+
+/**
+ * @dev Variant of {ReentrancyGuard} that uses transient storage.
+ *
+ * NOTE: This variant only works on networks where EIP-1153 is available.
+ */
+abstract contract ReentrancyGuardTransient {
+    using StorageSlot for *;
+
+    // keccak256(abi.encode(uint256(keccak256("openzeppelin.storage.ReentrancyGuard")) - 1)) & ~bytes32(uint256(0xff))
+    bytes32 private constant REENTRANCY_GUARD_STORAGE =
+        0x9b779b17422d0df92223018b32b4d1fa46e071723d6817e2486d003becc55f00;
+
+    /**
+     * @dev Unauthorized reentrant call.
+     */
+    error ReentrancyGuardReentrantCall();
+
+    /**
+     * @dev Prevents a contract from calling itself, directly or indirectly.
+     * Calling a `nonReentrant` function from another `nonReentrant`
+     * function is not supported. It is possible to prevent this from happening
+     * by making the `nonReentrant` function external, and making it call a
+     * `private` function that does the actual work.
+     */
+    modifier nonReentrant() {
+        _nonReentrantBefore();
+        _;
+        _nonReentrantAfter();
+    }
+
+    function _nonReentrantBefore() private {
+        // On the first call to nonReentrant, _status will be NOT_ENTERED
+        if (_reentrancyGuardEntered()) {
+            revert ReentrancyGuardReentrantCall();
+        }
+
+        // Any calls to nonReentrant after this point will fail
+        REENTRANCY_GUARD_STORAGE.asBoolean().tstore(true);
+    }
+
+    function _nonReentrantAfter() private {
+        REENTRANCY_GUARD_STORAGE.asBoolean().tstore(false);
+    }
+
+    /**
+     * @dev Returns true if the reentrancy guard is currently set to "entered", which indicates there is a
+     * `nonReentrant` function in the call stack.
+     */
+    function _reentrancyGuardEntered() internal view returns (bool) {
+        return REENTRANCY_GUARD_STORAGE.asBoolean().tload();
+    }
+}