[Bug]: Delegate-called logic contracts not returning proper `contract::address`

[0xNeshi]

What happened?

The following test case should pass, but it fails at the first assertion:

#[cfg(test)]
mod tests {
    use alloy_sol_macro::sol;
    use alloy_sol_types::SolCall;
    use motsu::prelude::*;
    use stylus_sdk::{alloy_primitives::Address, prelude::*, ArbResult};

    use super::*;
    use crate::proxy::{self, erc1967::Erc1967Proxy, IProxy};

    #[storage]
    #[entrypoint]
    struct Erc1967ProxyExample {
        erc1967: Erc1967Proxy,
    }

    unsafe impl TopLevelStorage for Erc1967ProxyExample {}

    #[public]
    impl Erc1967ProxyExample {
        #[constructor]
        fn constructor(
            &mut self,
            implementation: Address,
            data: Bytes,
        ) -> Result<(), proxy::erc1967::utils::Error> {
            self.erc1967.constructor(implementation, &data)
        }

        fn implementation(&self) -> Result<Address, Vec<u8>> {
            self.erc1967.implementation()
        }

        #[fallback]
        fn fallback(&mut self, calldata: &[u8]) -> ArbResult {
            unsafe { self.erc1967.do_fallback(calldata) }
        }
    }

    #[storage]
    struct LogicContract;

    unsafe impl TopLevelStorage for LogicContract {}

    #[public]
    impl LogicContract {
        fn contract_address(&self) -> Address {
            contract::address()
        }

        fn erc1967_utils_implementation(&self) -> Address {
            erc1967::utils::Erc1967Utils::get_implementation()
        }
    }

    sol! {
        interface LogicContractInterface {
            function contractAddress() external view returns (address);
            function erc1967UtilsImplementation() external view returns (address);
        }
    }

    #[motsu::test]
    fn test_addresses(
        proxy: Contract<Erc1967ProxyExample>,
        logic: Contract<LogicContract>,
        alice: Address,
    ) {
        proxy
            .sender(alice)
            .constructor(logic.address(), vec![].into())
            .expect("should be able to construct");

        let implementation = proxy
            .sender(alice)
            .implementation()
            .expect("should be able to get implementation");

        let call = LogicContractInterface::contractAddressCall {}.abi_encode();
        let contract_address = proxy
            .sender(alice)
            .fallback(&call)
            .map(|r| Address::from_slice(&r[12..]))
            .expect("should be able to get contract address");

        let call = LogicContractInterface::erc1967UtilsImplementationCall {}
            .abi_encode();
        let erc1967_utils_implementation = proxy
            .sender(alice)
            .fallback(&call)
            .map(|r| Address::from_slice(&r[12..]))
            .expect("should be able to get erc1967 utils implementation");

        assert_eq!(contract_address, proxy.address(), "contract_address isn't supposed to be equal to logic contract address: {:?}", logic.address());
        assert_eq!(implementation, logic.address()); // this passes
        assert_eq!(erc1967_utils_implementation, logic.address()); // this fails too
    }
}

Running this test prints:

--- STDERR:              openzeppelin-stylus proxy::utils::uups_upgradeable::tests::test_addresses ---

thread 'proxy::utils::uups_upgradeable::tests::test_addresses' panicked at contracts/src/proxy/utils/uups_upgradeable.rs:448:9:
assertion `left == right` failed: contract_address isn't supposed to be equal to logic contract address: 0x397bc5b97f629151e68146caedba62f10b47e426
  left: 0x397bc5b97f629151e68146caedba62f10b47e426
 right: 0x349381c1447d5247c9d0360cc20f41aad457c7bc
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

platform

• linux
• windows
• macos

Expected behavior

The above test case should pass.

Generally, delegate called contracts should use the caller storage instead of their own.

Contribution Guidelines

• I agree to follow this project's Contribution Guidelines

[0xNeshi]

Related #112, #111