diff --git a/contracts/utils/OnlyOnce.sol b/contracts/utils/OnlyOnce.sol new file mode 100644 index 0000000..0c21aa5 --- /dev/null +++ b/contracts/utils/OnlyOnce.sol @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2023 the ethier authors (github.com/divergencetech/ethier) +pragma solidity >=0.8.0 <0.9.0; + +import "forge-std/console2.sol"; + +contract OnlyOnce { + bytes32 private constant _LEFTMOST_FOUR_BYTES_MASK = + 0xFFFFFFFF00000000000000000000000000000000000000000000000000000000; + + /** + * @notice Thrown if a function can only be executed once. + */ + error FunctionAlreadyExecuted(bytes32 selector); + + /** + * @notice Keeps track of function executions. + */ + mapping(bytes32 => bool) private _functionAlreadyExecuted; + + /** + * @notice Ensures that a given identifier has not been marked as already + * executed and marks it as such. + */ + function _ensureOnlyOnce(bytes32 identifier) private { + if (_functionAlreadyExecuted[identifier]) { + revert FunctionAlreadyExecuted(identifier); + } + _functionAlreadyExecuted[identifier] = true; + } + + /** + * @notice Ensures that the modified function can only be executed once. + * @param identifier A generic UNIQUE identifier to mark the wrapped + * function as used. Typically the EVM function selector of the wrapped + * function (padded to the right with zeroes). + */ + modifier onlyOnceByIdentifier(bytes32 identifier) { + _ensureOnlyOnce(identifier); + _; + } + + /** + * @notice Ensures that the modified function can only be executed once. + * @dev This modifier MUST only be used on functions that are external (not + * public nor internal). The modifier uses the function selector of the + * current calldata context as identifier, which can have unintended + * side-effects for internally used functions. + */ + modifier onlyOnce() { + bytes32 selector; + assembly { + selector := and(calldataload(0), _LEFTMOST_FOUR_BYTES_MASK) + } + _ensureOnlyOnce(selector); + _; + } +} diff --git a/package.json b/package.json index 2bd6c63..3b1cba9 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@divergencetech/ethier", - "version": "0.46.0", + "version": "0.47.0", "description": "Golang and Solidity SDK to make Ethereum development ethier", "main": "\"\"", "scripts": { diff --git a/tests/utils/OnlyOnce.t.sol b/tests/utils/OnlyOnce.t.sol new file mode 100644 index 0000000..05ad520 --- /dev/null +++ b/tests/utils/OnlyOnce.t.sol @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2023 the ethier authors (github.com/divergencetech/ethier) +pragma solidity ^0.8.15; + +import "forge-std/Test.sol"; +import "forge-std/console2.sol"; + +import {OnlyOnce} from "../../contracts/utils/OnlyOnce.sol"; + +// solhint-disable no-empty-blocks +contract OnlyOnceConsumer is OnlyOnce { + bytes32 public constant SHARED_IDENTIFIER = keccak256("shared"); + + function autoLimited1() external onlyOnce {} + + function autoLimited2(uint256) external onlyOnce {} + + function explicitlyLimited1() + external + onlyOnceByIdentifier(OnlyOnceConsumer.explicitlyLimited1.selector) + {} + + function explicitlyLimited2() + external + onlyOnceByIdentifier(OnlyOnceConsumer.explicitlyLimited2.selector) + {} + + function sharedLimited1() + external + onlyOnceByIdentifier(SHARED_IDENTIFIER) + {} + + function sharedLimited2() + external + onlyOnceByIdentifier(SHARED_IDENTIFIER) + {} +} + +// solhint-enable no-empty-blocks + +contract OnlyOnceTest is Test { + OnlyOnceConsumer public c; + + function setUp() public { + c = new OnlyOnceConsumer(); + } + + function _expectRevertWithAlreadyExecuted(bytes32 identifier) internal { + vm.expectRevert( + abi.encodeWithSelector( + OnlyOnce.FunctionAlreadyExecuted.selector, + identifier + ) + ); + } + + function _testCannotCallTwice(function() external func) internal { + func(); + _expectRevertWithAlreadyExecuted(func.selector); + func(); + } + + function _testCannotCallTwice( + function(uint256) external func, + uint256 param1, + uint256 param2 + ) internal { + func(param1); + _expectRevertWithAlreadyExecuted(func.selector); + func(param2); + } + + function testCannotCallFunctionsTwice(uint256 param1, uint256 param2) + public + { + // This test also ensures that calling one limited function does not + // affect the executability of the other. + + _testCannotCallTwice(c.autoLimited1); + _testCannotCallTwice(c.autoLimited2, param1, param2); + _testCannotCallTwice(c.explicitlyLimited1); + _testCannotCallTwice(c.explicitlyLimited2); + } + + function testSharedIdentifier() public { + c.sharedLimited1(); + _expectRevertWithAlreadyExecuted(c.SHARED_IDENTIFIER()); + c.sharedLimited2(); + } +}