fix(security): timelock storage slot, drip mutex, inflation pool deregister (#116)

Code-level follow-ups from the pre-mainnet audit (#115).

TangleTimelock
- updateDelay restored to onlySelf (parent's policy); the previous override
  silently degraded it to onlyRole(DEFAULT_ADMIN_ROLE).
- _setMinDelay now writes to the correct OZ ERC-7201 namespaced slot for
  TimelockControllerStorage._minDelay (TIMELOCK_CONTROLLER_STORAGE_LOCATION
  + 1) instead of slot 0x33. The old write was a no-op for getMinDelay()
  and could silently corrupt slot 51 of the proxy.
- Bounds (MIN_DELAY=1d, MAX_DELAY=30d) still enforced.

StreamingPaymentManager
- dripAndGetChunk and dripOperatorStreams now nonReentrant. Closes the
  same-tx double-drip race where the transfer-to-distributor hook could
  re-enter and process additional chunks before the first frame finished.

InflationPool
- New deregisterOperator / deregisterOperators / deregisterCustomer /
  deregisterDeveloper paths. Inactive operators no longer accumulate in
  trackedOperators, capping distribution-loop gas growth over time. The
  swap-and-pop helper clears membership flags and registration epochs;
  pending rewards remain claimable. Re-registration is permitted.

Tests
- test/security/TimelockSetMinDelayTest.t.sol: 3 tests (caller auth,
  persistence through getMinDelay, bounds enforcement).
- test/security/InflationPoolDeregisterTest.t.sol: 3 tests (removal,
  non-admin revert, re-registration).
- Full suite: 1478/1478 passing.

Author: drewstone

src/governance/TangleTimelock.sol

@@ -72,29 +72,33 @@ contract TangleTimelock is Initializable, TimelockControllerUpgradeable, UUPSUpg
     // M-16 FIX: DELAY VALIDATION
     // ═══════════════════════════════════════════════════════════════════════════
 
-    /// @notice Update the minimum delay
-    /// @dev M-16 FIX: Validates bounds before allowing delay changes. Only callable by
-    ///      the timelock itself (via a governance proposal).
-    /// @param newDelay The new delay to set
-    function updateDelay(uint256 newDelay) external override onlyRole(DEFAULT_ADMIN_ROLE) {
-        // M-16 FIX: Enforce delay bounds to prevent timelock bypass
+    /// @notice ERC-7201 storage location for `TimelockControllerStorage`.
+    ///         keccak256(abi.encode(uint256(keccak256("openzeppelin.storage.TimelockController")) - 1)) & ~bytes32(uint256(0xff))
+    /// @dev Pinned to OpenZeppelin Contracts Upgradeable 5.1.0; verify before bumping.
+    bytes32 private constant TIMELOCK_CONTROLLER_STORAGE_LOCATION =
+        0x9a37c2aa9d186a0969ff8a8267bf4e07e864c2f2768f5040949e28a624fb3600;
+
+    /// @notice Update the minimum delay. Only callable by the timelock itself via a
+    ///         scheduled governance proposal (matches the parent's `onlySelf` policy).
+    /// @param newDelay The new delay to set, bounded by [MIN_DELAY, MAX_DELAY].
+    function updateDelay(uint256 newDelay) external override {
+        if (msg.sender != address(this)) revert TimelockUnauthorizedCaller(msg.sender);
         if (newDelay < MIN_DELAY) revert DelayTooShort(newDelay, MIN_DELAY);
         if (newDelay > MAX_DELAY) revert DelayTooLong(newDelay, MAX_DELAY);
 
-        // Parent implementation handles the actual update and event emission
-        // Note: This calls TimelockControllerUpgradeable.updateDelay which is onlyRole(DEFAULT_ADMIN_ROLE)
         emit MinDelayChange(getMinDelay(), newDelay);
         _setMinDelay(newDelay);
     }
 
-    /// @notice Internal function to set the minimum delay (matching OZ 5.x pattern)
+    /// @dev Writes `newDelay` directly into the OZ ERC-7201 namespaced storage slot for
+    ///      `TimelockControllerStorage._minDelay`. The struct layout is:
+    ///         slot 0: _timestamps mapping
+    ///         slot 1: _minDelay
+    ///      so the absolute slot is `TIMELOCK_CONTROLLER_STORAGE_LOCATION + 1`.
     function _setMinDelay(uint256 newDelay) private {
-        // Storage slot for _minDelay in TimelockControllerUpgradeable
-        // Using direct storage write since OZ 5.x TimelockControllerUpgradeable
-        // doesn't expose a setter function (it's an immutable in the non-upgradeable version)
+        bytes32 slot = bytes32(uint256(TIMELOCK_CONTROLLER_STORAGE_LOCATION) + 1);
         assembly {
-            // _minDelay storage slot (position 0x33 = 51 in TimelockControllerUpgradeable)
-            sstore(0x33, newDelay)
+            sstore(slot, newDelay)
         }
     }
 

src/rewards/InflationPool.sol

@@ -980,6 +980,56 @@ contract InflationPool is Initializable, UUPSUpgradeable, AccessControlUpgradeab
         }
     }
 
+    /// @notice Remove an operator from `trackedOperators` so distributions stop iterating it.
+    /// @dev Pending rewards stay claimable. The address can be re-registered later if it
+    ///      becomes active again. Without this path the iteration cost grows monotonically.
+    function deregisterOperator(address operator) external onlyRole(ADMIN_ROLE) {
+        _removeFromTrackedList(trackedOperators, isTrackedOperator, operator);
+        delete operatorRegistrationEpoch[operator];
+        delete operatorLastScoredEpoch[operator];
+    }
+
+    /// @notice Batch deregister inactive operators.
+    function deregisterOperators(address[] calldata operators) external onlyRole(ADMIN_ROLE) {
+        for (uint256 i = 0; i < operators.length; i++) {
+            _removeFromTrackedList(trackedOperators, isTrackedOperator, operators[i]);
+            delete operatorRegistrationEpoch[operators[i]];
+            delete operatorLastScoredEpoch[operators[i]];
+        }
+    }
+
+    /// @notice Remove a customer from `trackedCustomers`.
+    function deregisterCustomer(address customer) external onlyRole(ADMIN_ROLE) {
+        _removeFromTrackedList(trackedCustomers, isTrackedCustomer, customer);
+        delete customerRegistrationEpoch[customer];
+    }
+
+    /// @notice Remove a developer from `trackedDevelopers`.
+    function deregisterDeveloper(address developer) external onlyRole(ADMIN_ROLE) {
+        _removeFromTrackedList(trackedDevelopers, isTrackedDeveloper, developer);
+        delete developerRegistrationEpoch[developer];
+    }
+
+    /// @dev Swap-and-pop the address out of the tracked list and clear its membership flag.
+    function _removeFromTrackedList(
+        address[] storage list,
+        mapping(address => bool) storage membership,
+        address account
+    ) internal {
+        if (!membership[account]) return;
+        uint256 len = list.length;
+        for (uint256 i = 0; i < len; i++) {
+            if (list[i] == account) {
+                if (i != len - 1) {
+                    list[i] = list[len - 1];
+                }
+                list.pop();
+                break;
+            }
+        }
+        membership[account] = false;
+    }
+
     // ═══════════════════════════════════════════════════════════════════════════
     // ADMIN CONFIGURATION
     // ═══════════════════════════════════════════════════════════════════════════

src/rewards/StreamingPaymentManager.sol

@@ -223,15 +223,21 @@ contract StreamingPaymentManager is
         return (chunk, durationSeconds);
     }
 
-    /// @notice Drip a specific stream and return chunk info for distribution
-    /// @dev Called by ServiceFeeDistributor to get the drip amount. Transfers dripped tokens to caller.
+    /// @notice Drip a specific stream and return chunk info for distribution.
+    /// @dev `nonReentrant` is load-bearing: without it, a re-entrant call from the
+    ///      transfer hook (or a same-tx call to `dripOperatorStreams`) would observe
+    ///      `lastDripTime == block.timestamp` only after the first call mutated it,
+    ///      but the second call could otherwise re-enter and process additional state
+    ///      before the first frame finishes. Combined with timestamp-granular dripping
+    ///      this also closes the same-block double-drip race.
     function dripAndGetChunk(
         uint64 serviceId,
         address operator
     )
         external
         override
         onlyRole(DISTRIBUTOR_ROLE)
+        nonReentrant
         returns (uint256 amount, uint256 durationSeconds, uint64 blueprintId, address paymentToken)
     {
         StreamingPayment storage p = streamingPayments[serviceId][operator];
@@ -245,12 +251,16 @@ contract StreamingPaymentManager is
         }
     }
 
-    /// @notice Drip all active streams for an operator
-    /// @dev Called by ServiceFeeDistributor before score changes. Transfers dripped tokens to caller.
+    /// @notice Drip all active streams for an operator.
+    /// @dev `nonReentrant` mutex prevents `dripOperatorStreams` and `dripAndGetChunk`
+    ///      from executing in the same transaction. Without it, two distributor calls
+    ///      in the same block would each compute the same `durationSeconds` and pay
+    ///      the chunk twice.
     function dripOperatorStreams(address operator)
         external
         override
         onlyRole(DISTRIBUTOR_ROLE)
+        nonReentrant
         returns (
             uint64[] memory serviceIds,
             uint64[] memory blueprintIds,

test/security/InflationPoolDeregisterTest.t.sol

@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { Test } from "forge-std/Test.sol";
+import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import { InflationPool } from "../../src/rewards/InflationPool.sol";
+
+/// @title InflationPoolDeregisterTest
+/// @notice Regression: `deregisterOperator` removes an operator from `trackedOperators`,
+///         clears its registration epoch, and lets the index shrink instead of growing
+///         monotonically as inactive operators accumulate.
+contract InflationPoolDeregisterTest is Test {
+    InflationPool pool;
+    address admin = makeAddr("admin");
+    address tnt = address(0xDEAD);
+
+    function setUp() public {
+        InflationPool impl = new InflationPool();
+        ERC1967Proxy proxy = new ERC1967Proxy(
+            address(impl),
+            abi.encodeCall(InflationPool.initialize, (admin, tnt, address(0), address(0), 7 days))
+        );
+        pool = InflationPool(payable(address(proxy)));
+    }
+
+    function _trackedCount() internal view returns (uint256 count) {
+        // `trackedOperators` is `address[] public`; length is exposed indirectly.
+        // Iterate by slot until we hit a revert or zero - simpler to just probe with try/catch.
+        while (true) {
+            try pool.trackedOperators(count) returns (address) {
+                count++;
+            } catch {
+                return count;
+            }
+        }
+    }
+
+    function test_DeregisterOperator_RemovesFromTrackedList() public {
+        address[] memory ops = new address[](3);
+        ops[0] = makeAddr("op1");
+        ops[1] = makeAddr("op2");
+        ops[2] = makeAddr("op3");
+
+        vm.prank(admin);
+        pool.registerOperators(ops);
+
+        assertEq(_trackedCount(), 3, "all three tracked");
+        assertTrue(pool.isTrackedOperator(ops[1]));
+
+        vm.prank(admin);
+        pool.deregisterOperator(ops[1]);
+
+        assertEq(_trackedCount(), 2, "shrinks after deregister");
+        assertFalse(pool.isTrackedOperator(ops[1]));
+        assertEq(pool.operatorRegistrationEpoch(ops[1]), 0, "registration epoch cleared");
+    }
+
+    function test_DeregisterOperator_NonAdminReverts() public {
+        address attacker = makeAddr("attacker");
+        vm.prank(admin);
+        pool.registerOperator(makeAddr("op"));
+
+        vm.prank(attacker);
+        vm.expectRevert();
+        pool.deregisterOperator(makeAddr("op"));
+    }
+
+    function test_DeregisterOperator_AllowsReregistration() public {
+        address op = makeAddr("op");
+
+        vm.startPrank(admin);
+        pool.registerOperator(op);
+        pool.deregisterOperator(op);
+        pool.registerOperator(op);
+        vm.stopPrank();
+
+        assertTrue(pool.isTrackedOperator(op));
+        assertEq(_trackedCount(), 1);
+    }
+}

test/security/TimelockSetMinDelayTest.t.sol

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { Test } from "forge-std/Test.sol";
+import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import { TangleTimelock } from "../../src/governance/TangleTimelock.sol";
+import {
+    TimelockControllerUpgradeable
+} from "@openzeppelin/contracts-upgradeable/governance/TimelockControllerUpgradeable.sol";
+
+/// @title TimelockSetMinDelayTest
+/// @notice Regression test: `updateDelay` must (a) reject callers other than the timelock
+///         itself and (b) actually persist the new delay (i.e. write to the OZ ERC-7201
+///         namespaced `_minDelay` slot, not slot 0x33).
+contract TimelockSetMinDelayTest is Test {
+    TangleTimelock timelock;
+    address admin = makeAddr("admin");
+
+    function setUp() public {
+        TangleTimelock impl = new TangleTimelock();
+        address[] memory empty = new address[](0);
+        ERC1967Proxy proxy = new ERC1967Proxy(
+            address(impl),
+            abi.encodeCall(TangleTimelock.initialize, (1 days, empty, empty, admin))
+        );
+        timelock = TangleTimelock(payable(address(proxy)));
+    }
+
+    function test_UpdateDelay_RevertsWhenCallerNotTimelock() public {
+        vm.prank(admin);
+        vm.expectRevert(
+            abi.encodeWithSelector(TimelockControllerUpgradeable.TimelockUnauthorizedCaller.selector, admin)
+        );
+        timelock.updateDelay(2 days);
+    }
+
+    function test_UpdateDelay_PersistsNewValue() public {
+        assertEq(timelock.getMinDelay(), 1 days, "initial delay");
+
+        // Calling from the timelock itself (the only caller `updateDelay` accepts) must
+        // (1) succeed and (2) flow through `getMinDelay()`. If `_setMinDelay` writes to
+        // the wrong slot, `getMinDelay()` will continue to return the old value.
+        vm.prank(address(timelock));
+        timelock.updateDelay(3 days);
+
+        assertEq(timelock.getMinDelay(), 3 days, "delay must persist after updateDelay");
+    }
+
+    function test_UpdateDelay_BoundsEnforced() public {
+        vm.prank(address(timelock));
+        vm.expectRevert(abi.encodeWithSelector(TangleTimelock.DelayTooShort.selector, 1 hours, 1 days));
+        timelock.updateDelay(1 hours);
+
+        vm.prank(address(timelock));
+        vm.expectRevert(abi.encodeWithSelector(TangleTimelock.DelayTooLong.selector, 60 days, 30 days));
+        timelock.updateDelay(60 days);
+    }
+}