diff --git a/src/UntollHookSecurityExtensions.sol b/src/UntollHookSecurityExtensions.sol new file mode 100644 index 0000000..05c1ffe --- /dev/null +++ b/src/UntollHookSecurityExtensions.sol @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +library UntollHookSecurityExtensions { + error GraduationStateTransitionFailed(); + error PerBlockLimitExceeded(); + + struct HookState { + bool isGraduated; + uint256 blockVolume; + uint256 lastBlockNumber; + } + + function validateGraduationTransition(bool currentlyGraduated, bool targetGraduated) internal pure { + if (currentlyGraduated && !targetGraduated) { + revert GraduationStateTransitionFailed(); + } + } + + function checkBlockLimit(HookState storage state, uint256 swapAmount, uint256 maxBlockLimit) internal { + if (block.number > state.lastBlockNumber) { + state.lastBlockNumber = block.number; + state.blockVolume = 0; + } + + if (state.blockVolume + swapAmount > maxBlockLimit) { + revert PerBlockLimitExceeded(); + } + + state.blockVolume += swapAmount; + } +} diff --git a/test/UntollHookSecurity.t.sol b/test/UntollHookSecurity.t.sol new file mode 100644 index 0000000..3071788 --- /dev/null +++ b/test/UntollHookSecurity.t.sol @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "forge-std/Test.sol"; +import "../src/UntollHookSecurityExtensions.sol"; + +contract UntollHookSecurityTest is Test { + function testGraduationStateTransition() public pure { + UntollHookSecurityExtensions.validateGraduationTransition(false, true); + } + + function testFailInvalidGraduationRevert() public pure { + UntollHookSecurityExtensions.validateGraduationTransition(true, false); + } +}