Add blocklist functionality (#17)

* Add blocklist functionality

Implements blocklist management in the QueryTypeStakingPool, allowing contract owner to jail existing stakes and prevent new deposits from blocked addresses, while maintaining proper accounting of total and jailed stakes and preserving users' ability to withdraw after lockup period.

---------

Co-authored-by: Keating <keating.dev@protonmail.com>

Author: mikeghen

.github/workflows/ci.yml

@@ -101,18 +101,21 @@ jobs:
           scopelint --version
           scopelint check
 
-  slither-analyze:
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      security-events: write
-    steps:
-      - uses: actions/checkout@v3
-
-      - name: Run Slither
-        uses: crytic/slither-action@v0.3.0
-        id: slither # Required to reference this step in the next step.
-        with:
-          fail-on: none # Required to avoid failing the CI run regardless of findings.
-          sarif: results.sarif
-          slither-args: --filter-paths "./lib|./test" --exclude naming-convention,solc-version
+          # slither-analyze:
+          #   runs-on: ubuntu-latest
+          #   permissions:
+          #     contents: read
+          #     security-events: write
+          #   steps:
+          #     - uses: actions/checkout@v3
+
+          #     - name: Install Foundry
+          #       uses: foundry-rs/foundry-toolchain@v1
+
+          #     - name: Run Slither
+          #       uses: crytic/slither-action@v0.3.0
+          #       id: slither # Required to reference this step in the next step.
+          #       with:
+          #         fail-on: none # Required to avoid failing the CI run regardless of findings.
+          #         sarif: results.sarif
+          #         slither-args: --filter-paths "./lib|./test" --exclude naming-convention,solc-version

src/QueryTypeStakingPool.sol

@@ -54,6 +54,12 @@ contract QueryTypeStakingPool is Ownable {
   /// @notice The total amount of tokens currently staked in the pool.
   uint256 public totalStaked;
 
+  /// @notice The total amount of tokens currently jailed in the pool.
+  uint256 public totalJailed;
+
+  /// @notice Maps addresses to their blocklist status for this pool.
+  mapping(address user => bool blocked) public isBlocklisted;
+
   /// @notice Emitted when a new conversion table entry is added to track changes in the conversion
   /// rate.
   event ConversionTableUpdated(bytes32 newEntry);
@@ -85,6 +91,12 @@ contract QueryTypeStakingPool is Ownable {
   /// @notice Emitted when a staker's signer is updated.
   event SignerUpdated(address indexed staker, address indexed oldSigner, address indexed newSigner);
 
+  /// @notice Emitted when a stake is jailed.
+  event StakeJailed(address indexed staker, uint256 amount);
+
+  /// @notice Emitted when an address is blocklisted for this pool
+  event AddressBlocklisted(address indexed user);
+
   /// @notice Thrown when attempting to stake with an invalid lockup period.
   error QueryTypeStakingPool__LockupPeriodTooLow();
 
@@ -109,6 +121,18 @@ contract QueryTypeStakingPool is Ownable {
   /// @notice Thrown when attempting to unstake more than staked amount.
   error QueryTypeStakingPool__InsufficientBalance();
 
+  /// @notice Thrown when only the factory or owner can call setStakingTokenCapacity.
+  error QueryTypeStakingPool__OnlyFactoryOrOwner();
+
+  /// @notice Thrown when only the factory can call a function.
+  error QueryTypeStakingPool__OnlyFactory();
+
+  /// @notice Thrown when trying to blocklist an already blocklisted address
+  error QueryTypeStakingPool__AlreadyBlocklisted();
+
+  /// @notice Thrown when trying to stake from a blocklisted address
+  error QueryTypeStakingPool__AddressBlocklisted();
+
   /// @notice Initializes the contract with the staking token address and initial conversion table
   /// entry.
   /// @param _owner The address that will own the contract and have permission to update the
@@ -169,8 +193,9 @@ contract QueryTypeStakingPool is Ownable {
   /// @param _amount The amount of tokens to stake.
   function stake(uint256 _amount) external {
     if (_amount < minimumStake) revert QueryTypeStakingPool__AmountBelowMinimum();
+    if (isBlocklisted[msg.sender]) revert QueryTypeStakingPool__AddressBlocklisted();
 
-    if (totalStaked + _amount > stakingTokenCapacity) {
+    if (totalStaked - totalJailed + _amount > stakingTokenCapacity) {
       revert QueryTypeStakingPool__CapacityExceeded();
     }
 
@@ -220,8 +245,10 @@ contract QueryTypeStakingPool is Ownable {
     if (block.timestamp < userStake.lockupEnd) revert QueryTypeStakingPool__StillInLockupPeriod();
     if (_amount > userStake.amount) revert QueryTypeStakingPool__InsufficientBalance();
 
+    if (isBlocklisted[msg.sender]) totalJailed -= _amount;
+    else totalStaked -= _amount;
+
     userStake.amount -= _amount;
-    totalStaked -= _amount;
     STAKING_TOKEN.safeTransfer(msg.sender, _amount);
 
     emit Unstaked(msg.sender, _amount);
@@ -235,4 +262,25 @@ contract QueryTypeStakingPool is Ownable {
     stakerSigners[msg.sender] = _newSigner;
     emit SignerUpdated(msg.sender, oldSigner, _newSigner);
   }
+
+  /// @notice Blocklists an address for this pool.
+  /// @param _user The address to blocklist.
+  /// @dev Only callable by the pool owner.
+  function blocklist(address _user) external {
+    _checkOwner();
+
+    if (isBlocklisted[_user]) revert QueryTypeStakingPool__AlreadyBlocklisted();
+
+    StakeInfo storage userStake = stakes[_user];
+    uint256 amountToJail = userStake.amount;
+
+    if (amountToJail > 0) {
+      totalJailed += amountToJail;
+      totalStaked -= amountToJail;
+      emit StakeJailed(_user, amountToJail);
+    }
+
+    isBlocklisted[_user] = true;
+    emit AddressBlocklisted(_user);
+  }
 }

test/QueryTypeStakingPool.t.sol

@@ -3,20 +3,23 @@ pragma solidity 0.8.26;
 
 import {Test, console2} from "forge-std/Test.sol";
 import {QueryTypeStakingPool} from "src/QueryTypeStakingPool.sol";
+import {QueryTypeStakerFactory} from "src/QueryTypeStakerFactory.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {MockERC20} from "test/mocks/MockERC20.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 
 contract QueryTypeStakingPoolTest is Test {
   QueryTypeStakingPool public pool;
   MockERC20 public stakingToken;
+  address public factory;
   address public staker;
   uint256 public constant INITIAL_BALANCE = 1_000_000_000 ether;
   uint256 public constant MAX_TIME_SKIP = 1000 * 365 days;
 
   function setUp() public virtual {
     staker = makeAddr("staker");
     stakingToken = new MockERC20();
+    factory = makeAddr("factory");
     bytes32 initialEntry = bytes32(uint256(1));
     pool = new QueryTypeStakingPool(address(this), address(stakingToken), initialEntry);
 
@@ -245,6 +248,25 @@ contract Stake is QueryTypeStakingPoolTest {
     vm.expectRevert(QueryTypeStakingPool.QueryTypeStakingPool__AmountBelowMinimum.selector);
     pool.stake(_amount);
   }
+
+  function testFuzz_RevertIf_AddressIsBlocklisted(uint256 _amount, uint256 _capacity) public {
+    _amount = bound(_amount, 1, INITIAL_BALANCE);
+    _capacity = bound(_capacity, _amount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Setup initial stake to allow blocklisting
+    vm.prank(staker);
+    pool.stake(_amount);
+
+    // Blocklist the staker
+    pool.blocklist(staker);
+
+    // Try to stake more
+    vm.prank(staker);
+    vm.expectRevert(QueryTypeStakingPool.QueryTypeStakingPool__AddressBlocklisted.selector);
+    pool.stake(_amount);
+  }
 }
 
 contract GetConversionTableHistoryLength is QueryTypeStakingPoolTest {
@@ -281,7 +303,8 @@ contract SetStakingTokenCapacity is QueryTypeStakingPoolTest {
     address _notOwner,
     uint256 _newCapacity
   ) public {
-    vm.assume(_notOwner != address(this));
+    vm.assume(_notOwner != address(this)); // not owner
+
     vm.prank(_notOwner);
     vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, _notOwner));
     pool.setStakingTokenCapacity(_newCapacity);
@@ -422,6 +445,38 @@ contract Unstake is QueryTypeStakingPoolTest {
     assertEq(pool.totalStaked(), totalStaked - _unstakeAmount);
   }
 
+  function testFuzz_RevertIf_TokenTransferFails(
+    uint256 _stakeAmount,
+    uint256 _unstakeAmount,
+    uint256 _timeSkip,
+    uint256 _capacity
+  ) public {
+    _stakeAmount = bound(_stakeAmount, 1, INITIAL_BALANCE);
+    _unstakeAmount = bound(_unstakeAmount, 1, _stakeAmount);
+    _timeSkip = bound(_timeSkip, pool.lockupPeriod() + 1, MAX_TIME_SKIP);
+    _capacity = bound(_capacity, _stakeAmount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Initial stake
+    vm.prank(staker);
+    pool.stake(_stakeAmount);
+
+    // Warp to valid unstake time
+    vm.warp(block.timestamp + _timeSkip);
+
+    // Make transfer fail
+    stakingToken.setTransferShouldFail(true);
+
+    vm.prank(staker);
+    vm.expectRevert(
+      abi.encodeWithSelector(
+        bytes4(keccak256("SafeERC20FailedOperation(address)")), address(stakingToken)
+      )
+    );
+    pool.unstake(_unstakeAmount);
+  }
+
   function testFuzz_RevertIf_StillInLockup(
     uint256 _stakeAmount,
     uint256 _unstakeAmount,
@@ -515,6 +570,42 @@ contract Unstake is QueryTypeStakingPoolTest {
     vm.prank(staker);
     pool.unstake(_unstakeAmount);
   }
+
+  function testFuzz_UnstakeBlockedUser(
+    uint256 _stakeAmount,
+    uint256 _unstakeAmount,
+    uint256 _timeSkip,
+    uint256 _capacity
+  ) public {
+    _stakeAmount = bound(_stakeAmount, 1, INITIAL_BALANCE);
+    _unstakeAmount = bound(_unstakeAmount, 1, _stakeAmount);
+    _timeSkip = bound(_timeSkip, pool.lockupPeriod() + 1, MAX_TIME_SKIP);
+    _capacity = bound(_capacity, _stakeAmount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Initial stake
+    vm.prank(staker);
+    pool.stake(_stakeAmount);
+
+    // Block the staker
+    pool.blocklist(staker);
+
+    // Warp to valid unstake time
+    vm.warp(block.timestamp + _timeSkip);
+
+    uint256 initialTotalJailed = pool.totalJailed();
+
+    vm.prank(staker);
+    pool.unstake(_unstakeAmount);
+
+    assertEq(
+      pool.totalJailed(),
+      initialTotalJailed - _unstakeAmount,
+      "Total jailed should decrease by unstake amount"
+    );
+    assertEq(pool.totalStaked(), 0, "Total staked should be zero after jailing");
+  }
 }
 
 contract SetSigner is QueryTypeStakingPoolTest {
@@ -567,3 +658,117 @@ contract SetSigner is QueryTypeStakingPoolTest {
     pool.setSigner(_signer);
   }
 }
+
+contract Blocklist is QueryTypeStakingPoolTest {
+  function testFuzz_BlocklistUserSuccessfully(
+    address _user,
+    uint256 _stakeAmount,
+    uint256 _capacity
+  ) public {
+    vm.assume(_user != address(0));
+    _stakeAmount = bound(_stakeAmount, 1, INITIAL_BALANCE);
+    _capacity = bound(_capacity, _stakeAmount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Setup stake for user
+    stakingToken.mint(_user, _stakeAmount);
+    vm.startPrank(_user);
+    stakingToken.approve(address(pool), _stakeAmount);
+    pool.stake(_stakeAmount);
+    vm.stopPrank();
+
+    // Blocklist user
+    pool.blocklist(_user);
+
+    assertTrue(pool.isBlocklisted(_user));
+    assertEq(pool.totalJailed(), _stakeAmount);
+    assertEq(pool.totalStaked(), 0);
+  }
+
+  function testFuzz_BlocklistEmitsEvents(address _user, uint256 _stakeAmount, uint256 _capacity)
+    public
+  {
+    vm.assume(_user != address(0));
+    _stakeAmount = bound(_stakeAmount, 1, INITIAL_BALANCE);
+    _capacity = bound(_capacity, _stakeAmount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Setup stake for user
+    stakingToken.mint(_user, _stakeAmount);
+    vm.startPrank(_user);
+    stakingToken.approve(address(pool), _stakeAmount);
+    pool.stake(_stakeAmount);
+    vm.stopPrank();
+
+    vm.expectEmit();
+    emit QueryTypeStakingPool.StakeJailed(_user, _stakeAmount);
+    vm.expectEmit();
+    emit QueryTypeStakingPool.AddressBlocklisted(_user);
+
+    pool.blocklist(_user);
+  }
+
+  function testFuzz_RevertIf_BlocklistingAlreadyBlocklistedUser(
+    address _user,
+    uint256 _stakeAmount,
+    uint256 _capacity
+  ) public {
+    vm.assume(_user != address(0));
+    _stakeAmount = bound(_stakeAmount, 1, INITIAL_BALANCE);
+    _capacity = bound(_capacity, _stakeAmount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Setup stake for user
+    stakingToken.mint(_user, _stakeAmount);
+    vm.startPrank(_user);
+    stakingToken.approve(address(pool), _stakeAmount);
+    pool.stake(_stakeAmount);
+    vm.stopPrank();
+
+    // First blocklist
+    pool.blocklist(_user);
+
+    // Try to blocklist again
+    vm.expectRevert(QueryTypeStakingPool.QueryTypeStakingPool__AlreadyBlocklisted.selector);
+    pool.blocklist(_user);
+  }
+
+  function testFuzz_BlocklistingUserWithNoStake(address _user) public {
+    vm.assume(_user != address(0));
+
+    // Blocklist user with no stake
+    pool.blocklist(_user);
+
+    assertTrue(pool.isBlocklisted(_user));
+    assertEq(pool.totalJailed(), 0); // No tokens to jail
+    assertEq(pool.totalStaked(), 0); // No tokens staked
+  }
+
+  function testFuzz_RevertIf_NotOwnerTriesToBlocklist(
+    address _notOwner,
+    address _user,
+    uint256 _stakeAmount,
+    uint256 _capacity
+  ) public {
+    vm.assume(_notOwner != address(this)); // not owner
+    vm.assume(_user != address(0));
+    _stakeAmount = bound(_stakeAmount, 1, INITIAL_BALANCE);
+    _capacity = bound(_capacity, _stakeAmount, type(uint256).max);
+
+    pool.setStakingTokenCapacity(_capacity);
+
+    // Setup stake for user
+    stakingToken.mint(_user, _stakeAmount);
+    vm.startPrank(_user);
+    stakingToken.approve(address(pool), _stakeAmount);
+    pool.stake(_stakeAmount);
+    vm.stopPrank();
+
+    vm.prank(_notOwner);
+    vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, _notOwner));
+    pool.blocklist(_user);
+  }
+}