test: Security Policy Plugin Management

bcnmy · Oct 12, 2023 · 1dd2198 · 1dd2198
1 parent a72ec6b
commit 1dd2198
Show file tree

Hide file tree

Showing 4 changed files with 1,302 additions and 42 deletions.
diff --git a/contracts/smart-account/interfaces/modules/ISecurityPolicyManagerPlugin.sol b/contracts/smart-account/interfaces/modules/ISecurityPolicyManagerPlugin.sol
@@ -5,10 +5,7 @@ import {ISecurityPolicyPlugin} from "./ISecurityPolicyPlugin.sol";
 
 address constant SENTINEL_MODULE_ADDRESS = address(0x1);
 
-/// @title Security Policy Manager Plugin
-/// @author @ankurdubey521
-/// @dev Execution Phase Plugin responsible for enforcing security policies during plugin installation on the smart contract wallet
-interface ISecurityPolicyManagerPlugin {
+interface ISecurityPolicyManagerPluginEventsErrors {
     event SecurityPolicyEnabled(address indexed scw, address indexed policy);
     event SecurityPolicyDisabled(address indexed scw, address indexed policy);
 
@@ -17,7 +14,14 @@ interface ISecurityPolicyManagerPlugin {
     error InvalidSecurityPolicyAddress(address policy);
     error InvalidPointerAddress(address pointer);
     error EmptyPolicyList();
+}
 
+/// @title Security Policy Manager Plugin
+/// @author @ankurdubey521
+/// @dev Execution Phase Plugin responsible for enforcing security policies during plugin installation on the smart contract wallet
+interface ISecurityPolicyManagerPlugin is
+    ISecurityPolicyManagerPluginEventsErrors
+{
     /// @dev Enables the security policies for the smart contract wallet. Used during the setup process.
     /// @param _policy The security policy to be enabled
     function enableSecurityPolicy(ISecurityPolicyPlugin _policy) external;
@@ -56,8 +60,12 @@ interface ISecurityPolicyManagerPlugin {
 
     /// @dev Returns the security policy for the smart contract wallet.
     /// @param _scw The address of the smart contract wallet
-    /// @return SecurityPolicy The security policy of the smart contract wallet
-    function securityPolicies(
-        address _scw
-    ) external view returns (ISecurityPolicyPlugin[] memory);
+    /// @param _start The address of the first security policy in the list
+    /// @param _pageSize The number of security policies to be returned
+    /// @return enabledPolicies The list of enabled security policies
+    function securityPoliciesPaginated(
+        address _scw,
+        address _start,
+        uint256 _pageSize
+    ) external view returns (ISecurityPolicyPlugin[] memory enabledPolicies);
 }
diff --git a/contracts/smart-account/modules/SecurityPolicyManagerPlugin.sol b/contracts/smart-account/modules/SecurityPolicyManagerPlugin.sol
@@ -3,10 +3,25 @@ pragma solidity 0.8.17;
 
 import {ISecurityPolicyManagerPlugin, ISecurityPolicyPlugin, SENTINEL_MODULE_ADDRESS} from "contracts/smart-account/interfaces/modules/ISecurityPolicyManagerPlugin.sol";
 
+/// @title Security Policy Manager Plugin
+/// @author @ankurdubey521
+/// @dev Execution Phase Plugin responsible for enforcing security policies during plugin installation on the smart contract wallet
 contract SecurityPolicyManagerPlugin is ISecurityPolicyManagerPlugin {
     mapping(address => mapping(address => address))
         internal enabledSecurityPoliciesLinkedList;
 
+    ////////////////////////// PLUGIN INSTALLATION FUNCTIONS //////////////////////////
+
+    /// @inheritdoc ISecurityPolicyManagerPlugin
+    function checkSetupAndEnableModule(
+        address,
+        bytes calldata
+    ) external override returns (address) {
+        revert("Not implemented");
+    }
+
+    ////////////////////////// SECURITY POLICY MANAGEMENT FUNCTIONS //////////////////////////
+
     /// @inheritdoc ISecurityPolicyManagerPlugin
     function enableSecurityPolicy(
         ISecurityPolicyPlugin _policy
@@ -27,9 +42,10 @@ contract SecurityPolicyManagerPlugin is ISecurityPolicyManagerPlugin {
             revert SecurityPolicyAlreadyEnabled(address(_policy));
         }
 
-        enabledSecurityPolicies[address(_policy)] = enabledSecurityPolicies[
-            SENTINEL_MODULE_ADDRESS
-        ];
+        address head = enabledSecurityPolicies[SENTINEL_MODULE_ADDRESS];
+        enabledSecurityPolicies[address(_policy)] = head == address(0x0)
+            ? SENTINEL_MODULE_ADDRESS
+            : head;
         enabledSecurityPolicies[SENTINEL_MODULE_ADDRESS] = address(_policy);
 
         emit SecurityPolicyEnabled(msg.sender, address(_policy));
@@ -64,7 +80,9 @@ contract SecurityPolicyManagerPlugin is ISecurityPolicyManagerPlugin {
                 revert SecurityPolicyAlreadyEnabled(address(policy));
             }
 
-            enabledSecurityPolicies[policy] = head;
+            enabledSecurityPolicies[policy] = head == address(0x0)
+                ? SENTINEL_MODULE_ADDRESS
+                : head;
             head = policy;
 
             emit SecurityPolicyEnabled(msg.sender, policy);
@@ -110,6 +128,7 @@ contract SecurityPolicyManagerPlugin is ISecurityPolicyManagerPlugin {
         emit SecurityPolicyDisabled(msg.sender, address(_policy));
     }
 
+    /* solhint-disable code-complexity*/
     /// @inheritdoc ISecurityPolicyManagerPlugin
     function disableSecurityPoliciesRange(
         ISecurityPolicyPlugin _start,
@@ -145,16 +164,21 @@ contract SecurityPolicyManagerPlugin is ISecurityPolicyManagerPlugin {
 
         bool endFound = false;
         address current = address(_start);
-        while (current != address(_end) || current != SENTINEL_MODULE_ADDRESS) {
-            if (current == address(_end)) {
-                endFound = true;
-            }
-
+        while (true) {
             address next = enabledSecurityPolicies[current];
             delete enabledSecurityPolicies[current];
 
             emit SecurityPolicyDisabled(msg.sender, current);
 
+            if (current == address(_end)) {
+                endFound = true;
+                break;
+            }
+
+            if (current == SENTINEL_MODULE_ADDRESS) {
+                break;
+            }
+
             current = next;
         }
 
@@ -164,17 +188,42 @@ contract SecurityPolicyManagerPlugin is ISecurityPolicyManagerPlugin {
     }
 
     /// @inheritdoc ISecurityPolicyManagerPlugin
-    function checkSetupAndEnableModule(
-        address,
-        bytes calldata
-    ) external override returns (address) {
-        revert("Not implemented");
-    }
+    function securityPoliciesPaginated(
+        address _scw,
+        address _start,
+        uint256 _pageSize
+    ) external view returns (ISecurityPolicyPlugin[] memory enabledPolicies) {
+        enabledPolicies = new ISecurityPolicyPlugin[](_pageSize);
+        uint256 actualEnabledPoliciesLength;
 
-    /// @inheritdoc ISecurityPolicyManagerPlugin
-    function securityPolicies(
-        address
-    ) external view override returns (ISecurityPolicyPlugin[] memory) {
-        revert("Not implemented");
+        mapping(address => address)
+            storage enabledSecurityPolicies = enabledSecurityPoliciesLinkedList[
+                _scw
+            ];
+
+        if (_start == address(0)) {
+            _start = SENTINEL_MODULE_ADDRESS;
+        }
+
+        ISecurityPolicyPlugin current = ISecurityPolicyPlugin(_start);
+        do {
+            if (current != ISecurityPolicyPlugin(SENTINEL_MODULE_ADDRESS)) {
+                enabledPolicies[actualEnabledPoliciesLength] = current;
+                unchecked {
+                    ++actualEnabledPoliciesLength;
+                }
+            }
+            current = ISecurityPolicyPlugin(
+                enabledSecurityPolicies[address(current)]
+            );
+        } while (
+            actualEnabledPoliciesLength < _pageSize &&
+                current != ISecurityPolicyPlugin(SENTINEL_MODULE_ADDRESS) &&
+                current != ISecurityPolicyPlugin(address(0))
+        );
+
+        assembly {
+            mstore(enabledPolicies, actualEnabledPoliciesLength)
+        }
     }
 }