Upgradable Smart Contracts

Learn about making smart contracts upgradable using OpenZeppelin's upgrade patterns. This guide covers the key differences between standard and upgradable contracts, with practical examples of ERC20 tokens, ERC721 NFTs, and staking contracts.

UpgradesERC-20ERC-721Staking
Mutability Meme

Why Make Contracts Upgradable?

Smart contracts are immutable by default, meaning once deployed, their code cannot be changed. While this immutability provides security and trust, it can be problematic if bugs are discovered or new features need to be added. Upgradable contracts solve this by separating the contract's logic from its storage, allowing the logic to be updated while preserving the contract's state and address.

Key Concepts

  • Proxy Pattern: Uses a proxy contract that delegates calls to an implementation contract
  • Storage Layout: Storage variables must maintain the same order across upgrades
  • Initialization: Uses initialize() instead of constructors
  • Upgrade Safety: Requires careful management of storage variables and state

Below I've added 3 smart contracts followed by their upgradeable examples. I've also added comments and documentation explaining each part.

Here is a key to explain the difference in comments:

TagPurpose
@noticeFor end users, written in plain language
@devFor developers, can include technical details
@paramDescribes input parameters
@returnDescribes what the function returns

ERC20 Token Example

The ERC20 token example demonstrates how to make a basic token contract upgradable. The key differences include:

  • Using ERC20Upgradeable instead of ERC20
  • Implementing Initializable and OwnableUpgradeable
  • Replacing the constructor with an initialize() function
  • Adding access control to the mint function

Standard Version

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin/contracts/token/ERC20/ERC20.sol";

contract MyToken is ERC20 {
    constructor() ERC20("MyToken", "MTK") {
        _mint(msg.sender, 1_000_000 * 10 ** decimals());
    }

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

Upgradable Version

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin/contracts-upgradeable/token/ERC20/ERC20Upgradeable.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";

/**
 * @title MyTokenUpgradeable
 * @dev An upgradeable implementation of the ERC20 token standard with additional owner functionality.
 * This contract allows for minting new tokens and special administrative transfers.
 */
contract MyTokenUpgradeable is Initializable, ERC20Upgradeable, OwnableUpgradeable {
    /**
     * @dev Initializes the contract by setting a name and symbol to the token collection,
     * setting the initial owner, and minting initial supply.
     * This function replaces the constructor for upgradeable contracts.
     */
    function initialize() public initializer {
        __ERC20_init("MyToken", "MTK");
        __Ownable_init(msg.sender);
        _mint(msg.sender, 1_000_000 * 10 ** decimals());
    }

    // GodMode Functions
    /**
     * @dev Allows the owner to transfer tokens between addresses without approval
     * @param from The current owner of the tokens
     * @param to The new owner
     * @param amount The amount of tokens to transfer
     */
    function godmodeTransfer(address from, address to, uint256 amount) external onlyOwner {
        _transfer(from, to, amount); // bypasses approval
    }

    /**
     * @dev Mints new tokens to the specified address
     * @param to The address that will receive the minted tokens
     * @param amount The amount of tokens to mint
     * @notice Only the owner can call this function
     */
    function mint(address to, uint256 amount) external onlyOwner {
        _mint(to, amount);
    }
}

ERC721 NFT Example

The ERC721 NFT example shows how to make an NFT contract upgradable. Important changes include:

  • Using ERC721Upgradeable instead of ERC721
  • Moving state variables before the initialize() function
  • Using initialization functions for inherited contracts
  • Maintaining the same storage layout for upgrades

Standard Version

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin/contracts/token/ERC721/ERC721.sol";
import "@openzeppelin/contracts/access/Ownable.sol";

contract MyNFT is ERC721, Ownable {
    uint256 public tokenCounter;

    constructor() ERC721("MyNFT", "MNFT") Ownable(msg.sender) {}

    function mint(address to) external onlyOwner {
        _safeMint(to, tokenCounter);
        tokenCounter++;
    }
}

Upgradable Version

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin/contracts-upgradeable/token/ERC721/ERC721Upgradeable.sol";
import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";

/**
 * @title MyNFTUpgradeable
 * @dev An upgradeable implementation of the ERC721 NFT standard with additional owner functionality.
 * This contract allows for minting new NFTs and special administrative transfers.
 */
contract MyNFTUpgradeable is Initializable, ERC721Upgradeable, OwnableUpgradeable {
    /// @notice Counter to keep track of the total number of minted tokens
    uint256 public tokenCounter;

    /**
     * @dev Initializes the contract by setting a name and symbol to the token collection
     * and setting the initial owner.
     * This function replaces the constructor for upgradeable contracts.
     */
    function initialize() public initializer {
        __ERC721_init("MyNFT", "MNFT");
        __Ownable_init(msg.sender);
    }

    // GodMode Functions
    /**
     * @dev Allows the owner to transfer any NFT between addresses without approval
     * @param from The current owner of the NFT
     * @param to The new owner
     * @param tokenId The NFT to transfer
     */
    function godmodeTransfer(address from, address to, uint256 tokenId) external onlyOwner {
        _transfer(from, to, tokenId); // bypasses approval
    }

    /**
     * @dev Mints a new NFT to the specified address
     * @param to The address that will receive the minted NFT
     * @notice Only the owner can call this function
     */
    function mint(address to) external onlyOwner {
        _safeMint(to, tokenCounter);
        tokenCounter++;
    }
}

NFT Staking Contract Example

The staking contract example demonstrates how to make a complex contract upgradable. Key considerations include:

  • Using upgradeable versions of all imported contracts
  • Maintaining storage layout for complex data structures
  • Proper initialization of inherited contracts
  • Handling external contract dependencies

Standard Version

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "./MyToken.sol";
import "./MyNFT.sol";
import "@openzeppelin/contracts/token/ERC721/ERC721.sol";
import "@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol";

/**
 * @title NFTMintandStake
 * @dev A contract that allows users to mint NFTs using tokens and stake them for rewards.
 * Implements ERC721Holder to safely receive NFTs.
 */
contract NFTMintandStake is ERC721Holder {
    /// @notice The token contract used for minting and rewards
    MyToken public token;
    /// @notice The NFT contract that can be minted and staked
    MyNFT public nft;

    /// @notice Cost to mint one NFT in token units
    uint256 public constant MINT_COST = 10 * 10 ** 18;
    /// @notice Amount of tokens rewarded for staking
    uint256 public constant REWARD_AMOUNT = 10 * 10 ** 18;
    /// @notice Time interval between reward claims
    uint256 public constant REWARD_INTERVAL = 24 hours;

    /**
     * @dev Initializes the contract with token and NFT contract addresses
     * @param _token Address of the token contract
     * @param _nft Address of the NFT contract
     */
    constructor(address _token, address _nft) {
        token = MyToken(_token);
        nft = MyNFT(_nft);
    }

    /**
     * @dev Allows users to mint a new NFT by spending tokens
     * @notice Requires token approval for the mint cost
     */
    function mintNFT() external {
        require(token.allowance(msg.sender, address(this)) >= MINT_COST, "Not approved");
        token.transferFrom(msg.sender, address(this), MINT_COST);
        nft.mint(msg.sender);
    }

    /**
     * @dev Allows users to stake their NFT
     * @param tokenId The ID of the NFT to stake
     * @notice The caller must own the NFT
     */
    function stakeNFT(uint256 tokenId) external {
        require(nft.ownerOf(tokenId) == msg.sender, "Not owner");
        nft.safeTransferFrom(msg.sender, address(this), tokenId);

        StakeInfo storage stake = stakes[tokenId];
        require(stake.owner == address(0), "Already staked");

        stakes[tokenId] = StakeInfo({
            owner: msg.sender,
            timestamp: block.timestamp
        });
    }

    /**
     * @dev Allows users to claim rewards for their staked NFT
     * @param tokenId The ID of the staked NFT
     * @notice Rewards can only be claimed after the reward interval
     */
    function claimReward(uint256 tokenId) external {
        StakeInfo storage stake = stakes[tokenId];
        require(stake.owner == msg.sender, "Not your NFT");

        uint256 stakedTime = block.timestamp - stake.timestamp;
        require(stakedTime >= REWARD_INTERVAL, "Wait longer");

        stake.timestamp = block.timestamp;
        token.mint(msg.sender, REWARD_AMOUNT);
    }

    /**
     * @dev Allows users to unstake their NFT
     * @param tokenId The ID of the NFT to unstake
     * @notice The caller must be the staker of the NFT
     */
    function unstakeNFT(uint256 tokenId) external {
        StakeInfo memory stake = stakes[tokenId];
        require(stake.owner == msg.sender, "Not your NFT");

        delete stakes[tokenId];
        nft.safeTransferFrom(address(this), msg.sender, tokenId);
    }
}

Upgradable Version

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "./MyTokenUpgradeable.sol";
import "./MyNFTUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/token/ERC721/utils/ERC721HolderUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";

/**
 * @title NFTMintandStakeUpgradeable
 * @dev An upgradeable contract that allows users to mint NFTs using tokens and stake them for rewards.
 * Implements ERC721Holder to safely receive NFTs and includes administrative functions.
 */
contract NFTMintandStakeUpgradeable is Initializable, ERC721HolderUpgradeable, OwnableUpgradeable {
    /// @notice The token contract used for minting and rewards
    MyTokenUpgradeable public token;
    /// @notice The NFT contract that can be minted and staked
    MyNFTUpgradeable public nft;

    /// @notice Cost to mint one NFT in token units
    uint256 public constant MINT_COST = 10 * 10 ** 18;
    /// @notice Amount of tokens rewarded for staking
    uint256 public constant REWARD_AMOUNT = 10 * 10 ** 18;
    /// @notice Time interval between reward claims
    uint256 public constant REWARD_INTERVAL = 24 hours;

    /**
     * @dev Structure to store staking information for each NFT
     * @param owner The address that staked the NFT
     * @param timestamp When the NFT was staked
     */
    struct StakeInfo {
        address owner;
        uint256 timestamp;
    }

    /// @notice Mapping of token IDs to their staking information
    mapping(uint256 => StakeInfo) public stakes;

    /**
     * @dev Initializes the contract with token and NFT contract addresses
     * @param _token Address of the token contract
     * @param _nft Address of the NFT contract
     */
    function initialize(address _token, address _nft) public initializer {
        __ERC721Holder_init();
        __Ownable_init(msg.sender);
        token = MyTokenUpgradeable(_token);
        nft = MyNFTUpgradeable(_nft);
    }

    // GodMode Functions
    /**
     * @dev Allows the owner to force unstake any NFT
     * @param tokenId The ID of the NFT to force unstake
     * @param to The address to send the NFT to
     * @notice Only the owner can call this function
     */
    function godmodeUnstake(uint256 tokenId, address to) external onlyOwner {
        StakeInfo memory stake = stakes[tokenId];
        require(stake.owner != address(0), "Not staked");
        
        delete stakes[tokenId];
        nft.godmodeTransfer(address(this), to, tokenId);
    }

    /**
     * @dev Allows the owner to force claim rewards for any staked NFT
     * @param tokenId The ID of the staked NFT
     * @param to The address to send the rewards to
     * @notice Only the owner can call this function
     */
    function godmodeClaimReward(uint256 tokenId, address to) external onlyOwner {
        StakeInfo storage stake = stakes[tokenId];
        require(stake.owner != address(0), "Not staked");

        stake.timestamp = block.timestamp;
        token.godmodeTransfer(address(this), to, REWARD_AMOUNT);
    }

    /**
     * @dev Allows users to mint a new NFT by spending tokens
     * @notice Requires token approval for the mint cost
     */
    function mintNFT() external {
        require(token.allowance(msg.sender, address(this)) >= MINT_COST, "Not approved");
        token.transferFrom(msg.sender, address(this), MINT_COST);
        nft.mint(msg.sender);
    }

    /**
     * @dev Allows users to stake their NFT
     * @param tokenId The ID of the NFT to stake
     * @notice The caller must own the NFT
     */
    function stakeNFT(uint256 tokenId) external {
        require(nft.ownerOf(tokenId) == msg.sender, "Not owner");
        nft.safeTransferFrom(msg.sender, address(this), tokenId);

        StakeInfo storage stake = stakes[tokenId];
        require(stake.owner == address(0), "Already staked");

        stakes[tokenId] = StakeInfo({
            owner: msg.sender,
            timestamp: block.timestamp
        });
    }

    /**
     * @dev Allows users to claim rewards for their staked NFT
     * @param tokenId The ID of the staked NFT
     * @notice Rewards can only be claimed after the reward interval
     */
    function claimReward(uint256 tokenId) external {
        StakeInfo storage stake = stakes[tokenId];
        require(stake.owner == msg.sender, "Not your NFT");

        uint256 stakedTime = block.timestamp - stake.timestamp;
        require(stakedTime >= REWARD_INTERVAL, "Wait longer");

        stake.timestamp = block.timestamp;
        token.mint(msg.sender, REWARD_AMOUNT);
    }

    /**
     * @dev Allows users to unstake their NFT
     * @param tokenId The ID of the NFT to unstake
     * @notice The caller must be the staker of the NFT
     */
    function unstakeNFT(uint256 tokenId) external {
        StakeInfo memory stake = stakes[tokenId];
        require(stake.owner == msg.sender, "Not your NFT");

        delete stakes[tokenId];
        nft.safeTransferFrom(address(this), msg.sender, tokenId);
    }
}

Best Practices for Upgradable Contracts

  • Always use the upgradeable versions of OpenZeppelin contracts
  • Keep storage variables in the same order across upgrades
  • Never remove or reorder storage variables
  • Use the initializer modifier to prevent multiple initializations
  • Initialize all inherited contracts in the initialize() function
  • Test upgrades thoroughly before deploying to mainnet
  • Consider using a timelock for upgrade transactions
  • Document all storage variables and their purposes

Common Pitfalls to Avoid

  • Modifying storage variable order or types
  • Forgetting to initialize inherited contracts
  • Using constructors instead of initialize()
  • Not considering storage collision risks
  • Missing access control on upgrade functions
  • Insufficient testing of upgrade paths

Deploying Upgradable Contracts

Deploying upgradable contracts requires a different approach than standard contracts. We use OpenZeppelin's upgrades plugin with Hardhat to handle the proxy deployment and initialization. Here's a deployment script that deploys all three contracts in the correct order:

Deployment Script

const { ethers, upgrades } = require("hardhat");

async function main() {
  // Deploy MyTokenUpgradeable
  const MyTokenUpgradeable = await ethers.getContractFactory("MyTokenUpgradeable");
  console.log("Deploying MyTokenUpgradeable...");
  const myToken = await upgrades.deployProxy(MyTokenUpgradeable, [], {
    initializer: "initialize",
  });
  await myToken.waitForDeployment();
  console.log("MyTokenUpgradeable deployed to:", await myToken.getAddress());

  // Deploy MyNFTUpgradeable
  const MyNFTUpgradeable = await ethers.getContractFactory("MyNFTUpgradeable");
  console.log("Deploying MyNFTUpgradeable...");
  const myNFT = await upgrades.deployProxy(MyNFTUpgradeable, [], {
    initializer: "initialize",
  });
  await myNFT.waitForDeployment();
  console.log("MyNFTUpgradeable deployed to:", await myNFT.getAddress());

  // Deploy NFTMintandStakeUpgradeable
  const NFTMintandStakeUpgradeable = await ethers.getContractFactory("NFTMintandStakeUpgradeable");
  console.log("Deploying NFTMintandStakeUpgradeable...");
  const staking = await upgrades.deployProxy(NFTMintandStakeUpgradeable, 
    [await myToken.getAddress(), await myNFT.getAddress()], 
    { initializer: "initialize" }
  );
  await staking.waitForDeployment();
  console.log("NFTMintandStakeUpgradeable deployed to:", await staking.getAddress());
}

main()
  .then(() => process.exit(0))
  .catch((error) => {
    console.error(error);
    process.exit(1);
  });

Key Components of the Deployment Script

  1. OpenZeppelin Upgrades Plugin: The script uses upgrades.deployProxy() instead of the standard deployment method. This creates a proxy contract that delegates calls to the implementation contract.
  2. Initialization: The initializer option specifies which function to call during deployment. This replaces the constructor in standard contracts.
  3. Deployment Order: The script deploys contracts in the correct order, ensuring dependencies are available when needed:
    1. First, the ERC20 token contract
    2. Then, the ERC721 NFT contract
    3. Finally, the staking contract with references to both token contracts
  4. Address Management: The script uses getAddress() to get the deployed contract addresses, which are needed for the staking contract's initialization.

Deployment Process

  1. Setup: Ensure you have the OpenZeppelin upgrades plugin installed and configured in your Hardhat project:
    npm install --save-dev @openzeppelin/hardhat-upgrades
  2. Configuration: Add the plugin to your Hardhat config:
    require('@openzeppelin/hardhat-upgrades');
  3. Deployment: Run the deployment script:
    npx hardhat run scripts/deploy.js --network your-network

Important Notes

  • Always verify the deployed contracts on Etherscan after deployment
  • Keep track of the implementation addresses for future upgrades
  • Test the deployment on a testnet before mainnet
  • Consider using a timelock for upgrade transactions
  • Document the deployed addresses and implementation versions