check and update two slots

contracts/Nexus.sol

@@ -237,6 +237,20 @@
         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 @@
     }
 
     /// 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,8 +34,9 @@ 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");
     }
 

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);
     });