smom-dbis-138/contracts/smart-accounts/AccountWalletRegistryExtended.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin/contracts/access/AccessControl.sol";
import "../emoney/interfaces/IAccountWalletRegistry.sol";

/**
 * @title AccountWalletRegistryExtended
 * @notice Extends account-wallet registry with Smart Account (contract) support
 * @dev Links accountRefId to contract addresses (smart accounts); walletRefId = keccak256(abi.encodePacked(smartAccount))
 */
contract AccountWalletRegistryExtended is IAccountWalletRegistry, AccessControl {
    bytes32 public constant ACCOUNT_MANAGER_ROLE = keccak256("ACCOUNT_MANAGER_ROLE");

    address public smartAccountFactory;
    address public entryPoint;

    mapping(bytes32 => WalletLink[]) private _accountToWallets;
    mapping(bytes32 => bytes32[]) private _walletToAccounts;
    mapping(bytes32 => mapping(bytes32 => uint256)) private _walletIndex;
    mapping(bytes32 => mapping(bytes32 => bool)) private _walletAccountExists;
    mapping(bytes32 => address) private _smartAccountAddress; // walletRefId => smart account address

    event SmartAccountLinked(bytes32 indexed accountRefId, address indexed smartAccount, bytes32 provider);

    constructor(address admin, address smartAccountFactory_, address entryPoint_) {
        _grantRole(DEFAULT_ADMIN_ROLE, admin);
        _grantRole(ACCOUNT_MANAGER_ROLE, admin);
        smartAccountFactory = smartAccountFactory_;
        entryPoint = entryPoint_;
    }

    function setSmartAccountFactory(address factory) external onlyRole(DEFAULT_ADMIN_ROLE) {
        smartAccountFactory = factory;
    }

    function setEntryPoint(address entryPoint_) external onlyRole(DEFAULT_ADMIN_ROLE) {
        entryPoint = entryPoint_;
    }

    function linkAccountToWallet(bytes32 accountRefId, bytes32 walletRefId, bytes32 provider)
        external
        override
        onlyRole(ACCOUNT_MANAGER_ROLE)
    {
        require(accountRefId != bytes32(0), "zero accountRefId");
        require(walletRefId != bytes32(0), "zero walletRefId");
        require(provider != bytes32(0), "zero provider");

        if (_walletAccountExists[walletRefId][accountRefId]) {
            uint256 index = _walletIndex[accountRefId][walletRefId];
            require(index < _accountToWallets[accountRefId].length, "index out of bounds");
            WalletLink storage link = _accountToWallets[accountRefId][index];
            require(link.walletRefId == walletRefId, "link mismatch");
            link.active = true;
            link.linkedAt = uint64(block.timestamp);
        } else {
            WalletLink memory newLink = WalletLink({
                walletRefId: walletRefId,
                linkedAt: uint64(block.timestamp),
                active: true,
                provider: provider
            });
            _accountToWallets[accountRefId].push(newLink);
            _walletIndex[accountRefId][walletRefId] = _accountToWallets[accountRefId].length - 1;
            _walletAccountExists[walletRefId][accountRefId] = true;
            _walletToAccounts[walletRefId].push(accountRefId);
        }
        emit AccountWalletLinked(accountRefId, walletRefId, provider, uint64(block.timestamp));
    }

    function linkSmartAccount(bytes32 accountRefId, address smartAccount, bytes32 provider)
        external
        onlyRole(ACCOUNT_MANAGER_ROLE)
    {
        require(smartAccount != address(0), "AccountWalletRegistryExtended: zero smartAccount");
        require(_isContract(smartAccount), "AccountWalletRegistryExtended: not a contract");

        bytes32 walletRefId = keccak256(abi.encodePacked(smartAccount));
        this.linkAccountToWallet(accountRefId, walletRefId, provider);
        _smartAccountAddress[walletRefId] = smartAccount;
        emit SmartAccountLinked(accountRefId, smartAccount, provider);
    }

    function unlinkAccountFromWallet(bytes32 accountRefId, bytes32 walletRefId) external override onlyRole(ACCOUNT_MANAGER_ROLE) {
        require(_walletAccountExists[walletRefId][accountRefId], "not linked");
        uint256 index = _walletIndex[accountRefId][walletRefId];
        require(index < _accountToWallets[accountRefId].length, "index out of bounds");
        _accountToWallets[accountRefId][index].active = false;
        _walletAccountExists[walletRefId][accountRefId] = false;
        emit AccountWalletUnlinked(accountRefId, walletRefId);
    }

    function getWallets(bytes32 accountRefId) external view override returns (WalletLink[] memory) {
        return _accountToWallets[accountRefId];
    }

    function getAccounts(bytes32 walletRefId) external view override returns (bytes32[] memory) {
        return _walletToAccounts[walletRefId];
    }

    function isLinked(bytes32 accountRefId, bytes32 walletRefId) external view override returns (bool) {
        return _walletAccountExists[walletRefId][accountRefId];
    }

    function isActive(bytes32 accountRefId, bytes32 walletRefId) external view override returns (bool) {
        if (!_walletAccountExists[walletRefId][accountRefId]) return false;
        uint256 index = _walletIndex[accountRefId][walletRefId];
        return index < _accountToWallets[accountRefId].length && _accountToWallets[accountRefId][index].active;
    }

    function isSmartAccount(bytes32 walletRefId) public view returns (bool) {
        return _smartAccountAddress[walletRefId] != address(0);
    }

    function isSmartAccountAddress(address addr) public view returns (bool) {
        return _smartAccountAddress[keccak256(abi.encodePacked(addr))] == addr;
    }

    function _isContract(address account) internal view returns (bool) {
        return account.code.length > 0;
    }
}