Section 14: Understanding SmallProxy Contract Behavior with ImplementationA

[vascofcd]

Hello everyone,

I'm currently working through the lesson 14 about Foundry Upgrades and decided to do some experiments with the SmallProxy contract. Below is the code for the SmallProxy and ImplementationA contracts (it was modified!):

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

import "@openzeppelin/contracts/proxy/Proxy.sol";

contract SmallProxy is Proxy {
    address public admin; //added this

    bytes32 private constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;

    //added this
    function setSomeAddress(address newAdmin) external {
        admin = newAdmin;
    }

    function setImplementation(address newImplementation) public {
        assembly {
            sstore(_IMPLEMENTATION_SLOT, newImplementation)
        }
    }

    function _implementation() internal view override returns (address implementationAddress) {
        assembly {
            implementationAddress := sload(_IMPLEMENTATION_SLOT)
        }
    }

    // helper function
    function getDataToTransact(uint256 numberToUpdate) public pure returns (bytes memory) {
        return abi.encodeWithSignature("setValue(uint256)", numberToUpdate);
    }

    function readStorage() public view returns (uint256 valueAtStorageSlotZero) {
        assembly {
            valueAtStorageSlotZero := sload(0)
        }
    }
}

contract ImplementationA {
    address public owner;

    function setOwner() public {
        owner = msg.sender;
    }
}

Here is the unit test I wrote for these contracts:

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

import {Test, console} from "forge-std/Test.sol";
import {ImplementationA, SmallProxy} from "../../src/24-PuzzleWallet/ProxyExample/SmallProxy.sol";

contract SmallProxyTest is Test {
    ImplementationA implementationA;
    SmallProxy smallProxy;

    address owner = makeAddr("owner");
    address attacker = makeAddr("attacker");

    function setUp() public {
        vm.startPrank(owner, owner);

        // init impl contract
        implementationA = new ImplementationA();

        // setup proxy
        smallProxy = new SmallProxy();
        smallProxy.setImplementation(address(implementationA));
        implementationA = ImplementationA(address(smallProxy));
        
        vm.stopPrank();
    
        assertEq(address(0), smallProxy.admin()); // passed
        assertEq(address(0), implementationA.owner()); // passed
    }

    function testStorageLayout() public {
        // part 1 owner to attacker
        vm.startPrank(attacker, attacker);
        smallProxy.setSomeAddress(attacker);
        vm.stopPrank();

        assertEq(attacker, smallProxy.admin()); // passed
        assertEq(attacker, implementationA.owner()); // passed

        // part 2 attacker to owner
        vm.startPrank(owner, owner);
        implementationA.setOwner();
        vm.stopPrank();

        assertEq(owner, smallProxy.admin()); // passed
        assertEq(owner, implementationA.owner()); // passed
    }
}

First, I'm assuming that ImplementationA and SmallProxyContract are two independent contracts, each one with its own storage layout. I'm puzzled by the behavior of:

Changing the Admin Address

• SmallProxy::setSomeAddress is supposed to set a new address for the admin variable in slot 0 of SmallProxy.
• However, implementationA.owner is also changed to the attacker’s address, despite no delegatecall being present in SmallProxy::setSomeAddress.
• I expected only SmallProxy.admin to change, but both are being updated. Why is this happening?

Changing the Owner Address

• When calling implementationA.setOwner(), both admin and owner are updated to the owner's address.
• Again, there’s no delegatecall in this function, so why are both variables being updated?

Now this contract utilizes Proxy from OpenZeppelin, so my guess is that something is happening behind the scenes with the fallback and delegatecall function that I'm not able to understand.

2. Can someone provide a mental model to understand this behavior? What is really happening step by step?
3. Are there any resources or documentation that could help clarify this?

Thank you in advance!

[EngrPips]

Interesting experiment here i must say.

[vascofcd]

I just posted this into the wrong course repo, Ill put in foundry-full course and delete this one
edit: I cant delete :s

[EngrPips]

Yeah, you should do that

[vascofcd]

I don't think I have rights to delete this post. I dont see the delete option

[EngrPips]

Maybe just close it .

[Chinwuba22]

Yes, SmallProxyContract and ImplementationA are two different contracts.

Then, yes when SmallProxy delegateCalls ImplementationA, ImplementationA's code is expected to run using SmallProxys storage.

Also, when smallProxy.setSomeAddress(attacker); is called it is not going to make changes to implementationA.owner()), the reaason your test is passing is because you added implementationA = ImplementationA(address(smallProxy)); in your test setup which is making implementationA contract inherits the function of SmallProxy which ideally is not the case.

So, the reason your test is passing is because of implementationA = ImplementationA(address(smallProxy)); which I doubt is doing what you intend it to do.