refactor(ctb): Change to removeOwners

This commit introduces the ability to remove multiple owners at once in the LivenessModule.sol contract. A new function removeOwners has been added, which accepts arrays of previous owners and owners to be removed. The existing removeOwner function has been updated to use a new internal function _removeOwner. The _getPrevOwner function has been removed as it is no longer needed.

packages/contracts-bedrock/src/Safe/LivenessModule.sol

@@ -37,9 +37,6 @@ contract LivenessModule is ISemver {
     ///         keccak256("guard_manager.guard.address")
     uint256 internal constant GUARD_STORAGE_SLOT = 0x4a204f620c8c5ccdca3fd54d003badd85ba500436a431f0cbda4f558c93c34c8;
 
-    /// @notice The address of the first owner in the linked list of owners
-    address internal constant SENTINEL_OWNERS = address(0x1);
-
     /// @notice Semantic version.
     /// @custom:semver 1.0.0
     string public constant version = "1.0.0";
@@ -62,71 +59,73 @@ contract LivenessModule is ISemver {
         );
     }
 
+    function removeOwners(address[] memory _previousOwners, address[] memory _ownersToRemove) external {
+        require(_previousOwners.length == _ownersToRemove.length, "LivenessModule: arrays must be the same length");
+        for (uint256 i = 0; i < _previousOwners.length; i++) {
+            _removeOwner(_previousOwners[i], _ownersToRemove[i]);
+        }
+        _verifyFinalState();
+    }
+
     /// @notice This function can be called by anyone to remove an owner that has not signed a transaction
     ///         during the liveness interval. If the number of owners drops below the minimum, then the
     ///         ownership of the Safe is transferred to the fallback owner.
-    function removeOwner(address owner) external {
-        // Check that the owner to remove has not signed a transaction in the last 30 days
-        require(
-            LIVENESS_GUARD.lastLive(owner) < block.timestamp - LIVENESS_INTERVAL,
-            "LivenessModule: owner has signed recently"
-        );
-
+    function _removeOwner(address _prevOwner, address _ownerToRemove) internal {
         // Calculate the new threshold
         address[] memory owners = SAFE.getOwners();
         uint256 numOwnersAfter = owners.length - 1;
         if (_isAboveMinOwners(numOwnersAfter)) {
+            // Check that the owner to remove has not signed a transaction in the last 30 days
+            require(
+                LIVENESS_GUARD.lastLive(_ownerToRemove) < block.timestamp - LIVENESS_INTERVAL,
+                "LivenessModule: owner has signed recently"
+            );
             // Call the Safe to remove the owner and update the threshold
             uint256 thresholdAfter = get75PercentThreshold(numOwnersAfter);
-            address prevOwner = _getPrevOwner(owner, owners);
-            _removeOwner({ _prevOwner: prevOwner, _owner: owner, _threshold: thresholdAfter });
+            _removeOwnerSafeCall({ _prevOwner: _prevOwner, _owner: _ownerToRemove, _threshold: thresholdAfter });
         } else {
             // The number of owners is dangerously low, so we wish to transfer the ownership of this Safe
-            // to the fallback owner.
-
-            // Remove owners one at a time starting from the last owner.
-            // Since we're removing them in order from last to first, the ordering will remain constant,
-            // and we shouldn't need to query the list of owners again.
-            address prevOwner;
-            for (uint256 i = owners.length - 1; i > 0; i--) {
-                address currentOwner = owners[i];
-                prevOwner = _getPrevOwner(currentOwner, owners);
-
-                // Call the Safe to remove the owner
-                _removeOwner({ _prevOwner: prevOwner, _owner: currentOwner, _threshold: 1 });
-            }
-
-            prevOwner = _getPrevOwner(owners[0], owners);
+            // to the fallback owner. Therefore we no longer need to validate the liveness of the owners
+            // before removing them.
+            if (numOwnersAfter == 0) {
                 // Add the fallback owner as the sole owner of the Safe
-            _swapToFallbackOwner({ _prevOwner: prevOwner, _oldOwner: owners[0] });
+                _swapToFallbackOwnerSafeCall({ _prevOwner: _prevOwner, _oldOwner: _ownerToRemove });
+            } else {
+                // Remove the owner and set the threshold to 1
+                _removeOwnerSafeCall({ _prevOwner: _prevOwner, _owner: _ownerToRemove, _threshold: 1 });
+            }
         }
-
-        _verifyFinalState();
     }
 
     /// @notice Sets the fallback owner as the sole owner of the Safe with a threshold of 1
     /// @param _prevOwner Owner that pointed to the owner to be replaced in the linked list
     /// @param _oldOwner Owner address to be replaced.
-    function _swapToFallbackOwner(address _prevOwner, address _oldOwner) internal {
+    function _swapToFallbackOwnerSafeCall(address _prevOwner, address _oldOwner) internal {
+        require(
             SAFE.execTransactionFromModule({
                 to: address(SAFE),
                 value: 0,
                 operation: Enum.Operation.Call,
                 data: abi.encodeCall(OwnerManager.swapOwner, (_prevOwner, _oldOwner, FALLBACK_OWNER))
-        });
+            }),
+            "LivenessModule: failed to swap to fallback owner"
+        );
     }
 
     /// @notice Removes the owner `owner` from the Safe and updates the threshold to `_threshold`.
     /// @param _prevOwner Owner that pointed to the owner to be removed in the linked list
     /// @param _owner Owner address to be removed.
     /// @param _threshold New threshold.
-    function _removeOwner(address _prevOwner, address _owner, uint256 _threshold) internal {
+    function _removeOwnerSafeCall(address _prevOwner, address _owner, uint256 _threshold) internal {
+        require(
             SAFE.execTransactionFromModule({
                 to: address(SAFE),
                 value: 0,
                 operation: Enum.Operation.Call,
                 data: abi.encodeCall(OwnerManager.removeOwner, (_prevOwner, _owner, _threshold))
-        });
+            }),
+            "LivenessModule: failed to remove owner"
+        );
     }
 
     /// @notice A FREI-PI invariant check enforcing requirements on number of owners and threshold.
@@ -158,20 +157,6 @@ contract LivenessModule is ISemver {
         );
     }
 
-    /// @notice Get the previous owner in the linked list of owners
-    /// @param _owner The owner whose previous owner we want to find
-    /// @param _owners The list of owners
-    function _getPrevOwner(address _owner, address[] memory _owners) internal pure returns (address prevOwner_) {
-        for (uint256 i = 0; i < _owners.length; i++) {
-            if (_owners[i] != _owner) continue;
-            if (i == 0) {
-                prevOwner_ = SENTINEL_OWNERS;
-                break;
-            }
-            prevOwner_ = _owners[i - 1];
-        }
-    }
-
     /// @notice For a given number of owners, return the lowest threshold which is greater than 75.
     ///         Note: this function returns 1 for numOwners == 1.
     function get75PercentThreshold(uint256 _numOwners) public pure returns (uint256 threshold_) {

packages/contracts-bedrock/test/LivenessModule.t.sol

@@ -5,26 +5,57 @@ import { Test, StdUtils } from "forge-std/Test.sol";
 import { Safe } from "safe-contracts/Safe.sol";
 import { SafeProxyFactory } from "safe-contracts/proxies/SafeProxyFactory.sol";
 import { ModuleManager } from "safe-contracts/base/ModuleManager.sol";
+import { OwnerManager } from "safe-contracts/base/OwnerManager.sol";
 import { Enum } from "safe-contracts/common/Enum.sol";
 import "test/safe-tools/SafeTestTools.sol";
 
 import { LivenessModule } from "src/Safe/LivenessModule.sol";
 import { LivenessGuard } from "src/Safe/LivenessGuard.sol";
 
+contract OwnerSimulator is OwnerManager {
+    constructor(address[] memory _owners, uint256 _threshold) {
+        setupOwners(_owners, _threshold);
+    }
+
+    function removeOwnerWrapped(address prevOwner, address owner, uint256 _threshold) public {
+        OwnerManager(address(this)).removeOwner(prevOwner, owner, _threshold);
+    }
+}
+
 contract LivenessModule_TestInit is Test, SafeTestTools {
     using SafeTestLib for SafeInstance;
 
+    /// @notice The address of the first owner in the linked list of owners
+    address internal constant SENTINEL_OWNERS = address(0x1);
+
     event SignersRecorded(bytes32 indexed txHash, address[] signers);
 
     LivenessModule livenessModule;
     LivenessGuard livenessGuard;
     SafeInstance safeInstance;
+    OwnerSimulator ownerSimulator;
     address fallbackOwner;
 
+    /// @notice Get the previous owner in the linked list of owners
+    /// @param _owner The owner whose previous owner we want to find
+    /// @param _owners The list of owners
+    function _getPrevOwner(address _owner, address[] memory _owners) internal pure returns (address prevOwner_) {
+        for (uint256 i = 0; i < _owners.length; i++) {
+            if (_owners[i] != _owner) continue;
+            if (i == 0) {
+                prevOwner_ = SENTINEL_OWNERS;
+                break;
+            }
+            prevOwner_ = _owners[i - 1];
+        }
+    }
+
     function setUp() public {
         // Create a Safe with 10 owners
         (, uint256[] memory keys) = makeAddrsAndKeys(10);
         safeInstance = _setupSafe(keys, 8);
+        ownerSimulator = new OwnerSimulator(safeInstance.owners, 1);
+
         livenessGuard = new LivenessGuard(safeInstance.safe);
         fallbackOwner = makeAddr("fallbackOwner");
         livenessModule = new LivenessModule({
@@ -79,22 +110,34 @@ contract LivenessModule_Get75PercentThreshold_Test is LivenessModule_TestInit {
 contract LivenessModule_RemoveOwner_Test is LivenessModule_TestInit {
     function test_removeOwner_oneOwner_succeeds() external {
         uint256 ownersBefore = safeInstance.owners.length;
+        address[] memory prevOwners = new address[](1);
+        address[] memory ownersToRemove = new address[](1);
+        ownersToRemove[0] = safeInstance.owners[0];
+        prevOwners[0] = _getPrevOwner(safeInstance.owners[0], safeInstance.owners);
+
+        address ownerToRemove = safeInstance.owners[0];
         vm.warp(block.timestamp + 30 days);
-        livenessModule.removeOwner(safeInstance.owners[0]);
+
+        livenessModule.removeOwners(prevOwners, ownersToRemove);
         assertEq(safeInstance.safe.getOwners().length, ownersBefore - 1);
     }
 
     function test_removeOwner_allOwners_succeeds() external {
         vm.warp(block.timestamp + 30 days);
         uint256 numOwners = safeInstance.owners.length;
-        uint256 minOwners = livenessModule.minOwners();
-
-        // Remove enough owners to trigger the transfer to the fallbackOwner
-        uint256 numToRemove = numOwners - minOwners + 1;
 
-        for (uint256 i = 0; i < numToRemove; i++) {
-            livenessModule.removeOwner(safeInstance.owners[i]);
+        address[] memory prevOwners = new address[](numOwners);
+        address[] memory ownersToRemove = new address[](numOwners);
+        address[] memory currentOwners;
+        for (uint256 i = 0; i < numOwners; i++) {
+            currentOwners = ownerSimulator.getOwners();
+            ownersToRemove[i] = safeInstance.owners[i];
+            prevOwners[i] = _getPrevOwner(safeInstance.owners[i], currentOwners);
+            if (currentOwners.length == 1) break;
+            ownerSimulator.removeOwnerWrapped(prevOwners[i], ownersToRemove[i], 1);
         }
+
+        livenessModule.removeOwners(prevOwners, ownersToRemove);
         assertEq(safeInstance.safe.getOwners().length, 1);
         assertEq(safeInstance.safe.getOwners()[0], fallbackOwner);
         assertEq(safeInstance.safe.getThreshold(), 1);