TRST audit fixes for Subgraph Service contracts

packages/subgraph-service/contracts/SubgraphService.sol

@@ -312,11 +312,9 @@ contract SubgraphService is
     /**
      * @notice See {ISubgraphService.closeStaleAllocation}
      */
-    function forceCloseAllocation(address allocationId) external override {
+    function closeStaleAllocation(address allocationId) external override {
         Allocation.State memory allocation = allocations.get(allocationId);
-        bool isStale = allocation.isStale(maxPOIStaleness);
-        bool isOverAllocated_ = _isOverAllocated(allocation.indexer, delegationRatio);
-        require(isStale || isOverAllocated_, SubgraphServiceCannotForceCloseAllocation(allocationId));
+        require(allocation.isStale(maxPOIStaleness), SubgraphServiceCannotForceCloseAllocation(allocationId));
         require(!allocation.isAltruistic(), SubgraphServiceAllocationIsAltruistic(allocationId));
         _closeAllocation(allocationId);
     }
@@ -424,6 +422,7 @@ contract SubgraphService is
      * @return subgraphDeploymentId Subgraph deployment id for the allocation
      * @return tokens Amount of allocated tokens
      * @return accRewardsPerAllocatedToken Rewards snapshot
+     * @return accRewardsPending Rewards pending to be minted due to allocation resize
      */
     function getAllocationData(
         address allocationId

packages/subgraph-service/contracts/interfaces/ISubgraphService.sol

@@ -125,22 +125,20 @@ interface ISubgraphService is IDataServiceFees {
     error SubgraphServiceInvalidZeroStakeToFeesRatio();
 
     /**
-     * @notice Force close an allocation
-     * @dev This function can be permissionlessly called when the allocation is stale or
-     * if the indexer is over-allocated. This ensures that rewards for other allocations are
-     * not diluted by an inactive allocation, and that over-allocated indexers stop accumulating
-     * rewards with tokens they no longer have allocated.
+     * @notice Force close a stale allocation
+     * @dev This function can be permissionlessly called when the allocation is stale. This
+     * ensures that rewards for other allocations are not diluted by an inactive allocation.
      *
      * Requirements:
      * - Allocation must exist and be open
-     * - Allocation must be stale or indexer must be over-allocated
+     * - Allocation must be stale
      * - Allocation cannot be altruistic
      *
      * Emits a {AllocationClosed} event.
      *
      * @param allocationId The id of the allocation
      */
-    function forceCloseAllocation(address allocationId) external;
+    function closeStaleAllocation(address allocationId) external;
 
     /**
      * @notice Change the amount of tokens in an allocation

packages/subgraph-service/contracts/libraries/LegacyAllocation.sol

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity 0.8.27;
 
+import { IHorizonStaking } from "@graphprotocol/horizon/contracts/interfaces/IHorizonStaking.sol";
+
 /**
  * @title LegacyAllocation library
  * @notice A library to handle legacy Allocations.
@@ -22,20 +24,14 @@ library LegacyAllocation {
      * @notice Thrown when attempting to migrate an allocation with an existing id
      * @param allocationId The allocation id
      */
-    error LegacyAllocationExists(address allocationId);
+    error LegacyAllocationAlreadyExists(address allocationId);
 
     /**
      * @notice Thrown when trying to get a non-existent allocation
      * @param allocationId The allocation id
      */
     error LegacyAllocationDoesNotExist(address allocationId);
 
-    /**
-     * @notice Thrown when trying to migrate an allocation that has already been migrated
-     * @param allocationId The allocation id
-     */
-    error LegacyAllocationAlreadyMigrated(address allocationId);
-
     /**
      * @notice Migrate a legacy allocation
      * @dev Requirements:
@@ -52,7 +48,7 @@ library LegacyAllocation {
         address allocationId,
         bytes32 subgraphDeploymentId
     ) internal {
-        require(!self[allocationId].exists(), LegacyAllocationAlreadyMigrated(allocationId));
+        require(!self[allocationId].exists(), LegacyAllocationAlreadyExists(allocationId));
 
         State memory allocation = State({ indexer: indexer, subgraphDeploymentId: subgraphDeploymentId });
         self[allocationId] = allocation;
@@ -69,11 +65,19 @@ library LegacyAllocation {
 
     /**
      * @notice Revert if a legacy allocation exists
+     * @dev We first check the migrated mapping then the old staking contract.
+     * @dev TODO: after the transition period when all the allocations are migrated we can
+     * remove the call to the staking contract.
      * @param self The legacy allocation list mapping
      * @param allocationId The allocation id
      */
-    function revertIfExists(mapping(address => State) storage self, address allocationId) internal view {
-        require(!self[allocationId].exists(), LegacyAllocationExists(allocationId));
+    function revertIfExists(
+        mapping(address => State) storage self,
+        IHorizonStaking graphStaking,
+        address allocationId
+    ) internal view {
+        require(!self[allocationId].exists(), LegacyAllocationAlreadyExists(allocationId));
+        require(!graphStaking.isAllocation(allocationId), LegacyAllocationAlreadyExists(allocationId));
     }
 
     /**

packages/subgraph-service/contracts/utilities/AllocationManager.sol

@@ -212,7 +212,7 @@ abstract contract AllocationManager is EIP712Upgradeable, GraphDirectory, Alloca
 
         // Ensure allocation id is not reused
         // need to check both subgraph service (on allocations.create()) and legacy allocations
-        legacyAllocations.revertIfExists(_allocationId);
+        legacyAllocations.revertIfExists(_graphStaking(), _allocationId);
         Allocation.State memory allocation = allocations.create(
             _indexer,
             _allocationId,

packages/subgraph-service/test/shared/HorizonStakingShared.t.sol

@@ -5,11 +5,11 @@ import "forge-std/Test.sol";
 
 import { IGraphPayments } from "@graphprotocol/horizon/contracts/interfaces/IGraphPayments.sol";
 import { IHorizonStakingTypes } from "@graphprotocol/horizon/contracts/interfaces/internal/IHorizonStakingTypes.sol";
+import { IHorizonStakingExtension } from "@graphprotocol/horizon/contracts/interfaces/internal/IHorizonStakingExtension.sol";
 
 import { SubgraphBaseTest } from "../SubgraphBaseTest.t.sol";
 
 abstract contract HorizonStakingSharedTest is SubgraphBaseTest {
-
     /*
      * HELPERS
      */
@@ -45,11 +45,11 @@ abstract contract HorizonStakingSharedTest is SubgraphBaseTest {
     function _thawDeprovisionAndUnstake(address _indexer, address _verifier, uint256 _tokens) internal {
         // Initiate thaw request
         staking.thaw(_indexer, _verifier, _tokens);
-        
+
         // Skip thawing period
         IHorizonStakingTypes.Provision memory provision = staking.getProvision(_indexer, _verifier);
         skip(provision.thawingPeriod + 1);
-        
+
         // Deprovision and unstake
         staking.deprovision(_indexer, _verifier, 0);
         staking.unstake(_tokens);
@@ -64,6 +64,67 @@ abstract contract HorizonStakingSharedTest is SubgraphBaseTest {
         staking.setProvisionParameters(_indexer, _verifier, _maxVerifierCut, _thawingPeriod);
     }
 
+    function _setStorage_allocation_hardcoded(address indexer, address allocationId, uint256 tokens) internal {
+        IHorizonStakingExtension.Allocation memory allocation = IHorizonStakingExtension.Allocation({
+            indexer: indexer,
+            subgraphDeploymentID: bytes32("0x12344321"),
+            tokens: tokens,
+            createdAtEpoch: 1234,
+            closedAtEpoch: 1235,
+            collectedFees: 1234,
+            __DEPRECATED_effectiveAllocation: 1222234,
+            accRewardsPerAllocatedToken: 1233334,
+            distributedRebates: 1244434
+        });
+
+        // __DEPRECATED_allocations
+        uint256 allocationsSlot = 15;
+        bytes32 allocationBaseSlot = keccak256(abi.encode(allocationId, allocationsSlot));
+        vm.store(address(staking), allocationBaseSlot, bytes32(uint256(uint160(allocation.indexer))));
+        vm.store(address(staking), bytes32(uint256(allocationBaseSlot) + 1), allocation.subgraphDeploymentID);
+        vm.store(address(staking), bytes32(uint256(allocationBaseSlot) + 2), bytes32(tokens));
+        vm.store(address(staking), bytes32(uint256(allocationBaseSlot) + 3), bytes32(allocation.createdAtEpoch));
+        vm.store(address(staking), bytes32(uint256(allocationBaseSlot) + 4), bytes32(allocation.closedAtEpoch));
+        vm.store(address(staking), bytes32(uint256(allocationBaseSlot) + 5), bytes32(allocation.collectedFees));
+        vm.store(
+            address(staking),
+            bytes32(uint256(allocationBaseSlot) + 6),
+            bytes32(allocation.__DEPRECATED_effectiveAllocation)
+        );
+        vm.store(
+            address(staking),
+            bytes32(uint256(allocationBaseSlot) + 7),
+            bytes32(allocation.accRewardsPerAllocatedToken)
+        );
+        vm.store(address(staking), bytes32(uint256(allocationBaseSlot) + 8), bytes32(allocation.distributedRebates));
+
+        // _serviceProviders
+        uint256 serviceProviderSlot = 14;
+        bytes32 serviceProviderBaseSlot = keccak256(abi.encode(allocation.indexer, serviceProviderSlot));
+        uint256 currentTokensStaked = uint256(vm.load(address(staking), serviceProviderBaseSlot));
+        uint256 currentTokensProvisioned = uint256(
+            vm.load(address(staking), bytes32(uint256(serviceProviderBaseSlot) + 1))
+        );
+        vm.store(
+            address(staking),
+            bytes32(uint256(serviceProviderBaseSlot) + 0),
+            bytes32(currentTokensStaked + tokens)
+        );
+        vm.store(
+            address(staking),
+            bytes32(uint256(serviceProviderBaseSlot) + 1),
+            bytes32(currentTokensProvisioned + tokens)
+        );
+
+        // __DEPRECATED_subgraphAllocations
+        uint256 subgraphsAllocationsSlot = 16;
+        bytes32 subgraphAllocationsBaseSlot = keccak256(
+            abi.encode(allocation.subgraphDeploymentID, subgraphsAllocationsSlot)
+        );
+        uint256 currentAllocatedTokens = uint256(vm.load(address(staking), subgraphAllocationsBaseSlot));
+        vm.store(address(staking), subgraphAllocationsBaseSlot, bytes32(currentAllocatedTokens + tokens));
+    }
+
     /*
      * PRIVATE
      */

packages/subgraph-service/test/subgraphService/SubgraphService.t.sol

@@ -151,7 +151,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
         assertEq(afterSubgraphAllocatedTokens, _tokens);
     }
 
-    function _forceCloseAllocation(address _allocationId) internal {
+    function _closeStaleAllocation(address _allocationId) internal {
         assertTrue(subgraphService.isActiveAllocation(_allocationId));
 
         Allocation.State memory allocation = subgraphService.getAllocation(_allocationId);
@@ -168,7 +168,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
         );
 
         // close stale allocation
-        subgraphService.forceCloseAllocation(_allocationId);
+        subgraphService.closeStaleAllocation(_allocationId);
 
         // update allocation
         allocation = subgraphService.getAllocation(_allocationId);
@@ -380,6 +380,17 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
         }
     }
 
+    function _migrateLegacyAllocation(address _indexer, address _allocationId, bytes32 _subgraphDeploymentID) internal {
+        vm.expectEmit(address(subgraphService));
+        emit AllocationManager.LegacyAllocationMigrated(_indexer, _allocationId, _subgraphDeploymentID);
+
+        subgraphService.migrateLegacyAllocation(_indexer, _allocationId, _subgraphDeploymentID);
+
+        (address afterIndexer, bytes32 afterSubgraphDeploymentId) = subgraphService.legacyAllocations(_allocationId);
+        assertEq(afterIndexer, _indexer);
+        assertEq(afterSubgraphDeploymentId, _subgraphDeploymentID);
+    }
+
     /*
      * HELPERS
      */

packages/subgraph-service/test/subgraphService/allocation/forceClose.t.sol

@@ -3,38 +3,30 @@ pragma solidity 0.8.27;
 
 import "forge-std/Test.sol";
 
-import { IDataService } from "@graphprotocol/horizon/contracts/data-service/interfaces/IDataService.sol";
 import { IGraphPayments } from "@graphprotocol/horizon/contracts/interfaces/IGraphPayments.sol";
-import { ProvisionManager } from "@graphprotocol/horizon/contracts/data-service/utilities/ProvisionManager.sol";
-import { ProvisionTracker } from "@graphprotocol/horizon/contracts/data-service/libraries/ProvisionTracker.sol";
 
 import { Allocation } from "../../../contracts/libraries/Allocation.sol";
-import { AllocationManager } from "../../../contracts/utilities/AllocationManager.sol";
 import { ISubgraphService } from "../../../contracts/interfaces/ISubgraphService.sol";
-import { LegacyAllocation } from "../../../contracts/libraries/LegacyAllocation.sol";
 import { SubgraphServiceTest } from "../SubgraphService.t.sol";
 
 contract SubgraphServiceAllocationForceCloseTest is SubgraphServiceTest {
-
     address private permissionlessBob = makeAddr("permissionlessBob");
 
     /*
      * TESTS
      */
 
-    function test_SubgraphService_Allocation_ForceClose_Stale(
-        uint256 tokens
-    ) public useIndexer useAllocation(tokens) {
+    function test_SubgraphService_Allocation_ForceClose_Stale(uint256 tokens) public useIndexer useAllocation(tokens) {
         // Skip forward
         skip(maxPOIStaleness + 1);
 
         resetPrank(permissionlessBob);
-        _forceCloseAllocation(allocationID);
+        _closeStaleAllocation(allocationID);
     }
 
     function test_SubgraphService_Allocation_ForceClose_Stale_AfterCollecting(
         uint256 tokens
-    ) public useIndexer useAllocation(tokens) {      
+    ) public useIndexer useAllocation(tokens) {
         // Simulate POIs being submitted
         uint8 numberOfPOIs = 5;
         uint256 timeBetweenPOIs = 5 days;
@@ -52,43 +44,10 @@ contract SubgraphServiceAllocationForceCloseTest is SubgraphServiceTest {
 
         // Close the stale allocation
         resetPrank(permissionlessBob);
-        _forceCloseAllocation(allocationID);
-    }
-
-    function test_SubgraphService_Allocation_ForceClose_OverAllocated(
-        uint256 tokens
-    ) public useIndexer useAllocation(tokens) { 
-        // thaw some tokens to become over allocated
-        staking.thaw(users.indexer, address(subgraphService), tokens / 2);
-
-        resetPrank(permissionlessBob);
-        _forceCloseAllocation(allocationID);
-    }
-
-    function test_SubgraphService_Allocation_ForceClose_OverAllocated_AfterCollecting(
-        uint256 tokens
-    ) public useIndexer useAllocation(tokens) {      
-        // Simulate POIs being submitted
-        uint8 numberOfPOIs = 5;
-        uint256 timeBetweenPOIs = 5 days;
-
-        for (uint8 i = 0; i < numberOfPOIs; i++) {
-            // Skip forward
-            skip(timeBetweenPOIs);
-
-            bytes memory data = abi.encode(allocationID, bytes32("POI1"));
-            _collect(users.indexer, IGraphPayments.PaymentTypes.IndexingRewards, data);
-        }
-
-        // thaw some tokens to become over allocated
-        staking.thaw(users.indexer, address(subgraphService), tokens / 2);
-
-        // Close the over allocated allocation
-        resetPrank(permissionlessBob);
-        _forceCloseAllocation(allocationID);
+        _closeStaleAllocation(allocationID);
     }
 
-    function test_SubgraphService_Allocation_ForceClose_RevertIf_NotStaleOrOverAllocated(
+    function test_SubgraphService_Allocation_ForceClose_RevertIf_NotStale(
         uint256 tokens
     ) public useIndexer useAllocation(tokens) {
         // Simulate POIs being submitted
@@ -98,26 +57,24 @@ contract SubgraphServiceAllocationForceCloseTest is SubgraphServiceTest {
         for (uint8 i = 0; i < numberOfPOIs; i++) {
             // Skip forward
             skip(timeBetweenPOIs);
-            
+
             resetPrank(users.indexer);
 
             bytes memory data = abi.encode(allocationID, bytes32("POI1"));
             _collect(users.indexer, IGraphPayments.PaymentTypes.IndexingRewards, data);
-            
+
             resetPrank(permissionlessBob);
             vm.expectRevert(
                 abi.encodeWithSelector(
                     ISubgraphService.SubgraphServiceCannotForceCloseAllocation.selector,
                     allocationID
                 )
             );
-            subgraphService.forceCloseAllocation(allocationID);
+            subgraphService.closeStaleAllocation(allocationID);
         }
     }
 
-    function test_SubgraphService_Allocation_ForceClose_RevertIf_Altruistic(
-        uint256 tokens
-    ) public useIndexer {
+    function test_SubgraphService_Allocation_ForceClose_RevertIf_Altruistic(uint256 tokens) public useIndexer {
         tokens = bound(tokens, minimumProvisionTokens, MAX_TOKENS);
 
         _createProvision(users.indexer, tokens, maxSlashingPercentage, disputePeriod);
@@ -130,11 +87,8 @@ contract SubgraphServiceAllocationForceCloseTest is SubgraphServiceTest {
 
         resetPrank(permissionlessBob);
         vm.expectRevert(
-            abi.encodeWithSelector(
-                ISubgraphService.SubgraphServiceAllocationIsAltruistic.selector,
-                allocationID
-            )
+            abi.encodeWithSelector(ISubgraphService.SubgraphServiceAllocationIsAltruistic.selector, allocationID)
         );
-        subgraphService.forceCloseAllocation(allocationID);
+        subgraphService.closeStaleAllocation(allocationID);
     }
 }