From 606ab87b8640b792c2c65537ecf3d6351ea8e200 Mon Sep 17 00:00:00 2001
From: jihadMo <jihadMo@users.noreply.github.com>
Date: Thu, 6 Aug 2026 17:15:01 +0100
Subject: [PATCH] fix(security): add graduation state transition validation and
 unit tests (#1)

---
 src/UntollHookSecurityExtensions.sol | 32 ++++++++++++++++++++++++++++
 test/UntollHookSecurity.t.sol        | 15 +++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 src/UntollHookSecurityExtensions.sol
 create mode 100644 test/UntollHookSecurity.t.sol

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);
+    }
+}