fix/blacklist-low-stake (#33)

* make sure provider is blacklisted after work is invalidated if they lack sufficient stake

Author: mattdf

src/PrimeNetwork.sol

@@ -92,6 +92,10 @@ contract PrimeNetwork is AccessControlEnumerable {
         return domainId;
     }
 
+    function updateDomainValidationLogic(uint256 domainId, address validationLogic) external onlyRole(FEDERATOR_ROLE) {
+        domainRegistry.updateValidationLogic(domainId, validationLogic);
+    }
+
     function registerProvider(uint256 stake) external {
         uint256 stakeMinimum = stakeManager.getStakeMinimum();
         require(stake >= stakeMinimum, "Stake amount is below minimum");
@@ -168,19 +172,24 @@ contract PrimeNetwork is AccessControlEnumerable {
         emit ComputeNodeRemoved(provider, nodekey);
     }
 
-    function slash(address provider, uint256 amount, bytes calldata reason) external onlyRole(VALIDATOR_ROLE) {
-        uint256 slashed = stakeManager.slash(provider, amount, reason);
-        AIToken.transfer(msg.sender, slashed);
+    function _blacklistIfStakeTooLow(address provider) internal {
         if (stakeManager.getStake(provider) < calculateMinimumStake(provider, 0)) {
             computeRegistry.setWhitelistStatus(provider, false);
             emit ProviderBlacklisted(provider);
         }
     }
 
+    function slash(address provider, uint256 amount, bytes calldata reason) external onlyRole(VALIDATOR_ROLE) {
+        uint256 slashed = stakeManager.slash(provider, amount, reason);
+        AIToken.transfer(msg.sender, slashed);
+        _blacklistIfStakeTooLow(provider);
+    }
+
     function invalidateWork(uint256 poolId, uint256 penalty, bytes calldata data) external onlyRole(VALIDATOR_ROLE) {
         (address provider, address node) = computePool.invalidateWork(poolId, data);
-        try stakeManager.slash(provider, penalty, data) {}
-        catch {
+        try stakeManager.slash(provider, penalty, data) {
+            _blacklistIfStakeTooLow(provider);
+        } catch {
             // if slashing failed for whatever reason, blacklist provider to make sure they can't submit more work
             computeRegistry.setWhitelistStatus(provider, false);
             emit ProviderBlacklisted(provider);

test/PrimeNetwork.t.sol

@@ -16,6 +16,7 @@ import {RewardsDistributorFactory} from "../src/RewardsDistributorFactory.sol";
 import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
 import {IERC20Permit} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Permit.sol";
+import {SyntheticDataWorkValidator} from "../src/SyntheticDataWorkValidator.sol";
 
 contract PrimeNetworkTest is Test {
     using ECDSA for bytes32;
@@ -189,6 +190,14 @@ contract PrimeNetworkTest is Test {
         return primeNetwork.createDomain(name, IWorkValidation(address(0)), uri);
     }
 
+    function newDomainWithValidation(string memory name, string memory uri, address workValidator)
+        public
+        returns (uint256)
+    {
+        vm.startPrank(federator);
+        return primeNetwork.createDomain(name, IWorkValidation(workValidator), uri);
+    }
+
     function newPool(uint256 domainId, string memory name, string memory uri) public returns (uint256) {
         vm.startPrank(pool_creator);
         return computePool.createComputePool(domainId, computeManager, name, uri, 0);
@@ -729,4 +738,52 @@ contract PrimeNetworkTest is Test {
         withdrawStake(provider_good1);
         assertEq(AI.balanceOf(provider_good1), startingBalance);
     }
+
+    function test_stakeBlacklisting() public {
+        uint256 domain = newDomain("Decentralized Training", "https://primeintellect.ai/training/params");
+        uint256 pool = newPool(domain, "INTELLECT-1", "https://primeintellect.ai/pools/intellect-1");
+
+        SyntheticDataWorkValidator workValidator = new SyntheticDataWorkValidator(domain, address(computePool), 1 days);
+
+        vm.startPrank(federator);
+        primeNetwork.updateDomainValidationLogic(domain, address(workValidator));
+
+        uint256 minStakeNow = stakeManager.getStakeMinimum();
+        uint256 nodeComputeUnits = 1000;
+
+        bytes memory work_id = hex"1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
+
+        startPool(pool);
+
+        addProviderWithStake(provider_good1, minStakeNow);
+        whitelistProvider(provider_good1);
+
+        // add stake with some extra
+        increaseStake(provider_good1, nodeComputeUnits * 2 * minStakeNow);
+        addNodeWithCompute(provider_good1, node_good1, node_good1_sk, nodeComputeUnits);
+
+        // ensure that the stake is locked by the compute units
+        uint256 computeLockedStake = primeNetwork.calculateMinimumStake(provider_good1, 0);
+        assertEq(computeLockedStake, (nodeComputeUnits * minStakeNow) + minStakeNow);
+
+        validateNode(provider_good1, node_good1);
+        nodeJoin(domain, pool, provider_good1, node_good1);
+
+        computePool.submitWork(pool, node_good1, work_id);
+
+        // slash provider
+        uint256 stake = stakeManager.getStake(provider_good1);
+        vm.startPrank(validator);
+        primeNetwork.invalidateWork(pool, stake, work_id);
+
+        // check that stake has been slashed to 0
+        assertEq(stakeManager.getStake(provider_good1), 0);
+
+        // check that provider is now blacklisted
+        assertEq(computeRegistry.getProvider(provider_good1).isWhitelisted, false);
+
+        // check that provider is not allowed to add new nodes to the pool
+        vm.expectRevert();
+        nodeJoin(domain, pool, provider_good1, node_good1);
+    }
 }