Added BlockGuard (#63)

* init

* typos & tests

* changed version

contracts/mock/utils/BlockGuardMock.sol

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import "@openzeppelin/contracts/utils/Multicall.sol";
+
+import "../../utils/BlockGuard.sol";
+
+contract BlockGuardMock is Multicall, BlockGuard {
+    string public constant DEPOSIT_WITHDRAW_RESOURCE = "DEPOSIT_WITHDRAW";
+    string public constant LOCK_LOCK_RESOURCE = "LOCK_LOCK";
+
+    function deposit() external lockBlock(DEPOSIT_WITHDRAW_RESOURCE, msg.sender) {}
+
+    function withdraw() external checkBlock(DEPOSIT_WITHDRAW_RESOURCE, msg.sender) {}
+
+    function lock() external checkLockBlock(LOCK_LOCK_RESOURCE, msg.sender) {}
+
+    function getLatestLockBlock(
+        string memory resource_,
+        address key_
+    ) external view returns (uint256) {
+        return _getLatestLockBlock(resource_, key_);
+    }
+}

contracts/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@solarity/solidity-lib",
-  "version": "2.5.5",
+  "version": "2.5.6",
   "license": "MIT",
   "author": "Distributed Lab",
   "readme": "README.md",

contracts/utils/BlockGuard.sol

@@ -0,0 +1,59 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+/**
+ * @notice The BlockGuard module
+ *
+ * This module facilitates the flash-loan protection. Users may be prohibited from calling certain
+ * functions in the same block e.g. via the Multicall.
+ */
+abstract contract BlockGuard {
+    mapping(string => mapping(address => uint256)) private _lockedInBlocks;
+
+    modifier lockBlock(string memory resource_, address key_) {
+        _lockBlock(resource_, key_);
+        _;
+    }
+
+    modifier checkBlock(string memory resource_, address key_) {
+        _checkBlock(resource_, key_);
+        _;
+    }
+
+    modifier checkLockBlock(string memory resource_, address key_) {
+        _checkBlock(resource_, key_);
+        _lockBlock(resource_, key_);
+        _;
+    }
+
+    /**
+     * @notice The function to save the block when the resource key has been locked
+     * @param resource_ the id of the resource
+     * @param key_ the key of the resource
+     */
+    function _lockBlock(string memory resource_, address key_) internal {
+        _lockedInBlocks[resource_][key_] = block.number;
+    }
+
+    /**
+     * @notice The function to check if the resource key is called in the same block
+     * @param resource_ the id of the resource
+     * @param key_ the key of the resource
+     */
+    function _checkBlock(string memory resource_, address key_) internal view {
+        require(_lockedInBlocks[resource_][key_] != block.number, "BlockGuard: locked");
+    }
+
+    /**
+     * @notice The function to fetch the latest block when the resource key has been locked
+     * @param resource_ the id of the resource
+     * @param key_ the key of the resource
+     * @return block_ the block when the resource key has been locked
+     */
+    function _getLatestLockBlock(
+        string memory resource_,
+        address key_
+    ) internal view returns (uint256 block_) {
+        return _lockedInBlocks[resource_][key_];
+    }
+}

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@solarity/solidity-lib",
-  "version": "2.5.5",
+  "version": "2.5.6",
   "license": "MIT",
   "author": "Distributed Lab",
   "description": "Solidity Library by Distributed Lab",

test/helpers/block-helper.js

@@ -7,8 +7,12 @@ const setTime = async (time) => {
   await mine();
 };
 
+const getCurrentBlock = async () => {
+  return await web3.eth.getBlockNumber();
+};
+
 const getCurrentBlockTime = async () => {
-  return (await web3.eth.getBlock(await web3.eth.getBlockNumber())).timestamp;
+  return (await web3.eth.getBlock(await getCurrentBlock())).timestamp;
 };
 
 const mine = async (numberOfBlocks = 1) => {
@@ -18,6 +22,7 @@ const mine = async (numberOfBlocks = 1) => {
 };
 
 module.exports = {
+  getCurrentBlock,
   getCurrentBlockTime,
   setNextBlockTime,
   setTime,

test/utils/BlockGuard.js

@@ -0,0 +1,77 @@
+const { assert } = require("chai");
+const { accounts } = require("../../scripts/utils/utils");
+const truffleAssert = require("truffle-assertions");
+const { getCurrentBlock } = require("../helpers/block-helper");
+
+const BlockGuardMock = artifacts.require("BlockGuardMock");
+
+BlockGuardMock.numberFormat = "BigNumber";
+
+describe("BlockGuard", () => {
+  let mock;
+
+  let FIRST;
+  let SECOND;
+
+  before("setup", async () => {
+    FIRST = await accounts(0);
+    SECOND = await accounts(1);
+  });
+
+  beforeEach("setup", async () => {
+    mock = await BlockGuardMock.new();
+  });
+
+  describe("lockBlock", () => {
+    it("should return zero if the resource key hasn't been locked", async () => {
+      assert.equal((await mock.getLatestLockBlock(await mock.LOCK_LOCK_RESOURCE(), FIRST)).toFixed(), "0");
+    });
+
+    it("should return the current block if the resource key has been locked", async () => {
+      await mock.deposit();
+
+      assert.equal(
+        (await mock.getLatestLockBlock(await mock.DEPOSIT_WITHDRAW_RESOURCE(), FIRST)).toNumber(),
+        await getCurrentBlock()
+      );
+
+      assert.equal((await mock.getLatestLockBlock(await mock.DEPOSIT_WITHDRAW_RESOURCE(), SECOND)).toNumber(), "0");
+    });
+  });
+
+  describe("checkBlock", () => {
+    it("should allow to call in different blocks", async () => {
+      await mock.deposit();
+
+      await truffleAssert.passes(mock.withdraw());
+    });
+
+    it("should disallow to call in the same block", async () => {
+      await truffleAssert.reverts(
+        mock.multicall([mock.contract.methods.deposit().encodeABI(), mock.contract.methods.withdraw().encodeABI()]),
+        "BlockGuard: locked"
+      );
+    });
+  });
+
+  describe("checkLockBlock", () => {
+    it("should allow to call in different blocks", async () => {
+      await mock.lock();
+
+      const blockNumber = await getCurrentBlock();
+
+      assert.equal((await mock.getLatestLockBlock(await mock.LOCK_LOCK_RESOURCE(), FIRST)).toNumber(), blockNumber);
+
+      await mock.lock();
+
+      assert.equal((await mock.getLatestLockBlock(await mock.LOCK_LOCK_RESOURCE(), FIRST)).toNumber(), blockNumber + 1);
+    });
+
+    it("should disallow to call in the same block", async () => {
+      await truffleAssert.reverts(
+        mock.multicall([mock.contract.methods.lock().encodeABI(), mock.contract.methods.lock().encodeABI()]),
+        "BlockGuard: locked"
+      );
+    });
+  });
+});