neobase-one · zjesko · Dec 23, 2024
diff --git a/contracts/mocks/ExampleERC721.sol b/contracts/mocks/ExampleERC721.sol
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Ecosystem
+pragma solidity 0.8.25;
+
+import {ERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol";
+
+contract ExampleERC721 is ERC721 {
+    constructor() ERC721("Example NFT", "ENFT") {
+        _mint(msg.sender, 1);
+        _mint(msg.sender, 2);
+        _mint(msg.sender, 3);
+        _mint(msg.sender, 4);
+        _mint(msg.sender, 5);
+        _mint(msg.sender, 6);
+    }
+
+    function mint(address to, uint256 tokenId) external {
+        _mint(to, tokenId);
+    }
+
+}
diff --git a/contracts/validator-manager/ERC20TokenStakingManager.sol b/contracts/validator-manager/ERC20TokenStakingManager.sol
@@ -35,6 +35,7 @@ contract ERC20TokenStakingManager is
         IERC20Mintable _token;
         uint8 _tokenDecimals;
     }
+
     // solhint-enable private-vars-leading-underscore
 
     // keccak256(abi.encode(uint256(keccak256("avalanche-icm.storage.ERC20TokenStakingManager")) - 1)) & ~bytes32(uint256(0xff));
@@ -55,7 +56,20 @@ contract ERC20TokenStakingManager is
         }
     }
 
-    constructor(ICMInitializable init) {
+    function _addValidatorNft(bytes32 validationID, uint256 tokenId) internal override {}
+
+    function _addDelegatorNft(bytes32 delegationID, uint256 tokenId) internal override {}
+
+    function _deleteValidatorNft(
+        bytes32 validationID
+    ) internal override {}
+    function _deleteDelegatorNft(
+        bytes32 delegationID
+    ) internal override {}
+
+    constructor(
+        ICMInitializable init
+    ) {
         if (init == ICMInitializable.Disallowed) {
             _disableInitializers();
         }
@@ -138,7 +152,9 @@ contract ERC20TokenStakingManager is
      * @notice See {PoSValidatorManager-_unlock}
      * Note: Must be guarded with reentrancy guard for safe transfer.
      */
-    function _unlock(address to, uint256 value) internal virtual override {
+    function _unlock(address to, bytes32 id, bool isValidator) internal virtual override {
+        uint64 weight = isValidator ? getValidator(id).startingWeight : getDelegator(id).weight;
+        uint256 value = weightToValue(weight);
         _getERC20StakingManagerStorage()._token.safeTransfer(to, value);
     }
 

diff --git a/contracts/validator-manager/ERC721TokenStakingManager.sol b/contracts/validator-manager/ERC721TokenStakingManager.sol
@@ -0,0 +1,231 @@
+// (c) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+// SPDX-License-Identifier: Ecosystem
+
+pragma solidity 0.8.25;
+
+import {PoSValidatorManager} from "./PoSValidatorManager.sol";
+import {
+    PoSValidatorManagerSettings,
+    PoSValidatorManagerStorage
+} from "./interfaces/IPoSValidatorManager.sol";
+import {
+    ValidatorRegistrationInput, ValidatorManagerStorage
+} from "./interfaces/IValidatorManager.sol";
+import {IERC721TokenStakingManager} from "./interfaces/IERC721TokenStakingManager.sol";
+import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import {ICMInitializable} from "@utilities/ICMInitializable.sol";
+import {Initializable} from "@openzeppelin/[email protected]/proxy/utils/Initializable.sol";
+import {AccessControlUpgradeable} from "@openzeppelin/[email protected]/access/AccessControlUpgradeable.sol";
+
+/**
+ * @dev Implementation of the {IERC721TokenStakingManager} interface.
+ *
+ * @custom:security-contact https://github.com/ava-labs/icm-contracts/blob/main/SECURITY.md
+ */
+contract ERC721TokenStakingManager is
+    Initializable,
+    AccessControlUpgradeable,
+    PoSValidatorManager,
+    IERC721TokenStakingManager
+{
+    using SafeERC20 for IERC20;
+
+    // solhint-disable private-vars-leading-underscore
+    /// @custom:storage-location erc7201:avalanche-icm.storage.ERC721TokenStakingManager
+    struct ERC721TokenStakingManagerStorage {
+        IERC721 _token;
+        IERC20 _rewardToken;
+    }
+    // solhint-enable private-vars-leading-underscore
+
+    // keccak256(abi.encode(uint256(keccak256("avalanche-icm.storage.ERC721TokenStakingManager")) - 1)) & ~bytes32(uint256(0xff));
+    bytes32 public constant ERC721_STAKING_MANAGER_STORAGE_LOCATION =
+        0xf2d79c30881febd0da8597832b5b1bf1f4d4b2209b19059420303eb8fcab8a00;
+
+    error InvalidTokenAddress(address tokenAddress);
+    error InvalidRewardTokenAddress(address tokenAddress);
+
+
+    // solhint-disable ordering
+    function _getERC721StakingManagerStorage()
+        private
+        pure
+        returns (ERC721TokenStakingManagerStorage storage $)
+    {
+        assembly {
+            $.slot := ERC721_STAKING_MANAGER_STORAGE_LOCATION
+        }
+    }
+
+    function _addValidatorNft(bytes32 validationID, uint256 tokenId) internal override {
+        ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
+        $._validatorNFTs[validationID].nftIds.push(tokenId);
+    }
+
+    function _addDelegatorNft(bytes32 delegationID, uint256 tokenId) internal override {
+        PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
+        $._delegatorNFTs[delegationID].nftIds.push(tokenId);
+    }
+
+    function _deleteValidatorNft(
+        bytes32 validationID
+    ) internal override {
+        ValidatorManagerStorage storage $ = _getValidatorManagerStorage();
+        delete $._validatorNFTs[validationID];
+    }
+
+    function _deleteDelegatorNft(
+        bytes32 delegationID
+    ) internal override {
+        PoSValidatorManagerStorage storage $ = _getPoSValidatorManagerStorage();
+        delete $._delegatorNFTs[delegationID];
+    }
+
+    bytes32 public constant OPERATOR_ROLE = keccak256("OPERATOR_ROLE");
+
+    modifier onlyOperator() {
+        require(hasRole(OPERATOR_ROLE, msg.sender), "ERC721TokenStakingManager: caller is not an operator");
+        _;
+    }
+
+    constructor(ICMInitializable init) {
+        if (init == ICMInitializable.Disallowed) {
+            _disableInitializers();
+            _grantRole(DEFAULT_ADMIN_ROLE, msg.sender);
+        }
+    }
+
+    /**
+     * @notice Initialize the ERC721 token staking manager
+     * @dev Uses reinitializer(2) on the PoS staking contracts to make sure after migration from PoA, the PoS contracts can reinitialize with its needed values.
+     * @param settings Initial settings for the PoS validator manager
+     * @param stakingToken The ERC721 token to be staked
+     * @param rewardToken The ERC20 token to be used for rewards
+     */
+    function initialize(
+        PoSValidatorManagerSettings calldata settings,
+        IERC721 stakingToken,
+        IERC20 rewardToken
+    ) external reinitializer(2) {
+        __ERC721TokenStakingManager_init(settings, stakingToken, rewardToken);
+    }
+
+    // solhint-disable-next-line func-name-mixedcase
+    function __ERC721TokenStakingManager_init(
+        PoSValidatorManagerSettings calldata settings,
+        IERC721 stakingToken,
+        IERC20 rewardToken
+    ) internal onlyInitializing {
+        __POS_Validator_Manager_init(settings);
+        __ERC721TokenStakingManager_init_unchained(stakingToken, rewardToken);
+    }
+
+    // solhint-disable-next-line func-name-mixedcase
+    function __ERC721TokenStakingManager_init_unchained(
+        IERC721 stakingToken,
+        IERC20 rewardToken
+    ) internal onlyInitializing {
+        ERC721TokenStakingManagerStorage storage $ = _getERC721StakingManagerStorage();
+
+        if (address(stakingToken) == address(0)) {
+            revert InvalidTokenAddress(address(stakingToken));
+        }
+        if (address(rewardToken) == address(0)) {
+            revert InvalidRewardTokenAddress(address(rewardToken));
+        }
+
+        $._token = stakingToken;
+        $._rewardToken = rewardToken;
+    }
+
+    /**
+     * @notice See {IERC721TokenStakingManager-initializeValidatorRegistration}
+     */
+    function initializeValidatorRegistration(
+        ValidatorRegistrationInput calldata registrationInput,
+        uint16 delegationFeeBips,
+        uint64 minStakeDuration,
+        uint256 tokenId
+    ) external nonReentrant returns (bytes32 validationID) {
+        return _initializeValidatorRegistration(
+            registrationInput, delegationFeeBips, minStakeDuration, tokenId
+        );
+    }
+
+    /**
+     * @notice See {IERC721TokenStakingManager-initializeDelegatorRegistration}
+     */
+    function initializeDelegatorRegistration(
+        bytes32 validationID,
+        uint256 tokenId
+    ) external nonReentrant returns (bytes32) {
+        return _initializeDelegatorRegistration(validationID, _msgSender(), tokenId);
+    }
+
+    /**
+     * @notice Returns the ERC721 token being staked
+     */
+    function erc721() external view returns (IERC721) {
+        return _getERC721StakingManagerStorage()._token;
+    }
+
+    /**
+     * @notice Returns the ERC20 token used for rewards
+     */
+    function rewardToken() external view returns (IERC20) {
+        return _getERC721StakingManagerStorage()._rewardToken;
+    }
+
+    /**
+     * @notice See {PoSValidatorManager-_lock}
+     * Note: Must be guarded with reentrancy guard for safe transfer from.
+     */
+    function _lock(uint256 tokenId) internal virtual override returns (uint256) {
+        _getERC721StakingManagerStorage()._token.transferFrom(_msgSender(), address(this), tokenId);
+        return 1;
+    }
+
+    /**
+     * @notice See {PoSValidatorManager-_unlock}
+     * Note: Must be guarded with reentrancy guard for safe transfer.
+     */
+    function _unlock(address to, bytes32 id, bool isValidator) internal virtual override {
+        uint256[] memory nfts = isValidator ? getValidatorNfts(id) : getDelegatorNfts(id);
+        for (uint256 i = 0; i < nfts.length; i++) {
+            uint256 nftId = nfts[i];
+            _getERC721StakingManagerStorage()._token.safeTransferFrom(address(this), to, nftId);
+        }
+    }
+
+    /**
+     * @notice See {PoSValidatorManager-_reward}
+     * @dev Distributes ERC20 rewards to stakers
+     */
+    function _reward(address account, uint256 amount) internal virtual override {
+        ERC721TokenStakingManagerStorage storage $ = _getERC721StakingManagerStorage();
+        $._rewardToken.safeTransfer(account, amount);
+    }
+
+    /**
+     * @notice Allows the contract to receive reward tokens
+     * @dev Called by owner to fund rewards
+     * @param amount Amount of reward tokens to transfer to the contract
+     */
+    function fundRewards(uint256 amount) external onlyOperator {
+        ERC721TokenStakingManagerStorage storage $ = _getERC721StakingManagerStorage();
+        $._rewardToken.safeTransferFrom(_msgSender(), address(this), amount);
+    }
+
+    /**
+     * @notice Allows owner to recover excess reward tokens
+     * @param amount Amount of reward tokens to recover
+     */
+    function recoverRewardTokens(uint256 amount) external onlyOperator {
+        ERC721TokenStakingManagerStorage storage $ = _getERC721StakingManagerStorage();
+        $._rewardToken.safeTransfer(_msgSender(), amount);
+    }
+}
diff --git a/contracts/validator-manager/ExampleRewardCalculator.sol b/contracts/validator-manager/ExampleRewardCalculator.sol
@@ -6,7 +6,6 @@
 pragma solidity 0.8.25;
 
 import {IRewardCalculator} from "./interfaces/IRewardCalculator.sol";
-
 contract ExampleRewardCalculator is IRewardCalculator {
     uint256 public constant SECONDS_IN_YEAR = 31536000;
 
@@ -16,8 +15,11 @@ contract ExampleRewardCalculator is IRewardCalculator {
 
     uint64 public immutable rewardBasisPoints;
 
-    constructor(uint64 rewardBasisPoints_) {
+    uint64 public immutable decimals;
+
+    constructor(uint64 rewardBasisPoints_, uint64 decimals_) {
         rewardBasisPoints = rewardBasisPoints_;
+        decimals = decimals_;
     }
 
     /**
@@ -31,6 +33,7 @@ contract ExampleRewardCalculator is IRewardCalculator {
         uint64 stakingEndTime,
         uint64 uptimeSeconds
     ) external view returns (uint256) {
+
         // Equivalent to uptimeSeconds/(validator.endedAt - validator.startedAt) < UPTIME_REWARDS_THRESHOLD_PERCENTAGE/100
         // Rearranged to prevent integer division truncation.
         if (
@@ -40,7 +43,7 @@ contract ExampleRewardCalculator is IRewardCalculator {
             return 0;
         }
 
-        return (stakeAmount * rewardBasisPoints * (stakingEndTime - stakingStartTime))
+        return (stakeAmount * rewardBasisPoints * (stakingEndTime - stakingStartTime) * (10 ** decimals))
             / SECONDS_IN_YEAR / BIPS_CONVERSION_FACTOR;
     }
 }
diff --git a/contracts/validator-manager/NativeTokenStakingManager.sol b/contracts/validator-manager/NativeTokenStakingManager.sol
@@ -7,7 +7,9 @@ pragma solidity 0.8.25;
 
 import {PoSValidatorManager} from "./PoSValidatorManager.sol";
 import {PoSValidatorManagerSettings} from "./interfaces/IPoSValidatorManager.sol";
-import {ValidatorRegistrationInput} from "./interfaces/IValidatorManager.sol";
+import {
+    ValidatorRegistrationInput, ValidatorManagerStorage
+} from "./interfaces/IValidatorManager.sol";
 import {INativeTokenStakingManager} from "./interfaces/INativeTokenStakingManager.sol";
 import {INativeMinter} from
     "@avalabs/[email protected]/contracts/interfaces/INativeMinter.sol";
@@ -58,6 +60,17 @@ contract NativeTokenStakingManager is
     // solhint-disable-next-line func-name-mixedcase, no-empty-blocks
     function __NativeTokenStakingManager_init_unchained() internal onlyInitializing {}
 
+    function _addValidatorNft(bytes32 validationID, uint256 tokenId) internal override {}
+
+    function _addDelegatorNft(bytes32 delegationID, uint256 tokenId) internal override {}
+
+    function _deleteValidatorNft(
+        bytes32 validationID
+    ) internal override {}
+    function _deleteDelegatorNft(
+        bytes32 delegationID
+    ) internal override {}
+
     /**
      * @notice See {INativeTokenStakingManager-initializeValidatorRegistration}.
      */
@@ -93,7 +106,9 @@ contract NativeTokenStakingManager is
     /**
      * @notice See {PoSValidatorManager-_unlock}
      */
-    function _unlock(address to, uint256 value) internal virtual override {
+    function _unlock(address to, bytes32 id, bool isValidator) internal virtual override {
+        uint64 weight = isValidator ? getValidator(id).startingWeight : getDelegator(id).weight;
+        uint256 value = weightToValue(weight);
         payable(to).sendValue(value);
     }