tasks.py

import logging
from itertools import batched
from typing import Sequence, cast

from eth_typing import HexStr
from sw_utils import IpfsFetchClient, convert_to_mgno
from sw_utils.networks import GNO_NETWORKS
from web3 import Web3
from web3.types import BlockNumber, Gwei

from src.common.clients import execution_client
from src.common.contracts import VaultContract, validators_registry_contract
from src.common.execution import check_gas_price
from src.common.harvest import get_harvest_params
from src.common.metrics import metrics
from src.common.protocol_config import get_protocol_config
from src.common.typings import HarvestParams, ValidatorType
from src.config.settings import (
    MIN_ACTIVATION_BALANCE_GWEI,
    VALIDATORS_FUNDING_BATCH_SIZE,
    settings,
)
from src.validators.consensus import fetch_compounding_validators_balances
from src.validators.database import NetworkValidatorCrud
from src.validators.exceptions import EmptyRelayerResponseException, FundingException
from src.validators.execution import get_withdrawable_assets
from src.validators.keystores.base import BaseKeystore
from src.validators.metrics import update_unused_validator_keys_metric
from src.validators.oracles import poll_validation_approval
from src.validators.register_validators import fund_validators, register_validators
from src.validators.relayer import RelayerClient
from src.validators.typings import NetworkValidator, Validator
from src.validators.utils import get_validators_for_registration
from src.validators.validators_manager import get_validators_manager_signature

logger = logging.getLogger(__name__)


class ValidatorRegistrationSubtask:
    def __init__(
        self,
        keystore: BaseKeystore | None,
        relayer: RelayerClient | None,
    ):
        self.keystore = keystore
        self.relayer = relayer

    async def process(self) -> None:
        if self.keystore:
            await update_unused_validator_keys_metric(
                keystore=self.keystore,
            )
        harvest_params = await get_harvest_params(settings.vault)

        vault_assets = await get_vault_assets(harvest_params=harvest_params)
        vault_assets = Gwei(max(0, vault_assets - settings.vault_min_balance_gwei))

        if vault_assets < settings.min_deposit_amount_gwei:
            return

        if not await check_gas_price():
            return

        if settings.validator_type == ValidatorType.V1:
            if not settings.disable_validators_registration:
                await register_new_validators(
                    vault_assets=vault_assets,
                    harvest_params=harvest_params,
                    keystore=self.keystore,
                    relayer=self.relayer,
                )
            return

        if not settings.disable_validators_funding:
            try:
                vault_assets = await self.process_funding(
                    vault_assets=vault_assets,
                    harvest_params=harvest_params,
                )
            except FundingException as e:
                logger.warning('Funding failed: %s', e)
                return

        if not settings.disable_validators_registration:
            await register_new_validators(
                vault_assets=vault_assets,
                harvest_params=harvest_params,
                keystore=self.keystore,
                relayer=self.relayer,
            )

    async def process_funding(
        self,
        vault_assets: Gwei,
        harvest_params: HarvestParams | None,
    ) -> Gwei:
        """
        Processes funding compounding validators.
        Returns the remaining vault assets if funding is successful.
        Raises FundingException on failure.
        """
        compounding_validators_balances = await fetch_compounding_validators_balances()
        funding_amounts = _get_funding_amounts(
            compounding_validators_balances=compounding_validators_balances,
            vault_assets=vault_assets,
        )

        if not funding_amounts:
            return vault_assets

        if not await _is_funding_interval_passed():
            raise FundingException('Funding interval has not passed yet')

        for validator_fundings_chunk in batched(
            list(funding_amounts.items()), VALIDATORS_FUNDING_BATCH_SIZE
        ):
            try:
                tx_hash = await fund_compounding_validators(
                    validator_fundings=validator_fundings_chunk,
                    harvest_params=harvest_params,
                    relayer=self.relayer,
                )
            except EmptyRelayerResponseException as e:
                raise FundingException('Empty response from relayer') from e

            if not tx_hash:
                raise FundingException('Funding transaction failed')

            total_funded = sum(amount for _, amount in validator_fundings_chunk)
            vault_assets = Gwei(max(vault_assets - total_funded, 0))

        return vault_assets


async def fund_compounding_validators(
    validator_fundings: Sequence[tuple[HexStr, Gwei]],
    harvest_params: HarvestParams | None,
    relayer: RelayerClient | None = None,
) -> HexStr | None:
    """
    Funds vault compounding validators with the specified amount.
    """
    logger.info('Started funding of %d validator(s)', len(validator_fundings))
    validators_manager_signature = HexStr('0x')
    if settings.relayer_endpoint:
        # fetch validators and signature from relayer
        validators_response = await cast(RelayerClient, relayer).fund_validators(
            validator_fundings=validator_fundings,
        )

        if validators_response.validators_manager_signature:
            validators_manager_signature = validators_response.validators_manager_signature

    validators = []
    # the signature is not checked for funding active validators
    empty_signature = Web3.to_hex(bytes(96))
    pub_keys = []
    for public_key, amount in validator_fundings:
        validators.append(
            Validator(
                public_key=public_key,
                deposit_signature=empty_signature,
                amount=amount,
            )
        )
        pub_keys.append(public_key)

    tx_hash = await fund_validators(
        harvest_params=harvest_params,
        validators=validators,
        validators_manager_signature=validators_manager_signature,
    )
    if tx_hash:
        logger.info('Successfully funded validator(s) with public key(s) %s', ', '.join(pub_keys))
        tx_data = await execution_client.eth.get_transaction(tx_hash)
        metrics.last_funding_block.labels(network=settings.network).set(tx_data['blockNumber'])
    return tx_hash


# pylint: disable-next=too-many-locals
async def register_new_validators(
    vault_assets: Gwei,
    harvest_params: HarvestParams | None,
    keystore: BaseKeystore | None,
    relayer: RelayerClient | None = None,
) -> HexStr | None:
    validators_amounts = _get_deposits_amounts(vault_assets, settings.validator_type)
    validators_count = len(validators_amounts)
    if not validators_amounts:
        # not enough balance to register validators
        return None

    # get latest config
    protocol_config = await get_protocol_config()

    validators_batch_size = min(protocol_config.validators_approval_batch_limit, validators_count)
    validators: Sequence[Validator]
    validators_registry_root = await validators_registry_contract.get_registry_root()
    validators_manager_signature: HexStr | None = None

    if not settings.relayer_endpoint:
        validators = await get_validators_for_registration(
            keystore=cast(BaseKeystore, keystore),
            amounts=validators_amounts[:validators_batch_size],
        )
        if not validators:
            if not settings.disable_available_validators_warnings:
                logger.warning(
                    'There are no available public keys '
                    'to proceed with registration. '
                    'To register additional validators, you must generate new keystores.',
                )
            return None
    else:
        validators_response = await cast(RelayerClient, relayer).register_validators(
            amounts=validators_amounts[:validators_batch_size],
        )
        validators = validators_response.validators
        validators_manager_signature = validators_response.validators_manager_signature

        if not validators or not validators_manager_signature:
            # Missing signature indicates that relayer is not ready with the validators,
            # wait for the next cycle
            logger.info('Waiting for relayer validators')
            return None

    if not validators_manager_signature:
        validators_manager_signature = get_validators_manager_signature(
            vault=settings.vault,
            validators_registry_root=validators_registry_root,
            validators=validators,
        )

    if not await check_gas_price(high_priority=True):
        return None

    logger.info(
        'Started registration of %d %s validator(s)', len(validators), settings.validator_type.value
    )

    oracles_request, oracles_approval = await poll_validation_approval(
        keystore=keystore,
        validators=validators,
        validators_registry_root=validators_registry_root,
        validators_manager_signature=validators_manager_signature,
    )
    tx_hash = await register_validators(
        approval=oracles_approval,
        validators=validators,
        harvest_params=harvest_params,
        validators_registry_root=oracles_request.validators_root,
        validator_index=oracles_request.validator_index,
        validators_manager_signature=validators_manager_signature,
    )
    if tx_hash:
        pub_keys = ', '.join([val.public_key for val in validators])
        logger.info('Successfully registered validator(s) with public key(s) %s', pub_keys)
        tx_data = await execution_client.eth.get_transaction(tx_hash)
        metrics.last_registration_block.labels(network=settings.network).set(tx_data['blockNumber'])

    return tx_hash


async def get_vault_assets(harvest_params: HarvestParams | None) -> Gwei:
    vault_assets = await get_withdrawable_assets(settings.vault, harvest_params=harvest_params)
    if settings.network in GNO_NETWORKS:
        # apply GNO -> mGNO exchange rate
        vault_assets = convert_to_mgno(vault_assets)

    metrics.stakeable_assets.labels(network=settings.network).set(int(vault_assets))

    return Gwei(int(Web3.from_wei(vault_assets, 'gwei')))


async def load_genesis_validators() -> None:
    """
    Load consensus network validators from the ipfs dump.
    Used to speed up service startup
    """
    ipfs_hash = settings.network_config.GENESIS_VALIDATORS_IPFS_HASH
    if not (NetworkValidatorCrud().get_last_network_validator() is None and ipfs_hash):
        return

    logger.info('Downloading validators data from IPFS...')
    ipfs_fetch_client = IpfsFetchClient(
        ipfs_endpoints=settings.ipfs_fetch_endpoints,
        timeout=settings.genesis_validators_ipfs_timeout,
        retry_timeout=settings.genesis_validators_ipfs_retry_timeout,
    )
    data = await ipfs_fetch_client.fetch_bytes(ipfs_hash)
    genesis_validators: list[NetworkValidator] = []
    logger.info('Loading genesis validators into the database...')
    for i in range(0, len(data), 52):
        genesis_validators.append(
            NetworkValidator(
                public_key=Web3.to_hex(data[i + 4 : i + 52]),
                block_number=BlockNumber(int.from_bytes(data[i : i + 4], 'big')),
            )
        )

    NetworkValidatorCrud().save_network_validators(genesis_validators)
    logger.info('Loaded %d genesis validators', len(genesis_validators))


def _get_deposits_amounts(vault_assets: Gwei, validator_type: ValidatorType) -> list[Gwei]:
    """Returns a list of amounts in Gwei for each validator to be registered."""
    if vault_assets < MIN_ACTIVATION_BALANCE_GWEI:
        return []
    if validator_type == ValidatorType.V1:
        return [MIN_ACTIVATION_BALANCE_GWEI] * (vault_assets // MIN_ACTIVATION_BALANCE_GWEI)
    amounts = []
    num_full_validators, remainder = divmod(vault_assets, settings.max_validator_balance_gwei)
    amounts.extend([settings.max_validator_balance_gwei] * num_full_validators)
    if remainder >= MIN_ACTIVATION_BALANCE_GWEI:
        amounts.append(Gwei(remainder))

    return amounts


def _get_funding_amounts(
    compounding_validators_balances: dict[HexStr, Gwei], vault_assets: Gwei
) -> dict[HexStr, Gwei]:
    result = {}
    for public_key, balance in sorted(
        compounding_validators_balances.items(), key=lambda item: item[1], reverse=True
    ):
        remaining_capacity = settings.max_validator_balance_gwei - balance
        if remaining_capacity >= settings.min_deposit_amount_gwei:
            val_amount = min(remaining_capacity, vault_assets)
            result[public_key] = Gwei(val_amount)
            vault_assets = Gwei(vault_assets - val_amount)
        if vault_assets < settings.min_deposit_amount_gwei:
            break
    return result


async def _is_funding_interval_passed() -> bool:
    """
    Check if the required interval has passed since the last funding event.
    Mitigate gas griefing attack
    """
    blocks_delay = settings.min_deposit_delay // settings.network_config.SECONDS_PER_BLOCK
    to_block = await execution_client.eth.get_block_number()
    from_block = BlockNumber(to_block - blocks_delay)
    funding_events = await VaultContract(settings.vault).get_funding_events(from_block, to_block)
    return len(funding_events) == 0