@bloxchain/contracts 1.0.0-alpha.2 → 1.0.0-alpha.21

package/core/base/BaseStateMachine.sol

@@ -1,14 +1,14 @@
 // SPDX-License-Identifier: MPL-2.0
-pragma solidity 0.8.33;
+pragma solidity 0.8.35;
 
 // OpenZeppelin imports
 import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeable.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
-import "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
+import "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol";
 
 // Contracts imports
 import "../lib/EngineBlox.sol";
-import "../../utils/SharedValidation.sol";
+import "../lib/utils/SharedValidation.sol";
 import "./interface/IBaseStateMachine.sol";
 
 /**
@@ -36,7 +36,7 @@ import "./interface/IBaseStateMachine.sol";
  * - System configuration queries
  * - Event forwarding for external monitoring
  */
-abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, ReentrancyGuardUpgradeable {
+abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, ReentrancyGuardTransient {
     using EngineBlox for EngineBlox.SecureOperationState;
     using SharedValidation for *;
 
@@ -57,6 +57,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @param recovery The recovery address
      * @param timeLockPeriodSec The timelock period in seconds
      * @param eventForwarder The event forwarder address
+     * @dev Reentrancy protection inherits OpenZeppelin {ReentrancyGuardTransient} (EIP-1153 transient storage; requires Cancun+ / configured EVM such as Osaka). No initializer is required.
      */
     function _initializeBaseStateMachine(
         address initialOwner,
@@ -65,12 +66,14 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
         uint256 timeLockPeriodSec,
         address eventForwarder
     ) internal onlyInitializing {
-        __ERC165_init();
-        __ReentrancyGuard_init();
-        
-        _secureState.initialize(initialOwner, broadcaster, recovery, timeLockPeriodSec);
+        // Skip if already initialized (e.g. when multiple components call from Account.initialize)
+        if (!_secureState.initialized) {
+            __ERC165_init();
+
+            _secureState.initialize(initialOwner, broadcaster, recovery, timeLockPeriodSec);
 
-        _secureState.setEventForwarder(eventForwarder);
+            _secureState.setEventForwarder(eventForwarder);
+        }
     }
 
     // ============ SYSTEM ROLE QUERY FUNCTIONS ============
@@ -118,7 +121,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @dev Centralized function to request a transaction with common validation
      * @param requester The address requesting the transaction
      * @param target The target contract address
-     * @param value The ETH value to send (0 for standard function calls)
+     * @param value The ETH value to send (typically 0 for standard function calls; non-zero is supported for payable edge-case workflows)
      * @param gasLimit The gas limit for execution
      * @param operationType The type of operation
      * @param functionSelector The function selector for execution (NATIVE_TRANSFER_SELECTOR for simple native token transfers)
@@ -127,8 +130,9 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @notice Validates permissions for the calling function (request function), not the execution selector
      * @notice Execution functions are internal-only and don't need permission definitions
      * @notice This function is virtual to allow extensions to add hook functionality
-     * @notice For standard function calls: value=0, functionSelector=non-zero, params=encoded data
-     * @notice For simple native token transfers: value>0, functionSelector=NATIVE_TRANSFER_SELECTOR, params=""
+     * @notice Recommended standard calls: value=0, functionSelector=non-zero, params=encoded data
+     * @notice Flexible edge case: non-native selectors may intentionally forward ETH to payable targets
+     * @notice Native-only convenience flow: value>0, functionSelector=NATIVE_TRANSFER_SELECTOR, params=""
      */
     function _requestTransaction(
         address requester,
@@ -139,7 +143,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
         bytes4 functionSelector,
         bytes memory params
     ) internal virtual returns (EngineBlox.TxRecord memory) {
-        return EngineBlox.txRequest(
+        EngineBlox.TxRecord memory txRecord = EngineBlox.txRequest(
             _getSecureState(),
             requester,
             target,
@@ -150,6 +154,48 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
             functionSelector,
             params
         );
+        _postActionHook(txRecord);
+        return txRecord;
+    }
+
+    /**
+     * @dev Centralized function to request a transaction with payment details attached from the start
+     * @param requester The address requesting the transaction
+     * @param target The target contract address
+     * @param value The ETH value to send (typically 0 for standard function calls; non-zero is supported for payable edge-case workflows)
+     * @param gasLimit The gas limit for execution
+     * @param operationType The type of operation
+     * @param functionSelector The function selector for execution (NATIVE_TRANSFER_SELECTOR for simple native token transfers)
+     * @param params The encoded parameters for the function (empty for simple native token transfers)
+     * @param paymentDetails The payment details to attach to the transaction
+     * @return The created transaction record with payment set
+     * @notice Validates request permissions (same as request without payment)
+     * @notice This function is virtual to allow extensions to add hook functionality
+     */
+    function _requestTransactionWithPayment(
+        address requester,
+        address target,
+        uint256 value,
+        uint256 gasLimit,
+        bytes32 operationType,
+        bytes4 functionSelector,
+        bytes memory params,
+        EngineBlox.PaymentDetails memory paymentDetails
+    ) internal virtual returns (EngineBlox.TxRecord memory) {
+        EngineBlox.TxRecord memory txRecord = EngineBlox.txRequestWithPayment(
+            _getSecureState(),
+            requester,
+            target,
+            value,
+            gasLimit,
+            operationType,
+            bytes4(msg.sig),
+            functionSelector,
+            params,
+            paymentDetails
+        );
+        _postActionHook(txRecord);
+        return txRecord;
     }
 
     /**
@@ -159,12 +205,14 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @notice Validates permissions for the calling function (approval function selector), not the execution selector
      * @notice Execution functions are internal-only and don't need permission definitions
      * @notice This function is virtual to allow extensions to add hook functionality
-     * @notice Protected by ReentrancyGuard to prevent reentrancy attacks
+     * @notice Protected by ReentrancyGuardTransient to prevent reentrancy attacks
      */
     function _approveTransaction(
         uint256 txId
     ) internal virtual nonReentrant returns (EngineBlox.TxRecord memory) {
-        return EngineBlox.txDelayedApproval(_getSecureState(), txId, bytes4(msg.sig));
+        EngineBlox.TxRecord memory txRecord = EngineBlox.txDelayedApproval(_getSecureState(), txId, bytes4(msg.sig));
+        _postActionHook(txRecord);
+        return txRecord;
     }
 
     /**
@@ -174,12 +222,15 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @notice Validates permissions for the calling function (msg.sig) and handler selector from metaTx
      * @notice Uses EXECUTE_META_APPROVE action for permission checking
      * @notice This function is virtual to allow extensions to add hook functionality
-     * @notice Protected by ReentrancyGuard to prevent reentrancy attacks
+     * @notice Protected by ReentrancyGuardTransient to prevent reentrancy attacks
      */
     function _approveTransactionWithMetaTx(
         EngineBlox.MetaTransaction memory metaTx
     ) internal virtual nonReentrant returns (EngineBlox.TxRecord memory) {
-        return EngineBlox.txApprovalWithMetaTx(_getSecureState(), metaTx);
+        _validateMetaTxHandlerBinding(metaTx);
+        EngineBlox.TxRecord memory txRecord = EngineBlox.txApprovalWithMetaTx(_getSecureState(), metaTx);
+        _postActionHook(txRecord);
+        return txRecord;
     }
 
     /**
@@ -193,7 +244,9 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
     function _cancelTransaction(
         uint256 txId
     ) internal virtual returns (EngineBlox.TxRecord memory) {
-        return EngineBlox.txCancellation(_getSecureState(), txId, bytes4(msg.sig));
+        EngineBlox.TxRecord memory txRecord = EngineBlox.txCancellation(_getSecureState(), txId, bytes4(msg.sig));
+        _postActionHook(txRecord);
+        return txRecord;
     }
 
     /**
@@ -207,7 +260,10 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
     function _cancelTransactionWithMetaTx(
         EngineBlox.MetaTransaction memory metaTx
     ) internal virtual returns (EngineBlox.TxRecord memory) {
-        return EngineBlox.txCancellationWithMetaTx(_getSecureState(), metaTx);
+        _validateMetaTxHandlerBinding(metaTx);
+        EngineBlox.TxRecord memory txRecord = EngineBlox.txCancellationWithMetaTx(_getSecureState(), metaTx);
+        _postActionHook(txRecord);
+        return txRecord;
     }
 
     /**
@@ -217,26 +273,55 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @notice Validates permissions for the calling function (msg.sig) and handler selector from metaTx
      * @notice Uses EXECUTE_META_REQUEST_AND_APPROVE action for permission checking
      * @notice This function is virtual to allow extensions to add hook functionality
-     * @notice Protected by ReentrancyGuard to prevent reentrancy attacks
+     * @notice Protected by ReentrancyGuardTransient to prevent reentrancy attacks
      */
     function _requestAndApproveTransaction(
         EngineBlox.MetaTransaction memory metaTx
     ) internal virtual nonReentrant returns (EngineBlox.TxRecord memory) {
-        return EngineBlox.requestAndApprove(_getSecureState(), metaTx);
+        _validateMetaTxHandlerBinding(metaTx);
+        EngineBlox.TxRecord memory txRecord = EngineBlox.requestAndApprove(_getSecureState(), metaTx);
+        _postActionHook(txRecord);
+        return txRecord;
     }
 
     /**
-     * @dev Centralized function to update payment details for a pending transaction
-     * @param txId The transaction ID to update payment for
-     * @param paymentDetails The new payment details
-     * @notice This function is virtual to allow extensions to add hook functionality
+     * @dev Post-action hook invoked after any transaction operation that produces a TxRecord.
+     *      Override in derived contracts to add centralized post-tx logic (e.g. notifications, side effects).
+     * @param txRecord The transaction record produced by the operation
      */
-    function _updatePaymentForTransaction(
-        uint256 txId,
-        EngineBlox.PaymentDetails memory paymentDetails
-    ) internal virtual returns (EngineBlox.TxRecord memory) {
-        EngineBlox.updatePaymentForTransaction(_getSecureState(), txId, paymentDetails);
-        return _secureState.getTxRecord(txId);
+    function _postActionHook(EngineBlox.TxRecord memory txRecord) internal virtual {}
+
+    // ============ HOOK MANAGEMENT ============
+
+    /**
+     * @dev Sets the hook contract for a function selector (internal; no access control).
+     *      Extensions (e.g. HookManager) may expose an external setHook with owner check.
+     * @param functionSelector The function selector
+     * @param hook The hook contract address (must not be zero)
+     */
+    function _setHook(bytes4 functionSelector, address hook) internal {
+        EngineBlox.setHook(_getSecureState(), functionSelector, hook);
+    }
+
+    /**
+     * @dev Clears the hook contract for a function selector (internal; no access control).
+     *      Extensions may expose an external clearHook with owner check.
+     * @param functionSelector The function selector
+     * @param hook The hook contract address to remove (must not be zero)
+     */
+    function _clearHook(bytes4 functionSelector, address hook) internal {
+        EngineBlox.clearHook(_getSecureState(), functionSelector, hook);
+    }
+
+    /**
+     * @dev Returns all configured hooks for a function selector
+     * @param functionSelector The function selector
+     * @return hooks Array of hook contract addresses
+     * @notice Requires caller to have any role (via _validateAnyRole) to limit information visibility
+     */
+    function getHooks(bytes4 functionSelector) public view returns (address[] memory hooks) {
+        _validateAnyRole();
+        return EngineBlox.getHooks(_getSecureState(), functionSelector);
     }
 
     // ============ META-TRANSACTION UTILITIES ============
@@ -325,28 +410,29 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @param toTxId The ending transaction ID (inclusive)
      * @return The transaction history within the specified range
      * @notice Requires caller to have any role (via _validateAnyRole) to limit information visibility
+     * @notice Returns an empty array when **`txCounter == 0`** or when the range, after clamping to **1..txCounter**,
+     *         does not overlap any transaction ids (e.g. `fromTxId` entirely above the current counter).
      */
     function getTransactionHistory(uint256 fromTxId, uint256 toTxId) public view returns (EngineBlox.TxRecord[] memory) {
         _validateAnyRole();
-        
-        // Validate the range
+
+        uint256 counter = _secureState.txCounter;
+
+        // Normalize bounds to valid transaction id range [1, counter]
         fromTxId = fromTxId > 0 ? fromTxId : 1;
-        toTxId = toTxId > _secureState.txCounter ? _secureState.txCounter : toTxId;
-        
-        // Validate that fromTxId is less than toTxId
-        SharedValidation.validateLessThan(fromTxId, toTxId);
+        toTxId = toTxId > counter ? counter : toTxId;
+
+        // Empty history: invalid / non-overlapping clamped range (includes txCounter == 0 → toTxId == 0).
+        uint256 rangeSize = fromTxId > toTxId ? 0 : toTxId - fromTxId + 1;
+
+        if (rangeSize > 0) {
+            SharedValidation.validateRangeSize(rangeSize, 1000);
+        }
 
-        uint256 rangeSize = toTxId - fromTxId + 1;
-        
-        // For larger ranges, use paginated version
-        SharedValidation.validateRangeSize(rangeSize, 1000);
-        
         EngineBlox.TxRecord[] memory history = new EngineBlox.TxRecord[](rangeSize);
-        
         for (uint256 i = 0; i < rangeSize; i++) {
             history[i] = _secureState.getTxRecord(fromTxId + i);
         }
-        
         return history;
     }
 
@@ -368,7 +454,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      */
     function getPendingTransactions() public view returns (uint256[] memory) {
         _validateAnyRole();
-        return _secureState.getPendingTransactionsList();
+        return _secureState.getPendingTransactions();
     }
 
     // ============ ROLE AND PERMISSION QUERIES ============
@@ -377,7 +463,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @dev Gets the basic role information by its hash
      * @param roleHash The hash of the role to get
      * @return roleName The name of the role
-     * @return roleHashReturn The hash of the role
+     * @return hash The hash of the role
      * @return maxWallets The maximum number of wallets allowed for this role
      * @return walletCount The current number of wallets assigned to this role
      * @return isProtected Whether the role is protected from removal
@@ -385,7 +471,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      */
     function getRole(bytes32 roleHash) public view returns (
         string memory roleName,
-        bytes32 roleHashReturn,
+        bytes32 hash,
         uint256 maxWallets,
         uint256 walletCount,
         bool isProtected
@@ -429,29 +515,21 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @return Array of authorized wallet addresses
      * @notice Requires caller to have any role (via _validateAnyRole) to limit information visibility
      */
-    function getWalletsInRole(bytes32 roleHash) public view virtual returns (address[] memory) {
+    function getAuthorizedWallets(bytes32 roleHash) public view returns (address[] memory) {
         _validateAnyRole();
         _validateRoleExists(roleHash);
         return _getAuthorizedWallets(roleHash);
     }
 
     /**
-     * @dev Checks if a function schema exists
-     * @param functionSelector The function selector to check
-     * @return True if the function schema exists, false otherwise
-     */
-    function functionSchemaExists(bytes4 functionSelector) public view returns (bool) {
-        return _secureState.functions[functionSelector].functionSelector == functionSelector;
-    }
-
-    /**
-     * @dev Returns if an action is supported by a function
-     * @param functionSelector The function selector to check
-     * @param action The action to check
-     * @return True if the action is supported by the function, false otherwise
+     * @dev Gets function schema information
+     * @param functionSelector The function selector to get information for
+     * @return The full FunctionSchema struct (functionSignature, functionSelector, operationType, operationName, supportedActionsBitmap, enforceHandlerRelations, isProtected, isGrantRevocable, handlerForSelectors)
+     * @notice Reverts with ResourceNotFound if the schema does not exist
      */
-    function isActionSupportedByFunction(bytes4 functionSelector, EngineBlox.TxAction action) public view returns (bool) {
-        return _secureState.isActionSupportedByFunction(functionSelector, action);
+    function getFunctionSchema(bytes4 functionSelector) external view returns (EngineBlox.FunctionSchema memory) {
+        _validateAnyRole();
+        return _secureState.getFunctionSchema(functionSelector);
     }
 
     /**
@@ -485,7 +563,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      */
     function getSupportedOperationTypes() public view returns (bytes32[] memory) {
         _validateAnyRole();
-        return _secureState.getSupportedOperationTypesList();
+        return _secureState.getSupportedOperationTypes();
     }
 
     /**
@@ -495,7 +573,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      */
     function getSupportedRoles() public view returns (bytes32[] memory) {
         _validateAnyRole();
-        return _secureState.getSupportedRolesList();
+        return _secureState.getSupportedRoles();
     }
 
     /**
@@ -505,7 +583,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      */
     function getSupportedFunctions() public view returns (bytes4[] memory) {
         _validateAnyRole();
-        return _secureState.getSupportedFunctionsList();
+        return _secureState.getSupportedFunctions();
     }
 
     /**
@@ -542,15 +620,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @return Array of authorized wallet addresses
      */
     function _getAuthorizedWallets(bytes32 roleHash) internal view returns (address[] memory) {
-        EngineBlox.Role storage role = _secureState.roles[roleHash];
-        uint256 walletCount = role.walletCount;
-
-        address[] memory wallets = new address[](walletCount);
-        for (uint256 i = 0; i < walletCount; i++) {
-            wallets[i] = _getAuthorizedWalletAt(roleHash, i);
-        }
-
-        return wallets;
+        return EngineBlox.getAuthorizedWallets(_getSecureState(), roleHash);
     }
 
     /**
@@ -601,14 +671,14 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
     }
 
     /**
-     * @dev Centralized function to update assigned wallet for a role
+     * @dev Centralized function to update wallet for a role (replaces oldWallet with newWallet).
      * @param roleHash The role hash
      * @param newWallet The new wallet address
      * @param oldWallet The old wallet address
      * @notice This function is virtual to allow extensions to add hook functionality or additional validation
      */
-    function _updateAssignedWallet(bytes32 roleHash, address newWallet, address oldWallet) internal virtual {
-        EngineBlox.updateAssignedWallet(_getSecureState(), roleHash, newWallet, oldWallet);
+    function _updateWallet(bytes32 roleHash, address newWallet, address oldWallet) internal virtual {
+        EngineBlox.updateWallet(_getSecureState(), roleHash, newWallet, oldWallet);
     }
 
     /**
@@ -617,48 +687,56 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @notice This function is virtual to allow extensions to add hook functionality
      */
     function _updateTimeLockPeriod(uint256 newTimeLockPeriodSec) internal virtual {
+        uint256 oldPeriod = getTimeLockPeriodSec();
         EngineBlox.updateTimeLockPeriod(_getSecureState(), newTimeLockPeriodSec);
+        _logComponentEvent(abi.encode(oldPeriod, newTimeLockPeriodSec));
     }
 
     // ============ FUNCTION SCHEMA MANAGEMENT ============
 
     /**
-     * @dev Centralized function to create a function schema
+     * @dev Centralized function to register a function schema
      * @param functionSignature The function signature
      * @param functionSelector The function selector
      * @param operationName The operation name
      * @param supportedActionsBitmap The bitmap of supported actions
-     * @param isProtected Whether the function schema is protected
+     * @param enforceHandlerRelations Whether to enforce strict handler/schema alignment
+     * @param isProtected Whether the function schema is protected from unregister
+     * @param isGrantRevocable Whether role grants for this schema may be removed via `removeFunctionFromRole` (any role, including protected roles, when true)
      * @param handlerForSelectors Array of handler selectors
      * @notice This function is virtual to allow extensions to add hook functionality
      */
-    function _createFunctionSchema(
+    function _registerFunction(
         string memory functionSignature,
         bytes4 functionSelector,
         string memory operationName,
         uint16 supportedActionsBitmap,
+        bool enforceHandlerRelations,
         bool isProtected,
+        bool isGrantRevocable,
         bytes4[] memory handlerForSelectors
     ) internal virtual {
-        EngineBlox.createFunctionSchema(
+        EngineBlox.registerFunction(
             _getSecureState(),
             functionSignature,
             functionSelector,
             operationName,
             supportedActionsBitmap,
+            enforceHandlerRelations,
             isProtected,
+            isGrantRevocable,
             handlerForSelectors
         );
     }
 
     /**
-     * @dev Centralized function to remove a function schema
-     * @param functionSelector The function selector to remove
+     * @dev Centralized function to unregister a function schema
+     * @param functionSelector The function selector to unregister
      * @param safeRemoval Whether to perform safe removal (check for role references)
      * @notice This function is virtual to allow extensions to add hook functionality
      */
-    function _removeFunctionSchema(bytes4 functionSelector, bool safeRemoval) internal virtual {
-        EngineBlox.removeFunctionSchema(_getSecureState(), functionSelector, safeRemoval);
+    function _unregisterFunction(bytes4 functionSelector, bool safeRemoval) internal virtual {
+        EngineBlox.unregisterFunction(_getSecureState(), functionSelector, safeRemoval);
     }
 
     /**
@@ -686,6 +764,20 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
 
     // ============ PERMISSION VALIDATION ============
 
+    /**
+     * @dev Binds signed `MetaTxParams` to this wrapper's entrypoint (`msg.sig`) and verifying contract.
+     *      Must run in `BaseStateMachine` context, not inside linked `EngineBlox` library code (delegatecall
+     *      would make `msg.sig` refer to the library function, not the outer wrapper).
+     * @param metaTx The meta-transaction whose `params` are validated against `msg.sig` and `address(this)`.
+     */
+    function _validateMetaTxHandlerBinding(EngineBlox.MetaTransaction memory metaTx) internal view {
+        SharedValidation.validateMetaTxHandlerSelectorBinding(
+            metaTx.params.handlerSelector,
+            bytes4(msg.sig)
+        );
+        SharedValidation.validateMetaTxHandlerContractBinding(metaTx.params.handlerContract);
+    }
+
     /**
      * @dev Centralized function to validate that the caller has any role
      */
@@ -721,7 +813,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
     /**
      * @dev Adds a function selector to the system macro selectors set.
      *      Macro selectors are allowed to target address(this) for system-level operations.
-     * @param functionSelector The function selector to add (e.g. NATIVE_TRANSFER_SELECTOR, UPDATE_PAYMENT_SELECTOR).
+     * @param functionSelector The function selector to add (e.g. NATIVE_TRANSFER_SELECTOR).
      */
     function _addMacroSelector(bytes4 functionSelector) internal {
         EngineBlox.addMacroSelector(_getSecureState(), functionSelector);
@@ -763,8 +855,8 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @param target The target address to whitelist
      * @notice This function is virtual to allow extensions to add hook functionality
      */
-    function _addTargetToFunctionWhitelist(bytes4 functionSelector, address target) internal virtual {
-        _getSecureState().addTargetToFunctionWhitelist(functionSelector, target);
+    function _addTargetToWhitelist(bytes4 functionSelector, address target) internal virtual {
+        _getSecureState().addTargetToWhitelist(functionSelector, target);
     }
 
     /**
@@ -773,16 +865,18 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @param target The target address to remove
      * @notice This function is virtual to allow extensions to add hook functionality
      */
-    function _removeTargetFromFunctionWhitelist(bytes4 functionSelector, address target) internal virtual {
-        _getSecureState().removeTargetFromFunctionWhitelist(functionSelector, target);
+    function _removeTargetFromWhitelist(bytes4 functionSelector, address target) internal virtual {
+        _getSecureState().removeTargetFromWhitelist(functionSelector, target);
     }
 
     /**
-     * @dev Centralized function to get all whitelisted targets for a function selector
+     * @dev Gets all whitelisted targets for a function selector
      * @param functionSelector The function selector
      * @return Array of whitelisted target addresses
+     * @notice Requires caller to have any role (via _validateAnyRole) to limit information visibility
      */
-    function _getFunctionWhitelistTargets(bytes4 functionSelector) internal view returns (address[] memory) {
+    function getFunctionWhitelistTargets(bytes4 functionSelector) public view returns (address[] memory) {
+        _validateAnyRole();
         return _getSecureState().getFunctionWhitelistTargets(functionSelector);
     }
 
@@ -794,31 +888,31 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
      * @param functionSchemas Array of function schema definitions  
      * @param roleHashes Array of role hashes
      * @param functionPermissions Array of function permissions (parallel to roleHashes)
-     * @param allowProtectedSchemas Whether to allow protected function schemas (default: true for factory settings)
-     * @notice When allowProtectedSchemas is false, reverts if any function schema is protected
-     * @notice This allows custom definitions to be restricted from creating protected schemas
+     * @param requireProtected When true, all function schemas must be protected; reverts if any is not
+     * @notice When requireProtected is true, every function schema must have isProtected == true
      */
     function _loadDefinitions(
         EngineBlox.FunctionSchema[] memory functionSchemas,
         bytes32[] memory roleHashes,
         EngineBlox.FunctionPermission[] memory functionPermissions,
-        bool allowProtectedSchemas
+        bool requireProtected
     ) internal {
         // Load function schemas
         for (uint256 i = 0; i < functionSchemas.length; i++) {
-            // Validate protected schemas if not allowed (for custom definitions)
-            if (!allowProtectedSchemas && functionSchemas[i].isProtected) {
-                revert SharedValidation.CannotModifyProtected(
-                    bytes32(functionSchemas[i].functionSelector)
+            // When enforcing, require every schema to be protected
+            if (requireProtected && !functionSchemas[i].isProtected) {
+                revert SharedValidation.ContractFunctionMustBeProtected(
+                    functionSchemas[i].functionSelector
                 );
             }
-            EngineBlox.createFunctionSchema(
-                _getSecureState(),
+            _registerFunction(
                 functionSchemas[i].functionSignature,
                 functionSchemas[i].functionSelector,
                 functionSchemas[i].operationName,
                 functionSchemas[i].supportedActionsBitmap,
+                functionSchemas[i].enforceHandlerRelations,
                 functionSchemas[i].isProtected,
+                functionSchemas[i].isGrantRevocable,
                 functionSchemas[i].handlerForSelectors
             );
         }
@@ -826,11 +920,7 @@ abstract contract BaseStateMachine is Initializable, ERC165Upgradeable, Reentran
         // Load role permissions using parallel arrays
         SharedValidation.validateArrayLengthMatch(roleHashes.length, functionPermissions.length);
         for (uint256 i = 0; i < roleHashes.length; i++) {
-            EngineBlox.addFunctionToRole(
-                _getSecureState(),
-                roleHashes[i],
-                functionPermissions[i]
-            );
+            _addFunctionToRole(roleHashes[i], functionPermissions[i]);
         }
     }