🎨 fixes upon review

bcnmy · Oct 18, 2023 · 2312657 · 2312657
1 parent 95033fe
commit 2312657
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/contracts/smart-account/interfaces/modules/IMultiOwnedECDSAModule.sol b/contracts/smart-account/interfaces/modules/IMultiOwnedECDSAModule.sol
@@ -25,6 +25,7 @@ interface IMultiOwnedECDSAModule {
     error NotEOA(address account);
     error ZeroAddressNotAllowedAsOwner();
     error OwnerAlreadyUsedForSmartAccount(address owner, address smartAccount);
+    error NotAnOwner(address owner, address smartAccount);
 
     /**
      * @dev Initializes the module for a Smart Account.

diff --git a/contracts/smart-account/modules/MultiOwnedECDSAModule.sol b/contracts/smart-account/modules/MultiOwnedECDSAModule.sol
@@ -73,6 +73,10 @@ contract MultiOwnedECDSAModule is
         if (owner == address(0)) revert ZeroAddressNotAllowedAsOwner();
         if (owner == newOwner)
             revert OwnerAlreadyUsedForSmartAccount(newOwner, msg.sender);
+        if (!_smartAccountOwners[owner][msg.sender])
+            revert NotAnOwner(owner, msg.sender);
+        if (_smartAccountOwners[newOwner][msg.sender])
+            revert OwnerAlreadyUsedForSmartAccount(newOwner, msg.sender);
         _transferOwnership(msg.sender, owner, newOwner);
     }
 
@@ -84,11 +88,15 @@ contract MultiOwnedECDSAModule is
             revert OwnerAlreadyUsedForSmartAccount(owner, msg.sender);
 
         _smartAccountOwners[owner][msg.sender] = true;
-        numberOfOwners[msg.sender]++;
+        unchecked {
+           ++numberOfOwners[msg.sender];
+        }
     }
 
     /// @inheritdoc IMultiOwnedECDSAModule
     function removeOwner(address owner) external override {
+        if(!_smartAccountOwners[owner][msg.sender]) 
+            revert NotAnOwner(owner, msg.sender);
         _transferOwnership(msg.sender, owner, address(0));
         --numberOfOwners[msg.sender];
     }

diff --git a/test/bundler-integration/module/MultiOwnedECDSA.Module.specs.ts b/test/bundler-integration/module/MultiOwnedECDSA.Module.specs.ts
@@ -149,7 +149,6 @@ describe("MultiOwned ECDSA Module (with Bundler):", async () => {
   describe("transferOwnership: ", async () => {
     it("Call transferOwnership from userSA and it successfully changes owner ", async () => {
       const { multiOwnedECDSAModule, entryPoint, userSA } = await setupTests();
-      // console.log(await userSA.getImplementation());
 
       // Calldata to set Eve as owner
       const txnData1 = multiOwnedECDSAModule.interface.encodeFunctionData(