Upgradeable smart contracts using the Diamond Standard

2020-05-28This post is over 2 years old and may now be out of date

(5 minute read)

In a recent post I talked about how we used the Proxy + Eternal Storage pattern to enable upgradeable smart contracts. In this pattern the Solidity fallback function is used in a Proxy contract to delegate all incoming functions to a separate implementation contract:

After successfully implementing this pattern we then later on hit the problem of our contract bytecode size exceeding the allowed limit of 24 KB, due to adding numerous new features to our implementation.

We decided to use the nested delegation pattern to split our proxy implementation into several sub-contracts:

Splitting the implementation into sub-implementations to avoid contract size limits

Though this solved the contract size limit problem, it was only a temporary fix - the Implementation contract had to manually delegate each method to the corresponding sub-implementation contract, meaning it still had to contain an entry for each method. So it would still eventually exceed the contract size limit as more methods got added.

Moreover, we had trouble getting the manual sub-delegation of read-only methods to work properly, and thus we could only place read-write methods (i.e. methods called via a transaction) into the sub-implementations. This prevented us from building cohesive sub-implementations whereby each sub-implementation would contain all the related methods for a given aspect of the system.

In short, we needed a better solution and so we decided to give the Diamond Standard a go.

The shiny new approach

The Diamond Standard is designed to allow contracts to be of any size by having functionality split up into an unlimited number of implementation contracts. It works around the above issues by keeping track of which implementation contract contains which function:

In the Diamond Standard each implementation contract is called a facet, like the facets of a diamond. In the proxy's fallback function it looks up the facet which implements the incoming function call from the mapping table that's in storage and then delegates the call to it:

fallback() external payable {
  DiamondStorage storage ds = diamondStorage();
  address facet = address(bytes20(ds.facets[msg.sig]));
  require(facet != address(0), "Function does not exist.");
    
  assembly {
    let ptr := mload(0x40)
    calldatacopy(ptr, 0, calldatasize())
    let result := delegatecall(gas(), facet, ptr, calldatasize(), 0, 0)
    let size := returndatasize()
    returndatacopy(ptr, 0, size)
    switch result
    case 0 {revert(ptr, size)}
    default {return (ptr, size)}
  }
}

The diamondCut() function is used to upgrade the mapping table. It takes a a bytes[] argument which specifies the method -> facet mappings to update the internal mapping table with. An example of its use can be found in the Diamond Standard reference implementation:

bytes[] memory diamondCut = new bytes[](3);

// Adding cut function
diamondCut[0] = abi.encodePacked(diamondFacet, Diamond.diamondCut.selector);

// Adding diamond loupe functions                
diamondCut[1] = abi.encodePacked(
    diamondLoupeFacet,
    DiamondLoupe.facetFunctionSelectors.selector,
    DiamondLoupe.facets.selector,
    DiamondLoupe.facetAddress.selector,
    DiamondLoupe.facetAddresses.selector            
);    

// Adding supportsInterface function
diamondCut[2] = abi.encodePacked(address(this), ERC165.supportsInterface.selector);

// execute cut function
bytes memory cutFunction = abi.encodeWithSelector(Diamond.diamondCut.selector, diamondCut);
(bool success,) = address(diamondFacet).delegatecall(cutFunction);
require(success, "Adding functions failed.");

If at any point in future one wishes to freeze the current implementation then this is just a matter of upgrading the diamondCut() method itself (since it too is contained within a facet).

The benefits of the Diamond Standard over previous approaches are numerous. But some of the primary benefits are:

Incremental upgrades, i.e. only upgrade one method at a time.
No need for sub-delegation, and thus needing to manage a contract tree hierarchy that comes with that.
Cohesive facets, i.e. all methods related to a particular aspect of your system can be placed in the same facet, enabling a clean software architecture.

Integrating into our platform

In order to integrate the Diamond Standard into our platform we did, however, need to make some adjustments. Primarily, we needed to ensure that only system admins could perform contract upgrades, which meant controlling access to the diamondCut() method. At the same time we still wanted the ability to change the upgrade access rules at any point in time as well as being able to freeze upgrades in future.

We also wanted to make doing upgrades as easy as possible. Specifically, each facet should be responsible for telling diamondCut() which methods it exposes. This way the proxy Diamond contract wouldn't even need to know about what methods its various facets expose, enabler looser coupling.

Finally, our existing contracts already used Eternal Storage and we wanted to keep using this despite some of the latest advances in Solidity storage handling that were now available.

With these requirements in mind, we first modified the Diamond Standard base storage contract to inherit from our EternalStorage base contract:

import "./EternalStorage.sol";

contract DiamondStorageBase is EternalStorage {
    struct DiamondStorage {
        // maps function selectors to the facets that execute the functions.
        // and maps the selectors to the slot in the selectorSlots array.
        // and maps the selectors to the position in the slot.
        // func selector => address facet, uint64 slotsIndex, uint64 slotIndex
        mapping(bytes4 => bytes32) facets;
        // array of slots of function selectors.
        // each slot holds 8 function selectors.
        mapping(uint => bytes32) selectorSlots;
        // uint128 numSelectorsInSlot, uint128 selectorSlotsLength
        // selectorSlotsLength is the number of 32-byte slots in selectorSlots.
        // selectorSlotLength is the number of selectors in the last slot of
        // selectorSlots.
        uint selectorSlotsLength;
        // tracking initialization state
        // we use this to know whether a call to diamondCut() is part of the initial
        // construction or a later "upgrade" call
        bool initialized;
    }
        
    function diamondStorage() internal pure returns(DiamondStorage storage ds) {
        // ds_slot = keccak256("diamond.standard.diamond.storage");
        assembly { ds_slot := 0xc8fcad8db84d3cc18b4c41d551ea0ee66dd599cde068d998e57d5e09332c131c }
    }
}

We were then able to codify the base DiamondProxy contract from which all of our upgradeable contracts would inherit:

import "./DiamondStorageBase.sol";
import "./DiamondCutter.sol";
import "./DiamondLoupeFacet.sol";
import "./IDiamondFacet.sol";
import "./IDiamondProxy.sol";

abstract contract DiamondProxy is DiamondStorageBase, IDiamondProxy {
  constructor () public {
    DiamondCutter diamondCutter = new DiamondCutter();
    dataAddress["diamondCutter"] = address(diamondCutter);
        
    DiamondLoupeFacet diamondLoupeFacet = new DiamondLoupeFacet();
        
    address[] memory facets = new address[](1);
    facets[0] = address(diamondLoupeFacet);
    _registerFacets(facets);
  }
    
  // IDiamondProxy
  function registerFacets (address[] memory _facets) public override {
    require(msg.sender == address(this), 'external caller not allowed');
    _registerFacets(_facets);
  }
    
  // Internal methods
  function _registerFacets (address[] memory _facets) internal {
    bytes[] memory changes = new bytes[](_facets.length);
        
    for (uint i = 0; i < _facets.length; i += 1) {
      IDiamondFacet f = IDiamondFacet(_facets[i]);
      bytes memory selectors = f.getSelectors();
      changes[i] = abi.encodePacked(_facets[i], selectors);
    }
        
    _cut(changes);
  }
    
  function _cut (bytes[] memory _changes) internal {
    bytes memory cutFunction = abi.encodeWithSelector(DiamondCutter.diamondCut.selector, _changes);
    (bool success,) = dataAddress["diamondCutter"].delegatecall(cutFunction);
    require(success, "Adding functions failed.");
  }
    
  // Finds facet for function that is called and executes the
  // function if it is found and returns any value.
  fallback() external payable {
    DiamondStorage storage ds = diamondStorage();
    address facet = address(bytes20(ds.facets[msg.sig]));
    require(facet != address(0), "Function does not exist.");
        
    assembly {
      let ptr := mload(0x40)
      calldatacopy(ptr, 0, calldatasize())
      let result := delegatecall(gas(), facet, ptr, calldatasize(), 0, 0)
      let size := returndatasize()
      returndatacopy(ptr, 0, size)
      switch result
      case 0 {revert(ptr, size)}
      default {return (ptr, size)}
    }
  }
    
  receive() external payable {
  }
}

The key point to notice in DiamondProxy is that the diamondCut() method is not externally callable. In fact, it's not even added as part of a facet but instead is hard-coded via a special DiamondCutter contract instance. This means it cannot be overridden in an upgrade.

We also added an internal method called _registerFacets(). This method iterates through the passed-in contract addresses and generates the method -> facet mappings to pass to diamondCut(). Note that it uses the IDiamondFacet interface to obtain this information from each facet:

interface IDiamondFacet {
  function getSelectors () external pure returns (bytes memory);
}

In our system every facet implement this interface and uses it to tell the proxy contract what methods it exposes. For example:

contract TestEntityFacet is IDiamondFacet {
  function getSelectors () public pure override returns (bytes memory) {
    return abi.encodePacked(
      TestEntityFacet.getNumPolicies.selector
    );
  }
    
  function getNumPolicies() public view override returns (uint256) {
    return 666;
  }
}

Having this interface means that when we wish to upgrade a proxy contract we can simply pass in a list of deployed facet addresses and the _registerFacets() method will iterate through them and obtain the method selectors to then call diamondCut() with. The proxy contract doesn't itself need to know what methods are exposed by each facet.

But how are upgrades accomplished if diamondCut() is no longer externally callable and _registerFacets() is internal?

Every proxy contract in our implementation has a facet that is specifically for handling upgrades. And these "upgrade facets" all inherit from an IDiamondUpgradeFacet base contract which implements the IDiamondFacet interface:

abstract contract IDiamondUpgradeFacet is IDiamondFacet {
  // IDiamondFacet
  function getSelectors () public pure override returns (bytes memory) {
    return abi.encodePacked(
      IDiamondUpgradeFacet.upgrade.selector
    );
  }
    
  // methods
  function upgrade (address[] memory _facets) public virtual;
}

For example, our insurance policy contracts have a PolicyUpgradeFacet that looks like this (simplified form):

contract PolicyUpgradeFacet is EternalStorage, IDiamondUpgradeFacet {
  modifier assertIsAdmin () { /* access control logic here */ }

  // IDiamondUpgradeFacet
  function upgrade (address[] memory _facets) public override assertIsAdmin {
    IDiamondProxy(address(this)).registerFacets(_facets);
  }
}

Notice that the upgrade() method calls through to registerFacets() in the proxy contract. Let's take a look at that method:

function registerFacets (address[] memory _facets) public override {
    require(msg.sender == address(this), 'external caller not allowed');
  _registerFacets(_facets);
}

So registerFacets() is a public method which calls through to the internal _registerFacets() method to do the actual upgrade. But it first checks to see that the caller is the current contract itself. Thus, the registerFacets() method cannot be called directly by an external caller - it must be called by the proxy contract itself. Which means that in order to upgrade the proxy contract you have to call the IDiamondUpgradeFacet.upgrade() method. And since that method is in a facet, we can change how we control access to it at any time. Or even disable it and thus prevent future upgrades.

The final piece of the puzzle is how we set the initial facets for the Diamond proxy. Taking our policy insurance contracts as an example, here is how that code looks (simplified form):

import "./base/DiamondProxy.sol";

contract Policy is DiamondProxy {
  constructor (address[] memory _initialFacets) DiamondProxy() public {
    // set implementations
    _registerFacets(_initialFacets);
  }
}

Thus, on deployment we call the internal _registerFacets() method directly to setup the initial method -> facet mapping table. And we will hopefully have passed in the address of an IDiamondUpgradeFacet facet to enable future upgrades. Subsequent upgrades will then involve calling IDiamondUpgradeFacet.upgrade() on the proxy, which will then call registerFacets(), which will then call _registerFacets() to do the upgrade.

Further thoughts

Once we were able to modify the Diamond Standard to enable access control to upgrades and looser coupling between proxies and facets we were very quickly able to re-implement our contract code using it. And we've since then already performed an upgrade across Mainnet-deployed contracts to fix a bug.

You can view all of our code at https://github.com/nayms/contracts.