Transparent Upgradable Proxy

A Transparent Upgradable Proxy using assembly to return from fallback and to write to any storage slot
Explore the docs »

View Demo · Report Bug · Request Feature

Table of Contents

1. About The Project
   • Built With
2. Getting Started
   • Prerequisites
   • Installation
3. Usage
4. Roadmap
5. Contributing
6. License
7. Contact
8. Acknowledgments

About The Project

(back to top)

Built With

• 
• 

(back to top)

Getting Started

To get a local copy up and running follow these simple example steps.

Getting Started

To get a local copy up and running follow these simple example steps.

Prerequisites

• npm

  npm install npm@latest -g

• hardhat

  npm install --save-dev hardhat

  npm install @nomiclabs/hardhat-ethers @nomiclabs/hardhat-waffle

  run:

  npx hardhat

Installation

1. Clone the repo

   git clone https://github.com/Aboudoc/Transparent-Upgradable-Proxy-assembly.git

2. Install NPM packages

   npm install

(back to top)

Transparent Uppgradable Proxy

This project is focusing on assembly

• Check this repository to learn how to upgrade using hardhat-deploy following the Transparent Upgradable Proxy pattern from Open Zeppelin

• Check this repository to learn how to write to any storage slot, read any slot using StorageSlot library from Open Zeppelin

• You can find some proxy vulnerabilities in this repository. The first exploit is Wallet Hijack and the second one is Misaligned Storage

In this project, we built our own library

To build our upgradable proxy, we'll try to follow these steps

1. Intro: the wrong way to implement proxy
2. Return data from fallback
3. Storage for implementation and admin
4. Separate user / admin interfaces
5. Proxy admin contract

Intro: the wrong way to implement proxy

Let's consider the following upgradable contract with two separate implementations: CounterV1 and CounterV2:

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

contract BuggyProxy {
    address public  implementation;
    address public admin;

    constructor() {
        admin = msg.sender;
    }

    function _delegate() private  {
        (bool ok, bytes memory res) = implementation.delegatecall(msg.data);
        require(ok);
    }

    receive() external payable{
        _delegate();
    }

    fallback() external payable {
        _delegate();
    }

    function upgradeTo(address _implementation) external {
        require(msg.sender == admin);
        implementation = _implementation;
    }
}

contract CounterV1 {
    address public implementation;
    address public admin;
    uint public count;

    function inc() external {
        count += 1;
    }
}

contract CounterV2 {
    address public implementation;
    address public admin;
    uint public count;

    function inc() external {
        count += 1;
    }

    function dec() external {
        count -= 1;
    }
}

Two problems to fix:

1. All of the implementation contract must have the same storage layout as the proxy contract
2. The fallbackcan not return any data, so we can not get the count of the counter.

Return data from fallback

We can increment the count using inc function, but if you try to get the count, it will return 0

fallback function and receivefunction both calls an internal function called _delegate

We'll be modifying this function to call delegatecall but after it will return the data even though the function signature does say that it return any data.

This will allow us to get the count from the implementation

To return data, we'll need to use assembly

  /**
   * @dev Delegates execution to an implementation contract.
   * This is a low level function that doesn't return to its internal call site.
   * It will return to the external caller whatever the implementation returns.
   * @param implementation Address to delegate.
   */
  function _delegate(address implementation) internal {
    assembly {
      // Copy msg.data. We take full control of memory in this inline assembly
      // block because it will not return to Solidity code. We overwrite the
      // Solidity scratch pad at memory position 0.
      calldatacopy(0, 0, calldatasize())

      // Call the implementation.
      // out and outsize are 0 because we don't know the size yet.
      let result := delegatecall(gas(), implementation, 0, calldatasize(), 0, 0)

      // Copy the returned data.
      returndatacopy(0, 0, returndatasize())

      switch result
      // delegatecall returns 0 on error.
      case 0 { revert(0, returndatasize()) }
      default { return(0, returndatasize()) }
    }
  }

This code comes from the following repo from Open Zeppelin

Notes that if you compile you'll get the following error : only local variables are supported

so we pass in the storage variable _implementation as argument:

function _delegate(address _implementation) internal {}

Then, we change the state variable to a local:

let result := delegatecall(gas(), _implementation, 0, calldatasize(), 0, 0)

Finally, pass in the state variable to fallback and receive:

    fallback() external payable {
        _delegate(_implementation);
    }

    receive() external payable {
        _delegate(_implementation);
    }

The basic idea for the code inside assembly is that we're gonna be copying the data, then manually delegatecall, and once we call delegatecall, we have our data stored in returned data. We'll copy this returned data into memory, then manually return it.

Let's see in details what the code inside assembly does

calldatacopy(t, f, s)

calldatacopy copy the calldata at memory 0 (t) starting from the calldata from 0 (f) to calldatasize

calldatacopy(0, 0, calldatasize())

Basically we're copying all of the calldata ounto memory at 0th position

delegatecall(g, a, in, insize, out, outsize)

g means we are forwarding all of the gas a is the address of the implementation to execut delegatecall on in and insize means the data is stored inside the memory from 0 to datasize. outand outside says to store the result of delegatecall to memory from out, to output size. But since we don't know the size of the output before delegating, we are simply ignoring the output, saying 0 and 0.

let result := delegatecall(gas(), implementation, 0, calldatasize(), 0, 0)

We were ignoring the output, next we will handle it by copying the output (data returned)

returndatacopy(t, f, s)

returndatacopy copies s bytes from returndata at position f to memory at position t

returndatacopy(0, 0, returndatasize())

Basically we're copying the data that was returned to memory 0 (t) starting at the 0th position (f) of the memory, and the size of the data to copy is stored in the returndatasize (returndatasize() returns the size of the last returndata)

The last step is to handle wether the delegatecall was successfull or not

switch result
            case 0 {
                revert(0, returndatasize())
            }
            default {
                return(0, returndatasize())
            }

If the result is 0, this mean there was an error, so we will revert and return the all of the output that was returned from delegatecall, from 0 to returndatasize

If the result is not equal to 0, we return the data stored in memory (that was copied using returndatacopy) from 0 to returndatasize

Storage for implementation and admin

Our goal is to store the address of the implementation and admin somewhere else, besides the 0th and the 1st slot

First of all, remove the address of the implementation and the admin from CounterV1 and `CounterV2``

contract CounterV1 {
    uint public count;

    function inc() external {
        count += 1;
    }
}
contract CounterV2 {
    uint public count;

    function inc() external {
        count += 1;
    }

    function dec() external {
        count -= 1;
    }
}

Now we need to write to any storage slot. We'll do it by creating a library StorageSlot:

library StorageSlot {
    struct AddressSlot {
        address value;
    }

    function getAddressSlot(bytes32 slot) internal pure returns (AddressSlot storage r) {
        assembly {
            r.slot := slot
        }
    }
}

Keep in mind that the storage of a solidity smart contract is an array of length 2^256, and to each slot we can store up to 32 bytes

To use this trick we can not simply pass in an address, we'll need to wrap our address in a struct

Next, we write a function to get the pointer of the storage, taking a single input that will specify the pointer that we want to get (in bytes32). This function will return the pointer to the address slot

Basically, this function will return the pointer to the storage r, located at slot from the input

To do this, we will use assembly

assembly {
    r.slot := slot
}

We can use StorageSlot library to get an address stored at any slot

StorageSlot.getAddressSlot(SLOT).value

This will return a pointer that is storing SLOT address struct

We can use StorageSlot library to store an address to any slot

StorageSlot.getAddressSlot(SLOT).value = _addr

Basically it says to get the storage pointer at the slot from the input

By using StorageSlot library, our proxy is not "Buggy" anymore. We'll be following the Open Zeppelin Transparent Upgradable Proxy contract, so we store implementation and admin in a special place

bytes32 private constant IMPLEMENTATION_SLOT =
        bytes32(uint256(keccak256("eip1967.proxy.implementation")) - 1);

bytes32 private constant ADMIN_SLOT = bytes32(uint256(keccak256("eip1967.proxy.admin")) - 1);

We're substracting 1 by casting the hash into uint to make some king of Hash Collision Attack difficult to pull off because we don't know the pre-image of the hash

We'll need a getter and a setter for these! (check code from line 34 to 50) and use these in constructor(), fallback(), receive() and upgradeTo()

Then, write public functions to get the admin, and get the implementation: admin() using _getAdmin() and implementation() using _getImplementation

Separate user / admin interfaces

If the caller is an admin, then we'll let him execute some function on the proxy contract, otherwise, if the caller is a user, we'll forward all of the calls to the implementation

Notice that we have some public functions in our Proxy contract (IMPLEMENTATION_SLOT, ADMIN_SLOT, upgradeTo and admin and implementation)

The problem is if we were to had these function also inside the implementation contract, then even though we might want to call the function inside the implementation contract, these functions inside the Proxy contract will always be executed (even if it's not the same function, as long as the function selectors are the same)

Inside the Proxy contract, we'll make all the public functions private, but we'll keep upgradeTo public since since we need, so we keep it external and create isAdmin modifier using _getAdmin internal function. Remember, admin is no longer a state variable, it's stored in the admin slot

In the modifier, if msg.sender is not admin, we'll forward the request to fallback

modifier ifAdmin() {
    if (msg.sender == _getAdmin()) {
        _;
    } else {
        _fallback();
    }
}

Notice that fallback function is external, so we can not call it directly, instead we use an intenal function named _fallback

function _fallback() private {
    _delegate(_getImplementation());
}

We also use this internal function inside fallback and receive

Then, we expose upgradeTo() only if msg.sender is admin

function upgradeTo(address _implementation) external ifAdmin {
    _setImplementation(_implementation);
}

Add modifier also for admin() function so we can keep it public

Proxy admin contract

Finally, we'll write a ProxyAdmin contract, and this contract will be the admin of the Proxy contract

The problem is that since the admin() and the implementation() also exist inside the Proxy contract, if the admin wants to call one of these function inside the implementation contract, we wouldn't be able to call it

So 1st thing that we need into the Proxy contract is to set the admin of the contract

function changeAdmin(address _admin) external ifAdmin {
    _setAdmin(_admin);
}

Now we'll write the ProxyAdmin contract, and when we deploy both, the Proxy contract and the ProxyAdmin contract, we'll call changeAdmin() passing-in the address of the ProxyAdmin contract

Note1 Inside of changeProxyAdmin, proxy address passed in as argument is marqued as payable because the Proxy contract has a fallback and a receive

function changeProxyAdmin(address payable proxy, address _admin) external onlyOwner {
    TransparentUpgradeableProxy(proxy).changeAdmin(_admin);
}

Note2 We wrote read-only function to read the address of the implementation and the address of the admin because in Proxy contract admin() and implementation() are not read-only functions. We had to remore the view declaration since we have a ifAdmin modifier which have the potential to call the fallback

The way we'll call a function that is not read-only: staticcall to make it read-only function

staticcall is like call except that it does not write anything into the blockchain. Note that we're passing-in empty parenthesis inside abi.encodeCall because there is no input to pass-in inside the function admin()

Finally, we know that when call the admin() it will return an address, so we abi.decode the response into address

    function getProxyAdmin(address proxy) external view returns (address) {

        (bool ok, bytes memory res) = proxy.staticcall(
            abi.encodeCall(TransparentUpgradeableProxy.admin, ())
        );
        require(ok, "call failed");

        return abi.decode(res, (address));
    }

We did something similar to get the address of the implementation inside getProxyImplementation()

Further reading

(...soon)

Sources

(back to top)

Roadmap

• Further reading

See the open issues for a full list of proposed features (and known issues).

(back to top)

Contributing

Contributions are what make the open source community such an amazing place to learn, inspire, and create. Any contributions you make are greatly appreciated.

If you have a suggestion that would make this better, please fork the repo and create a pull request. You can also simply open an issue with the tag "enhancement". Don't forget to give the project a star! Thanks again!

1. Fork the Project
2. Create your Feature Branch (git checkout -b feature/AmazingFeature)
3. Commit your Changes (git commit -m 'Add some AmazingFeature')
4. Push to the Branch (git push origin feature/AmazingFeature)
5. Open a Pull Request

(back to top)

License

Distributed under the MIT License. See LICENSE.txt for more information.

(back to top)

Contact

Reda Aboutika - @twitter - reda.aboutika@gmail.com

Project Link: https://github.com/Aboudoc/Transparent-Upgradable-Proxy-assembly.git

(back to top)

Acknowledgments

• Smart Contract Engineer
• Open Zeppelin

(back to top)