Making Smart Contracts Upgradeable: Fixing Locked Funds and Bugs

Introduction

In the blockchain world, smart contracts are immutable. While this immutability ensures security, it creates challenges when bugs or vulnerabilities arise. Many smart contracts have encountered these issues:

• Locked funds: Bugs that prevent users from withdrawing.
• Design flaws: Ether mistakenly sent to the contract becomes irretrievable.

In this article, we’ll address these problems by demonstrating how to make smart contracts upgradeable, allowing flexibility and recoverability.

Technical Stack

For this guide, we’ll use:

• Solidity: Version 0.8.26 for smart contract development.
• Hardhat: A robust development and testing environment for Ethereum.

Demonstration

We will:

1. Deploy a non-upgradable smart contract and highlight its limitations.
2. Explore two key issues:
   • Logic flaws causing transaction failures.
   • Funds being locked in the contract with no way to withdraw them.
3. Fix these issues by upgrading the contract.

The Problematic Contract

The smart contract in question is a token bulk sender that allows users to airdrop tokens to multiple recipients while optimizing gas fees. However, it contains two critical issues:

1. Incorrect Ether handling in the onlyAllowedAccount modifier.

   • When non-VIP users interact with the contract, the modifier transfers the entire msg.value (sent Ether) to the contract owner instead of just the transaction fee (txFee). This causes transaction failures.

2. Inability to withdraw mistakenly sent Ether.

   • If a user accidentally sends Ether directly to the contract or the DApp sends excess Ether, the funds become permanently stuck.

Here’s the original code:

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

import "@openzeppelin/contracts/access/Ownable.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC1155/IERC1155.sol";
import "@openzeppelin/contracts/token/ERC721/IERC721.sol";
import "./IBulkSender.sol";

contract BulkSender is Ownable, IBulkSender {
    address public receiverAddress;

    uint public txFee = 0.007 ether;
    uint public VIPFee = 0.1 ether;

    mapping(address => bool) private vipList;

    modifier onlyAllowedAccount() {
        require(isVIP(msg.sender) || msg.value >= txFee, NotAllowedAccount());
        if (!isVIP(msg.sender)) {
            payable(receiverAddress).transfer(msg.value); // Issue: Transfers entire msg.value
        }
        _;
    }

    constructor(address _receiverAddress) Ownable(msg.sender) {
        receiverAddress = _receiverAddress;
    }

    function registerVIP() public payable {
        require(msg.value >= VIPFee, InsufficientFunds(msg.value, VIPFee));
        require(!isVIP(msg.sender), AlreadyVIP());
        payable(receiverAddress).transfer(msg.value);
        vipList[msg.sender] = true;
        emit LogVIPRegistered(msg.sender, msg.value);
    }
}

Full Code source here

Issue 1: Incorrect Modifier Logic

The onlyAllowedAccount modifier transfers all Ether sent (msg.value) to the receiver address. This can cause transaction failures when the intended payment is only the transaction fee (txFee).

Original Code

modifier onlyAllowedAccount() {
    require(isVIP(msg.sender) || msg.value >= txFee, NotAllowedAccount());
    if (!isVIP(msg.sender)) {
        payable(receiverAddress).transfer(msg.value); // Issue: Transfers entire msg.value
    }
    _;
}

Corrected Version of the Modifier

The fix ensures that only the txFee is transferred to the receiver address instead of the full msg.value.

modifier onlyAllowedAccount() {
    require(isVIP(msg.sender) || msg.value >= txFee, NotAllowedAccount());
    if (!isVIP(msg.sender)) {
        payable(receiverAddress).transfer(txFee); // Fix: Transfer only the transaction fee
    }
    _;
}

Issue 2: Stuck Ether in the Contract

Another significant issue arises when Ether is accidentally sent to the contract, either directly or via a bug in the DApp. Since the contract lacks a withdrawal mechanism, these funds become permanently stuck.

Example

Here’s a screenshot showing Ether stuck in the contract:

Solution: Upgradeable Smart Contracts

To resolve these issues, we make the contract upgradeable using the UUPS proxy pattern.

What is the UUPS Proxy Pattern?

The UUPS proxy pattern separates the storage (managed by a proxy contract) from the implementation (logic contract). This allows developers to deploy a new implementation and link it to the proxy, effectively upgrading the contract without losing stored data.

How It Works

1. Proxy Contract: Handles all user interactions and stores the state variables.
2. Implementation Contract: Contains the business logic. This can be replaced with an upgraded version when necessary.

For more details, refer to the OpenZeppelin upgradeable contracts documentation.

UUPS Proxy

In this example, we will use the UUPS (Universal Upgradeable Proxy Standard) proxy, which integrates the upgrade logic directly into the implementation contract. This approach simplifies upgrades by allowing the implementation contract to manage its own upgrade process. Below are the details:

The proxy contract acts as a mediator, delegating all function calls to the implementation contract. Crucially, the proxy can modify its own storage based on the logic defined in the implementation. This enables seamless upgrades while preserving the proxy's state.

Refer to the diagram below for a high-level overview of how the proxy interacts with the implementation and manages storage.

Make BulkSender.sol contracts upgreadable

To upgrade the contract as we use @openzeppelin/contracts we have to install @openzeppelin/contracts-upgradeable, as below:

$ npm install @openzeppelin/contracts-upgradeable

then, replace the variables storage by a struct as below:

    /// @custom:storage-location bulksendtokens.xyz.bulksender.storage.bulksender
    struct BulkSenderStorage {
        address _receiverAddress;
        uint _txFee;
        uint _vipFee;
        mapping(address => bool) _vipList;
    }

    // keccak256(abi.encode(uint256(keccak256("bulksendtokens.xyz.bulksender.storage.bulksender")) - 1)) & ~bytes32(uint256(0xff))
    bytes32 private constant BulkSenderStorageLocation =
        0xa9b8ea93cd1a4e28b0276278267515f30a34f7de34d3bc6de92b1e97a9a6b700;

Create a function to get the current storage:

    function _getBulkSenderStorage()
        private
        pure
        returns (BulkSenderStorage storage $)
    {
        assembly {
            $.slot := BulkSenderStorageLocation
        }
    }

Then all variable must be initialized in storage instead of constructor, as below:

    function initialize(address receiverAddress) public initializer {
        __Ownable_init(msg.sender);
        __UUPSUpgradeable_init();
        BulkSenderStorage storage $ = _getBulkSenderStorage();
        $._receiverAddress = receiverAddress;
        $._txFee = 0.007 ether;
        $._vipFee = 0.1 ether;
    }

A hook should be added to restrict who can upgrade the contract or to implement custom upgrade logic. For example, you can ensure that only the contract owner is allowed to upgrade the contract.

    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}

We can deploy both the proxy and the implementation using the Hardhat script provided below:

var pkg = require('hardhat');
const { upgrades, ethers } = pkg;

async function main() {
  const BulkSender = await ethers.getContractFactory('BulkSender');

  const bulkSender = await upgrades.deployProxy(BulkSender, [process.env.OWNER_ADDRESS]);
  await bulkSender.waitForDeployment();
  console.log("Box deployed to:", await bulkSender.getAddress());
}

main();

You should save the logged address of the proxy, it will be needed in next steps:

Upgrade the bulksender.sol

We can now update the implementation by adding a function to unlock locked ethers and fixing the issue in the OnlyAllowedAccount modifier.

  function withdrawEther() external onlyOwner {
        require(address(this).balance > 0, "No Ether to withdraw");

        // Transfer all Ether to the owner
        payable(owner).transfer(address(this).balance);
    }

export the proxy address, as below:

export PROXY_ADDRESS=YOUR_PROXY_ADDRESS

Now that the contract has been updated, we can upgrade the implementation using the Hardhat scripts provided below:

 const BulkSenderV2= await ethers.getContractFactory("BulkSenderV2");
  const proxy = await upgrades.upgradeProxy(PROXY_ADDRESS, BulkSenderV2);
  console.log("BulkSender upgraded");

Our contracts are now fixed, and we have maintained the same address for all customers. Additionally, we can withdraw the locked funds.

Conclusion

By making the contract upgradeable, we’ve resolved the two primary issues:

1. Transaction Failures: Fixed the onlyAllowedAccount modifier to deduct only the required txFee.
2. Stuck Funds: Enabled a mechanism to introduce withdrawal functions or other fixes in future versions.

Upgrading contracts ensures flexibility and safeguards against unforeseen issues. As the blockchain ecosystem evolves, adopting patterns like the UUPS proxy is essential for building robust and user-friendly DApps.

Open Code Sources:

No Upgreadbale Bulk Sender Contracts

Upgreadbale Bulk Sender Contracts