check and update two slots

contracts/Nexus.sol

@@ -237,6 +237,20 @@ contract Nexus is INexus, EIP712, BaseAccount, ExecutionHelper, ModuleManager, U
         return IValidator(validator).isValidSignatureWithSender(msg.sender, computeHash, truncatedSignature);
     }
 
+    /// @notice Retrieves the address of the current implementation from the EIP-1967 slot.
+    /// @notice Checks the 1967 implementation slot, if not found then checks the slot defined by address (Biconomy V2 smart account)
+    /// @return implementation The address of the current contract implementation.
+    function getImplementation() external view returns (address implementation) {
+        assembly {
+            implementation := sload(_ERC1967_IMPLEMENTATION_SLOT)
+        }
+        if(implementation == address(0)) {
+            assembly {
+            implementation := sload(address())
+            }
+        }
+    }
+
     /// @notice Checks if a specific module type is supported by this smart account.
     /// @param moduleTypeId The identifier of the module type to check.
     /// @return True if the module type is supported, false otherwise.
@@ -285,9 +299,22 @@ contract Nexus is INexus, EIP712, BaseAccount, ExecutionHelper, ModuleManager, U
     }
 
     /// Upgrades the contract to a new implementation and calls a function on the new contract.
+    /// @notice Updates two slots 1. ERC1967 slot and 
+    /// 2. address() slot in case if it's potentially upgraded earlier from Biconomy V2 account,
+    /// as Biconomy v2 Account (proxy) reads implementation from the slot that is defined by its address
     /// @param newImplementation The address of the new contract implementation.
     /// @param data The calldata to be sent to the new implementation.
-    function upgradeToAndCall(address newImplementation, bytes calldata data) public payable virtual override {
+    function upgradeToAndCall(address newImplementation, bytes calldata data) public payable virtual override onlyEntryPointOrSelf {
+        if(newImplementation == address(0)) revert InvalidImplementationAddress();
+        bool res;
+        assembly {
+            res := gt(extcodesize(newImplementation), 0)
+        }
+        if(res == false) revert InvalidImplementationAddress();
+        // update the address() storage slot as well.
+        assembly {
+            sstore(address(), newImplementation)
+        }
         UUPSUpgradeable.upgradeToAndCall(newImplementation, data);
     }
 

contracts/interfaces/INexus.sol

@@ -33,6 +33,12 @@ interface INexus is IERC4337Account, IERC7579Account, INexusEventsAndErrors {
 
     /// @notice Throws if zero address has been provided as Entry Point address
     error EntryPointCannotBeZero();
+
+    /// @notice Throws if the implementation address is invalid
+    error InvalidImplementationAddress();
+
+    /// @notice Throws if the implementation is not a contract
+    error ImplementationIsNotAContract();
 
     /// @notice Initializes the smart account with a validator and custom data.
     /// @dev This method sets up the account for operation, linking it with a validator and initializing it with specific data.

test/foundry/fork/arbitrum/ArbitrumSmartAccountUpgradeTest.t.sol

@@ -55,8 +55,9 @@ contract ArbitrumSmartAccountUpgradeTest is NexusTest_Base, ArbitrumSettings {
         address beforeUpgradeImplementation = IBiconomySmartAccountV2(SMART_ACCOUNT_V2_ADDRESS).getImplementation();
         assertNotEq(beforeUpgradeImplementation, address(newImplementation), "Implementation address does not match before upgrade.");
         test_UpgradeV2ToV3AndInitialize();
-        bytes32 SLOT = bytes32(uint256(uint160(SMART_ACCOUNT_V2_ADDRESS)));
-        address afterUpgradeImplementation = address(uint160(uint256(vm.load(SMART_ACCOUNT_V2_ADDRESS, SLOT)))); 
+        // bytes32 SLOT = bytes32(uint256(uint160(SMART_ACCOUNT_V2_ADDRESS)));
+        // address afterUpgradeImplementation = address(uint160(uint256(vm.load(SMART_ACCOUNT_V2_ADDRESS, SLOT)))); 
+        address afterUpgradeImplementation = Nexus(payable(SMART_ACCOUNT_V2_ADDRESS)).getImplementation();
         address expectedImplementation = address(newImplementation);
         assertEq(afterUpgradeImplementation, expectedImplementation, "Implementation address does not match after upgrade.");
     }

test/foundry/integration/UpgradeSmartAccountTest.t.sol

@@ -17,8 +17,9 @@ contract UpgradeSmartAccountTest is NexusTest_Base {
 
     /// @notice Tests that the current implementation address is correct
     function test_currentImplementationAddress() public {
-        bytes32 _ERC1967_IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
-        address currentImplementation = address(uint160(uint256(vm.load(address(BOB_ACCOUNT), _ERC1967_IMPLEMENTATION_SLOT)))); 
+        // bytes32 _ERC1967_IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+        // address currentImplementation = address(uint160(uint256(vm.load(address(BOB_ACCOUNT), _ERC1967_IMPLEMENTATION_SLOT)))); 
+        address currentImplementation = BOB_ACCOUNT.getImplementation();
         assertEq(currentImplementation, address(ACCOUNT_IMPLEMENTATION), "Current implementation address mismatch");
     }
 
@@ -33,11 +34,77 @@ contract UpgradeSmartAccountTest is NexusTest_Base {
 
         PackedUserOperation[] memory userOps = buildPackedUserOperation(BOB, BOB_ACCOUNT, EXECTYPE_DEFAULT, execution, address(VALIDATOR_MODULE));
         ENTRYPOINT.handleOps(userOps, payable(address(BOB.addr)));
-         bytes32 _ERC1967_IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
-        address newImplementation = address(uint160(uint256(vm.load(address(BOB_ACCOUNT), _ERC1967_IMPLEMENTATION_SLOT)))); 
+        // bytes32 _ERC1967_IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+        // address newImplementation = address(uint160(uint256(vm.load(address(BOB_ACCOUNT), _ERC1967_IMPLEMENTATION_SLOT)))); 
+        address newImplementation = BOB_ACCOUNT.getImplementation();
         assertEq(newImplementation, address(newSmartAccount), "New implementation address mismatch");
     }
 
+    /// @notice Tests the upgrade of the smart account implementation with invalid call data
+    function test_upgradeImplementation_invalidCallData() public {
+        address _ENTRYPOINT = 0x0000000071727De22E5E9d8BAf0edAc6f37da032;
+        Nexus newSmartAccount = new Nexus(_ENTRYPOINT);
+        bytes memory callData = abi.encodeWithSelector(Nexus.upgradeToAndCall.selector, address(newSmartAccount), bytes(hex"1234"));
+        Execution[] memory execution = new Execution[](1);
+        execution[0] = Execution(address(BOB_ACCOUNT), 0, callData);
+        PackedUserOperation[] memory userOps = buildPackedUserOperation(BOB, BOB_ACCOUNT, EXECTYPE_DEFAULT, execution, address(VALIDATOR_MODULE));
+        bytes memory expectedRevertReason = abi.encodeWithSelector(MissingFallbackHandler.selector, bytes4(hex"1234"));
+        bytes32 userOpHash = ENTRYPOINT.getUserOpHash(userOps[0]);
+        // Expect the UserOperationRevertReason event
+        vm.expectEmit(true, true, true, true);
+        emit UserOperationRevertReason(
+            userOpHash, // userOpHash
+            address(BOB_ACCOUNT), // sender
+            userOps[0].nonce, // nonce
+            expectedRevertReason
+        );
+        ENTRYPOINT.handleOps(userOps, payable(address(BOB.addr)));
+    }
+
+    /// @notice Tests the upgrade of the smart account implementation with an invalid address
+    function test_upgradeImplementation_InvalidAddress() public {
+        /// @note "" means empty calldata. this will just update the implementation but not setup the account.
+        bytes memory callData = abi.encodeWithSelector(Nexus.upgradeToAndCall.selector, address(0), "");
+        Execution[] memory execution = new Execution[](1);
+        execution[0] = Execution(address(BOB_ACCOUNT), 0, callData);
+        PackedUserOperation[] memory userOps = buildPackedUserOperation(BOB, BOB_ACCOUNT, EXECTYPE_DEFAULT, execution, address(VALIDATOR_MODULE));
+        bytes memory expectedRevertReason = abi.encodeWithSignature("InvalidImplementationAddress()");
+        bytes32 userOpHash = ENTRYPOINT.getUserOpHash(userOps[0]);
+        // Expect the UserOperationRevertReason event
+        vm.expectEmit(true, true, true, true);
+        emit UserOperationRevertReason(
+            userOpHash, // userOpHash
+            address(BOB_ACCOUNT), // sender
+            userOps[0].nonce, // nonce
+            expectedRevertReason
+        );
+        ENTRYPOINT.handleOps(userOps, payable(address(BOB.addr)));
+    }
+
+    /// @notice Tests the upgrade of the smart account implementation with an invalid address
+    function test_upgradeImplementation_InvalidAddress_NotAContract() public {
+        /// @note "" means empty calldata. this will just update the implementation but not setup the account.
+        bytes memory callData = abi.encodeWithSelector(Nexus.upgradeToAndCall.selector, BOB.addr, "");
+        Execution[] memory execution = new Execution[](1);
+        execution[0] = Execution(address(BOB_ACCOUNT), 0, callData);
+        PackedUserOperation[] memory userOps = buildPackedUserOperation(BOB, BOB_ACCOUNT, EXECTYPE_DEFAULT, execution, address(VALIDATOR_MODULE));
+        bytes memory expectedRevertReason = abi.encodeWithSignature("InvalidImplementationAddress()");
+        bytes32 userOpHash = ENTRYPOINT.getUserOpHash(userOps[0]);
+        // Expect the UserOperationRevertReason event
+        vm.expectEmit(true, true, true, true);
+        emit UserOperationRevertReason(
+            userOpHash, // userOpHash
+            address(BOB_ACCOUNT), // sender
+            userOps[0].nonce, // nonce
+            expectedRevertReason
+        );
+        ENTRYPOINT.handleOps(userOps, payable(address(BOB.addr)));
+    }
+
+    /// Could add...
+    /// Access control on upgrades
+    /// send setup data instead of empty data
+
     /// @notice Tests the entire upgrade process
     function test_upgradeSmartAccount() public {
         test_proxiableUUIDSlot();

test/foundry/unit/fuzz/TestFuzz_ERC4337Account.t.sol

@@ -83,6 +83,7 @@ contract TestFuzz_ERC4337Account is NexusTest_Base {
     /// @param amount The amount to withdraw.
     function testFuzz_WithdrawDepositTo(address to, uint256 amount) public {
         vm.assume(!isContract(to)); // Valid 'to' address and skip precompiles
+        vm.assume(uint160(address(to)) > 10);
         vm.assume(amount > 0.01 ether && amount <= 50 ether); // Restricting the amount to a reasonable upper limit and ensuring it's greater than 0
         vm.assume(to.balance == 0);
         // Fund the BOB_ACCOUNT with more than just the deposit amount to cover potential transaction fees

test/foundry/utils/EventsAndErrors.sol

@@ -17,6 +17,8 @@ contract EventsAndErrors {
 
     // Define all errors
     // General Errors
+    error MissingFallbackHandler(bytes4 sig);
+    error InvalidImplementationAddress();
     error AccountInitializationFailed();
     error AccountAccessUnauthorized();
     error ExecutionFailed();

test/hardhat/smart-account/MSA.Basics.specs.ts

@@ -120,7 +120,7 @@ describe("Nexus Basic Specs", function () {
     });
 
     it("Should get implementation address of smart account", async () => {
-      const slot =
+      /*const slot =
         "0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";
       // Get the provider (default to Hardhat's local network)
       const provider = ethers.provider;
@@ -131,7 +131,8 @@ describe("Nexus Basic Specs", function () {
         "latest",
       ]);
       // Convert the storage value to an address
-      const saImplementation = ethers.getAddress(toHex(BigInt(storageValue)));
+      const saImplementation = ethers.getAddress(toHex(BigInt(storageValue)));*/
+      const saImplementation = await smartAccount.getImplementation();
       console.log("Implementation Address: ", saImplementation);
       expect(saImplementation).to.not.equal(ZeroAddress);
     });