feat: update _checkOwner() (#41)

0xSplits · Aug 15, 2024 · e23179e · e23179e
1 parent 1fed001
commit e23179e
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 0 deletions.
diff --git a/packages/smart-vaults/src/vault/SmartVault.sol b/packages/smart-vaults/src/vault/SmartVault.sol
@@ -306,6 +306,25 @@ contract SmartVault is IAccount, Ownable, UUPSUpgradeable, MultiSignerAuth, ERC1
 
     function _authorize() internal view override(MultiSignerAuth, FallbackManager, ModuleManager) onlySelf { }
 
+    /**
+     * @dev Conditions for a valid owner check:
+     *      if owner is non zero, caller must be owner.
+     *      If owner is address(0), contract can call itself.
+     */
+    function _checkOwner() internal view override {
+        address owner;
+        address caller = msg.sender;
+
+        /// @solidity memory-safe-assembly
+        assembly {
+            owner := sload(_OWNER_SLOT)
+        }
+
+        if (owner == caller) return;
+        if (owner == address(0) && caller == address(this)) return;
+        revert Unauthorized();
+    }
+
     /// @dev Get light user op hash of the Packed user operation.
     function _getLightUserOpHash(PackedUserOperation calldata userOp_) internal view returns (bytes32) {
         return keccak256(abi.encode(userOp_.hashLight(), entryPoint(), block.chainid));

diff --git a/packages/smart-vaults/test/SmartVault.t.sol b/packages/smart-vaults/test/SmartVault.t.sol
@@ -6,6 +6,7 @@ import "@web-authn/../test/Utils.sol";
 import "@web-authn/WebAuthn.sol";
 import { FCL_Elliptic_ZZ } from "FreshCryptoLib/FCL_elliptic.sol";
 
+import { UUPSUpgradeable } from "solady/utils/UUPSUpgradeable.sol";
 import { UserOperationLib } from "src/library/UserOperationLib.sol";
 import { MultiSignerLib } from "src/signers/MultiSigner.sol";
 
@@ -821,6 +822,19 @@ contract SmartVaultTest is BaseTest {
         vault.transferOwnership(BOB.addr);
     }
 
+    function test_transferOwnership_when_ownerIsZero() public {
+        vm.prank(vault.owner());
+        vault.renounceOwnership();
+
+        Caller.Call memory call =
+            Caller.Call(address(vault), 0, abi.encodeWithSelector(Ownable.transferOwnership.selector, BOB.addr));
+
+        vm.prank(ENTRY_POINT);
+        vault.execute(call);
+
+        assertEq(vault.owner(), BOB.addr);
+    }
+
     /* -------------------------------------------------------------------------- */
     /*                             SIGNER SET UPDATES                             */
     /* -------------------------------------------------------------------------- */
@@ -1096,4 +1110,71 @@ contract SmartVaultTest is BaseTest {
         vm.prank(caller_);
         vault.executeFromModule(calls_);
     }
+
+    /* -------------------------------------------------------------------------- */
+    /*                               ONLY SELF TESTS                              */
+    /* -------------------------------------------------------------------------- */
+
+    function test_onlySelf_addSigner() public {
+        Caller.Call memory call = Caller.Call(
+            address(vault), 0, abi.encodeWithSelector(MultiSignerAuth.addSigner.selector, createSigner(BOB.addr), 4)
+        );
+
+        vm.prank(ENTRY_POINT);
+        vault.execute(call);
+
+        assertEq(vault.getSigner(4), createSigner(BOB.addr));
+        assertEq(vault.getSignerCount(), 4);
+    }
+
+    function test_onlySelf_removeSigner() public {
+        Caller.Call memory call =
+            Caller.Call(address(vault), 0, abi.encodeWithSelector(MultiSignerAuth.removeSigner.selector, 0));
+
+        vm.prank(ENTRY_POINT);
+        vault.execute(call);
+
+        assertEq(vault.getSigner(0), createSigner(address(0)));
+        assertEq(vault.getSignerCount(), 2);
+    }
+
+    function test_onlySelf_updateThreshold() public {
+        Caller.Call memory call =
+            Caller.Call(address(vault), 0, abi.encodeWithSelector(MultiSignerAuth.updateThreshold.selector, 2));
+
+        vm.prank(ENTRY_POINT);
+        vault.execute(call);
+
+        assertEq(vault.getThreshold(), 2);
+    }
+
+    function test_onlySelf_upgradeImplementation_when_ownerIsZero() public {
+        vm.prank(vault.owner());
+        vault.renounceOwnership();
+
+        address newImplementation = address(new SmartVault());
+        Caller.Call memory call = Caller.Call(
+            address(vault),
+            0,
+            abi.encodeWithSelector(UUPSUpgradeable.upgradeToAndCall.selector, newImplementation, new bytes(0))
+        );
+
+        vm.prank(ENTRY_POINT);
+        vault.execute(call);
+
+        assertEq(vault.getImplementation(), newImplementation);
+    }
+
+    function test_onlySelf_upgradeImplementation_revertsWhen_ownerIsNotZero() public {
+        address newImplementation = address(new SmartVault());
+        Caller.Call memory call = Caller.Call(
+            address(vault),
+            0,
+            abi.encodeWithSelector(UUPSUpgradeable.upgradeToAndCall.selector, newImplementation, new bytes(0))
+        );
+
+        vm.prank(ENTRY_POINT);
+        vm.expectRevert(abi.encodeWithSelector(Unauthorized.selector));
+        vault.execute(call);
+    }
 }