Testing real-world contract upgrades

In this post, we’ll use Ganache to create a fork of a chain and set up a playground to test our contract upgrades using ZeppelinOS, before we actually execute them in production.

Deploying to production can be a stressful experience, especially if you don’t have a proper suite of tests to ensure that what you built will run smoothly. This encompasses not just testing the code itself, but also the process for upgrading the current system in production to the new one, including both code and state.

This becomes even more important in the context of smart contracts, where a single error can cost millions of dollars. ZeppelinOS helps you patch any errors by allowing you to upgrade your contracts to new versions that you may have tested extensively. But the question of how to test the upgrade itself remains.

With this in mind, we’ll set up a sample project to actually test in a local environment a ZeppelinOS upgrade for a contract on the Ethereum network. Let’s get to it.

Getting started

We’ll set up a sample ZeppelinOS project using zos 2.1.0 and install a few dependencies as well.

$ npm init -y
$ npm install zos@2.1.0 zos-lib@2.1.0 truffle@5.0.1
$ npx zos init sample-erc20

We’ll be using the upgradeable ERC20 contracts provided by the openzeppelin-eth package for this project.

$ npx zos link openzeppelin-eth
Installing openzeppelin-eth via npm...

For the sake of this example, we’ll create an ERC20 token contract that adds a transferMany function for sending funds to multiple recipients simultaneously:

// contracts/CustomERC20.sol
pragma solidity ^ 0.5 .0;

import "zos-lib/contracts/Initializable.sol";
import "openzeppelin-eth/contracts/token/ERC20/ERC20.sol";
import "openzeppelin-eth/contracts/token/ERC20/ERC20Detailed.sol";

contract CustomERC20 is Initializable, ERC20, ERC20Detailed {
  function initialize(
    string memory name, string memory symbol, uint8 decimals, uint256 initialSupply, address initialHolder
  ) public initializer {
    require(initialSupply > 0);
    ERC20Detailed.initialize(name, symbol, decimals);
    _mint(initialHolder, initialSupply);
  }

  function transferMany(address[] memory tos, uint256 value) public {
    for (uint256 i = 0; i < tos.length; i++) {
      _transfer(msg.sender, tos[i], value);
    }
  }
}

To build our contract, we’re extending from the ERC20 implementation provided by the openzeppelin-eth package. Note that we’re also using the Initializable contract from zos-lib to flag our contract as initializable (remember that upgradeable contracts use initializer functions instead of constructors).

Deploying the contract

We can now deploy and create an instance of our contract on a network. We’ll use Rinkeby for this sample, so we need to add an entry to our truffle config for that network. See this tutorial for more info, and make sure to unlock more than one account by adding the last two parameters to the HDWalletProvider.

networks: {
rinkeby: {
provider: new HDWalletProvider(mnemonic, 'https://ropsten.infura.io/v3/' + infuraToken, 0, 5)
network_id: 4
}
}

On the following commands, make sure to replace $ADMIN with the sender account you’ll be using, and $USER with another account that will hold the initial supply of tokens.

Also, if you are following this tutorial in a network that is not Rinkeby, Ropsten, Kovan, or Mainnet, make sure to add the flag --deploy-dependencies to the zos push command.

Remember to assign the $ADMIN an address that is not the first account from your mnemonic or on your node!! Otherwise, you will get a nasty “Cannot call fallback function from the proxy admin” error later. See here for more information on why this is needed. Note that this will no longer be required starting on version 2.2.

$ npx zos add CustomERC20
Compiling contracts with Truffle...
Adding CustomERC20

$ npx zos session --network rinkeby --from $ADMIN
Using network rinkeby, sender address $ADMIN by default.

$ npx zos push
Using session with network rinkeby, sender address $ADMIN
Compiling contracts with Truffle...
Validating contract CustomERC20
Uploading CustomERC20 contract as CustomERC20
Deploying logic contract for CustomERC20
Updated zos.rinkeby.json

$ npx zos create CustomERC20 --init --args "MyToken","MYT",8,"1000000000000","$USER"
Using session with network rinkeby, sender address $ADMIN
Creating proxy to logic contract and initializing by calling initialize with:
- name (string): "MyToken"
- symbol (string): "MYT"
- decimals (uint8): "8"
- initialSupply (uint256): "1000000000000"
- initialHolder (address): "$USER"
Instance created at $TOKEN
Updated zos.rinkeby.json

We now have our upgradeable contract deployed on Rinkeby. Let’s fire up a Truffle console via truffle console --network rinkeby and interact with it. Remember to replace $TOKEN with the address returned by the create command.

truffle(rinkeby)> const erc20 = await CustomERC20.at($TOKEN)

truffle(rinkeby)> erc20.balanceOf($USER).then(x => x.toNumber() / 1e8)
10000

truffle(rinkeby)> erc20.totalSupply().then(x => x.toNumber() / 1e8)
10000

truffle(rinkeby)> erc20.name()
“MyToken”

If you’re getting a VM revert when running these tests, or just empty return values, doublecheck that you didn’t set $ADMIN to be the first account on your node. If you did, you’ll need to perform these calls from another account. For example: erc20.balanceOf($USER, { from: $USER }). See here to learn more about this.

Everything should be running smoothly at this point. Feel free to run a few transactions as well to further test that the token is working as expected. All data from this deployment is stored in the zos.rinkeby.json file in the root of your project.

Modifying the contract

Let’s suppose that we want to make our token burnable. We can leverage OpenZeppelin’s ERC20 burnable contract to do this, just by extending from it on our contract.

// contracts/CustomERC20.sol
pragma solidity ^0.5.0;

import "zos-lib/contracts/Initializable.sol";
import "openzeppelin-eth/contracts/token/ERC20/ERC20.sol";
import "openzeppelin-eth/contracts/token/ERC20/ERC20Detailed.sol";
import "openzeppelin-eth/contracts/token/ERC20/ERC20Burnable.sol";

contract CustomERC20 is Initializable, ERC20, ERC20Burnable, ERC20Detailed {
 // ...
}

But before we actually push this change, let’s make sure we test that the upgrade will work as intended. To do this, we’ll start a new Ganache instance forking off from Rinkeby. This will set up a local environment with exactly the same state as the entire network where our current contract is running, which we will use for testing.

Setting up the testing environment

Install Ganache via npm install -g ganache-cli if needed, and start it using the following command, where $RINKEBYNODE should be the path to the Rinkeby node (such as Infura).

$ ganache-cli --fork $RINKEBYNODE --unlock "$USER" --unlock "$ADMIN" --port 9545 --networkId 1004 --deterministic

What we’ve done here is start a new chain with id 1004, forking off from Rinkeby, and unlock the USER and ADMIN accounts to use them freely within this Ganache instance, which is listening on port 9545.

The next step is to actually trick ZeppelinOS into thinking that the state on this new network is the same as the one on Rinkeby. To do this, we just need to copy the zos.rinkeby.json file to the one corresponding to a development network with id 1004.

$ cp zos.rinkeby.json zos.dev-1004.json

The last piece is to add a connection to our Ganache instance in our Truffle configuration file. We’ll name it rinkeby-test.

networks: {
 "rinkeby-test": {
   host: 'localhost',
   port: 9545,
   network_id: "1004",
   gasPrice: 10e9,
   gas: 50000
 }
}

We now have our environment ready and can test our upgrade. Try firing up a Truffle console at rinkeby-test, and repeat the queries we did before directly on Rinkeby. You’ll see that you get the same results without having needed to re-upload your contracts! This is because we’ve effectively forked off an existing chain, so we get to keep all the previous state from it on our new chain.

truffle(rinkeby-test)> const erc20 = await CustomERC20.at($TOKEN)
truffle(rinkeby-test)> erc20.balanceOf($USER).then(x => x.toNumber() / 1e8)
10000
truffle(rinkeby-test)> erc20.totalSupply().then(x => x.toNumber() / 1e8)
10000
truffle(rinkeby-test)> erc20.name()
“MyToken”

Upgrading the contract

Let’s start by switching the current zos session to one based on rinkeby-test, so we interact in our Ganache playground. Make sure you use the same admin address as before, which should be unlocked in the Ganache instance.

$ npx zos session --network rinkeby-test --from $ADMIN
Using network rinkeby-test, sender address $ADMIN by default.

We can now push the modified implementation contract to the network using zos push.

$ npx zos push
Using session with network rinkeby-test, sender address $ADMIN
Compiling contracts with Truffle...
Validating contract CustomERC20
- New variable 'uint256[50] ______gap' was inserted in contract ERC20Detailed in openzeppelin-eth/contracts/token/ERC20/ERC20Detailed.sol:1.a
You should only add new variables at the end of your contract.
See https://docs.zeppelinos.org/docs/writing_contracts.html#modifying-your-contracts for more info.

One or more contracts have validation errors.

Please review the items listed above and fix them, or run this command again with the --force option.

Dang! We’re getting a validation error! It seems that one of the changes we did on our contract altered the contract storage layout, which could potentially break our current contract instance. For the sake of this exercise, let’s pretend we don’t care about this warning for now (we’ll see what happens later!). We can use the --force flag to deploy anyway and update our contract instance to the new implementation.

$ npx zos push --force
Using session with network rinkeby-test, sender address $ADMIN
Compiling contracts with Truffle...
Validating contract CustomERC20
- New variable 'uint256[50] ______gap' was inserted in contract ERC20Detailed in openzeppelin-eth/contracts/token/ERC20/ERC20Detailed.sol:1. You should only add new variables at the end of your contract.
See https://docs.zeppelinos.org/docs/writing_contracts.html#modifying-your-contracts for more info.
Uploading CustomERC20 contract as CustomERC20

Deploying logic contract for CustomERC20

Updated zos.dev-1004.json
$ npx zos update CustomERC20
Using session with network rinkeby-test, sender address $ADMIN
Upgrading proxy to logic contract $IMPLEMENTATION

Instance at $TOKEN upgraded

Updated zos.dev-1004.json

Let’s start a Truffle console at Rinkeby test and see if everything worked as expected:

truffle(rinkeby)> const erc20 = await CustomERC20.at($TOKEN)

truffle(rinkeby)> erc20.balanceOf($USER).then(x => x.toNumber() / 1e8)
10000

truffle(rinkeby)> erc20.totalSupply().then(x => x.toNumber() / 1e8)
10000

truffle(rinkeby)> erc20.name()
""

truffle(rinkeby)> erc20.symbol()
""

truffle(rinkeby)> erc20.decimals().then(n => n.toNumber())
0

While the total supply and balances seem to be OK, the token’s details are gone! We’ve effectively flunked our upgrade.

As a side note, the reason behind this is related to the error displayed by the zos push operation. By inserting a new base contract in the middle of the inheritance chain (note that Burnable was added in between ERC20 and Detailed), we altered the storage layout. This moved the slots where the contract thinks that name, symbol, and decimals are stored. Instead, using Initializable, ERC20, ERC20Detailed, ERC20Burnable (in that order) keeps the same storage layout and works fine.

Luckily, we ran these tests on a disposable Ganache instance instead of on the actual network! This means that we can kill the Ganache instance, fix our contract, and try again by following the same steps.

Summing up

As we mentioned initially, upgrading a contract on mainnet can be a scary process. Regardless of the unit tests we run on our updated contract on our local development environments, we can’t catch any issues that may arise from the migration process itself, which requires carefully reproducing the state of the network where our contract is running.

Nevertheless, by setting up a disposable carbon copy of the blockchain on our workstations or CIs, we can work on a safe environment that has exactly the same characteristics as mainnet, allowing us to test in conditions as close as possible to those in the real world.