espresso_node/request_response/catchup/

state.rs

use std::sync::Arc;

use alloy::primitives::U256;
use anyhow::Context;
use async_trait::async_trait;
use committable::{Commitment, Committable};
use espresso_types::{
    BackoffParams, BlockMerkleTree, FeeAccount, FeeAccountProof, FeeMerkleCommitment, Leaf2,
    NodeState, PubKey, SeqTypes,
    traits::{SequencerPersistence, StateCatchup},
    v0_3::{ChainConfig, RewardAccountProofV1, RewardAccountV1, RewardMerkleCommitmentV1},
    v0_4::{
        PermittedRewardMerkleTreeV2, RewardAccountV2, RewardMerkleCommitmentV2,
        forgotten_accounts_include,
    },
};
use hotshot::traits::NodeImplementation;
use hotshot_types::{
    data::ViewNumber, message::UpgradeLock,
    simple_certificate::LightClientStateUpdateCertificateV2, stake_table::HSStakeTable,
    traits::network::ConnectedNetwork, utils::verify_leaf_chain,
};
use jf_merkle_tree_compat::{ForgetableMerkleTreeScheme, MerkleTreeScheme};
use request_response::RequestType;
use tokio::time::timeout;
use versions::EPOCH_VERSION;

use crate::{
    api::RewardMerkleTreeV2Data,
    request_response::{
        RequestResponseProtocol,
        request::{Request, Response},
    },
};

#[async_trait]
impl<I: NodeImplementation<SeqTypes>, N: ConnectedNetwork<PubKey>, P: SequencerPersistence>
    StateCatchup for RequestResponseProtocol<I, N, P>
{
    async fn try_fetch_leaf(
        &self,
        _retry: usize,
        height: u64,
        stake_table: HSStakeTable<SeqTypes>,
        success_threshold: U256,
    ) -> anyhow::Result<Leaf2> {
        // Timeout after a few batches
        let timeout_duration = self.config.request_batch_interval * 3;

        // Fetch the leaf
        timeout(
            timeout_duration,
            self.fetch_leaf(height, stake_table, success_threshold),
        )
        .await
        .with_context(|| "timed out while fetching leaf")?
    }

    async fn try_fetch_accounts(
        &self,
        _retry: usize,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        fee_merkle_tree_root: FeeMerkleCommitment,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        // Timeout after a few batches
        let timeout_duration = self.config.request_batch_interval * 3;

        // Fetch the accounts
        timeout(
            timeout_duration,
            self.fetch_accounts(
                instance,
                height,
                view,
                fee_merkle_tree_root,
                accounts.to_vec(),
            ),
        )
        .await
        .with_context(|| "timed out while fetching accounts")?
    }

    async fn try_remember_blocks_merkle_tree(
        &self,
        _retry: usize,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        // Timeout after a few batches
        let timeout_duration = self.config.request_batch_interval * 3;

        // Remember the blocks merkle tree
        timeout(
            timeout_duration,
            self.remember_blocks_merkle_tree(instance, height, view, mt),
        )
        .await
        .with_context(|| "timed out while remembering blocks merkle tree")?
    }

    async fn try_fetch_chain_config(
        &self,
        _retry: usize,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        // Timeout after a few batches
        let timeout_duration = self.config.request_batch_interval * 3;

        // Fetch the chain config
        timeout(timeout_duration, self.fetch_chain_config(commitment))
            .await
            .with_context(|| "timed out while fetching chain config")?
    }

    async fn try_fetch_reward_merkle_tree_v2(
        &self,
        _retry: usize,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV2,
        accounts: Arc<Vec<RewardAccountV2>>,
    ) -> anyhow::Result<PermittedRewardMerkleTreeV2> {
        // Timeout after a few batches
        let timeout_duration = self.config.request_batch_interval * 3;

        // Fetch the reward accounts
        timeout(
            timeout_duration,
            self.fetch_reward_merkle_tree_v2(height, view, reward_merkle_tree_root, accounts),
        )
        .await
        .with_context(|| "timed out while fetching reward merkle tree v2")?
    }

    async fn try_fetch_reward_accounts_v1(
        &self,
        _retry: usize,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV1,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        // Timeout after a few batches
        let timeout_duration = self.config.request_batch_interval * 3;

        // Fetch the reward accounts
        timeout(
            timeout_duration,
            self.fetch_reward_accounts_v1(
                instance,
                height,
                view,
                reward_merkle_tree_root,
                accounts.to_vec(),
            ),
        )
        .await
        .with_context(|| "timed out while fetching reward accounts")?
    }

    async fn try_fetch_state_cert(
        &self,
        _retry: usize,
        epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        let timeout_duration = self.config.request_batch_interval * 3;

        timeout(timeout_duration, self.fetch_state_cert(epoch))
            .await
            .with_context(|| "timed out while fetching state cert")?
    }

    fn backoff(&self) -> &BackoffParams {
        unreachable!()
    }

    fn name(&self) -> String {
        "request-response".to_string()
    }

    async fn fetch_accounts(
        &self,
        _instance: &NodeState,
        height: u64,
        view: ViewNumber,
        fee_merkle_tree_root: FeeMerkleCommitment,
        accounts: Vec<FeeAccount>,
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        tracing::info!("Fetching accounts for height: {height}, view: {view}");

        // Clone things we need in the first closure
        let accounts_clone = accounts.clone();
        let response_validation_fn = move |_request: &Request, response: Response| {
            // Clone again
            let accounts_clone = accounts_clone.clone();

            async move {
                // Make sure the response is an accounts response
                let Response::Accounts(fee_merkle_tree) = response else {
                    return Err(anyhow::anyhow!("expected accounts response"));
                };

                // Verify the merkle proofs
                let mut proofs = Vec::new();
                for account in accounts_clone {
                    let (proof, _) = FeeAccountProof::prove(&fee_merkle_tree, account.into())
                        .with_context(|| format!("response was missing account {account}"))?;
                    proof.verify(&fee_merkle_tree_root).with_context(|| {
                        format!(
                            "invalid proof for fee account {account}, root: {fee_merkle_tree_root}"
                        )
                    })?;
                    proofs.push(proof);
                }

                Ok(proofs)
            }
        };

        // Wait for the protocol to send us the accounts
        let response = self
            .request_indefinitely(
                Request::Accounts(height, *view, accounts),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request accounts")?;

        tracing::info!("Fetched accounts for height: {height}, view: {view}");

        Ok(response)
    }

    async fn fetch_leaf(
        &self,
        height: u64,
        stake_table: HSStakeTable<SeqTypes>,
        success_threshold: U256,
    ) -> anyhow::Result<Leaf2> {
        tracing::info!("Fetching leaf for height: {height}");

        // Clone things we need in the first closure
        let stake_table_clone = stake_table.clone();
        let response_validation_fn = move |_request: &Request, response: Response| {
            // Clone again
            let stake_table_clone = stake_table_clone.clone();

            async move {
                // Make sure the response is a leaf response
                let Response::Leaf(leaf_chain) = response else {
                    return Err(anyhow::anyhow!("expected leaf response"));
                };

                // Verify the leaf chain
                let leaf = verify_leaf_chain(
                    leaf_chain,
                    &stake_table_clone,
                    success_threshold,
                    height,
                    &UpgradeLock::<SeqTypes>::new(versions::Upgrade::trivial(EPOCH_VERSION)),
                )
                .await
                .with_context(|| "leaf chain verification failed")?;

                Ok(leaf)
            }
        };

        // Wait for the protocol to send us the accounts
        let response = self
            .request_indefinitely(
                Request::Leaf(height),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request leaf")?;

        tracing::info!("Fetched leaf for height: {height}");

        Ok(response)
    }

    async fn fetch_chain_config(
        &self,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        tracing::info!("Fetching chain config with commitment: {commitment}");

        // Create the response validation function
        let response_validation_fn = move |_request: &Request, response: Response| {
            async move {
                // Make sure the response is a chain config response
                let Response::ChainConfig(chain_config) = response else {
                    return Err(anyhow::anyhow!("expected chain config response"));
                };

                // Make sure the commitments match
                if commitment != chain_config.commit() {
                    return Err(anyhow::anyhow!("chain config commitment mismatch"));
                }

                Ok(chain_config)
            }
        };

        // Wait for the protocol to send us the chain config
        let response = self
            .request_indefinitely(
                Request::ChainConfig(commitment),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request chain config")?;

        tracing::info!("Fetched chain config with commitment: {commitment}");

        Ok(response)
    }

    async fn remember_blocks_merkle_tree(
        &self,
        _instance: &NodeState,
        height: u64,
        view: ViewNumber,
        mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        tracing::info!("Fetching blocks frontier for height: {height}, view: {view}");

        // Clone the merkle tree
        let mt_clone = mt.clone();

        // Create the response validation function
        let response_validation_fn = move |_request: &Request, response: Response| {
            // Clone the merkle tree
            let mut block_merkle_tree = mt_clone.clone();

            async move {
                // Make sure the response is a blocks frontier response
                let Response::BlocksFrontier(blocks_frontier) = response else {
                    return Err(anyhow::anyhow!("expected blocks frontier response"));
                };

                // Get the leaf element associated with the proof
                let leaf_elem = blocks_frontier
                    .elem()
                    .with_context(|| "provided frontier is missing leaf element")?;

                // Verify the block proof
                block_merkle_tree
                    .remember(
                        block_merkle_tree.num_leaves() - 1,
                        *leaf_elem,
                        blocks_frontier,
                    )
                    .with_context(|| "merkle tree verification failed")?;

                // Return the verified merkle tree
                Ok(block_merkle_tree)
            }
        };

        // Wait for the protocol to send us the blocks frontier
        let response = self
            .request_indefinitely(
                Request::BlocksFrontier(height, *view),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request blocks frontier")?;

        // Replace the merkle tree
        *mt = response;

        tracing::info!("Fetched blocks frontier for height: {height}, view: {view}");

        Ok(())
    }

    async fn fetch_reward_merkle_tree_v2(
        &self,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV2,
        accounts: Arc<Vec<RewardAccountV2>>,
    ) -> anyhow::Result<PermittedRewardMerkleTreeV2> {
        tracing::info!("Fetching RewardMerkleTreeV2 for height: {height}");

        // Create the response validation function
        let response_validation_fn = move |_request: &Request, response: Response| {
            let accounts = accounts.clone();
            async move {
                // Make sure the response is a reward accounts response
                let Response::RewardMerkleTreeV2(tree_bytes) = response else {
                    return Err(anyhow::anyhow!("expected reward accounts response"));
                };

                let tree_data = bincode::deserialize::<RewardMerkleTreeV2Data>(&tree_bytes)
                    .context(
                        "Failed to deserialize RewardMerkleTreeV2 for height {height} from remote",
                    )?;

                // Verify the tree's commitment
                let reward_merkle_tree: PermittedRewardMerkleTreeV2 =
                    PermittedRewardMerkleTreeV2::try_from_kv_set(tree_data.balances).await?;

                anyhow::ensure!(reward_merkle_tree.tree.commitment() == reward_merkle_tree_root);
                anyhow::ensure!(!forgotten_accounts_include(&reward_merkle_tree, &accounts));

                Ok(reward_merkle_tree)
            }
        };

        // Wait for the protocol to send us the reward accounts
        let response = self
            .request_indefinitely(
                Request::RewardMerkleTreeV2(height, *view),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request reward accounts")?;

        tracing::info!("Fetched RewardMerkleTreeV2 for height: {height}");

        Ok(response)
    }

    async fn fetch_reward_accounts_v1(
        &self,
        _instance: &NodeState,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV1,
        accounts: Vec<RewardAccountV1>,
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        tracing::info!("Fetching v1 reward accounts for height: {height}, view: {view}");

        // Clone things we need in the first closure
        let accounts_clone = accounts.clone();

        // Create the response validation function
        let response_validation_fn = move |_request: &Request, response: Response| {
            // Clone again
            let accounts_clone = accounts_clone.clone();

            async move {
                // Make sure the response is a reward accounts response
                let Response::RewardAccountsV1(reward_merkle_tree) = response else {
                    return Err(anyhow::anyhow!("expected v1 reward accounts response"));
                };

                // Verify the merkle proofs
                let mut proofs = Vec::new();
                for account in accounts_clone {
                    let (proof, _) =
                        RewardAccountProofV1::prove(&reward_merkle_tree, account.into())
                            .with_context(|| format!("response was missing account {account}"))?;
                    proof.verify(&reward_merkle_tree_root).with_context(|| {
                        format!(
                            "invalid proof for v1 reward account {account}, root: \
                             {reward_merkle_tree_root}"
                        )
                    })?;
                    proofs.push(proof);
                }

                Ok(proofs)
            }
        };

        // Wait for the protocol to send us the reward accounts
        let response = self
            .request_indefinitely(
                Request::RewardAccountsV1(height, *view, accounts),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request v1 reward accounts")?;

        tracing::info!("Fetched v1 reward accounts for height: {height}, view: {view}");

        Ok(response)
    }

    async fn fetch_state_cert(
        &self,
        epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        tracing::info!("Fetching state cert for epoch: {epoch}");

        // Create the response validation function
        let response_validation_fn = move |_request: &Request, response: Response| async move {
            // Make sure the response is a state cert response
            let Response::StateCert(state_cert) = response else {
                return Err(anyhow::anyhow!("expected state cert response"));
            };

            Ok(state_cert)
        };

        // Wait for the protocol to send us the state cert
        let response = self
            .request_indefinitely(
                Request::StateCert(epoch),
                RequestType::Batched,
                response_validation_fn,
            )
            .await
            .with_context(|| "failed to request state cert")?;

        tracing::info!("Fetched state cert for epoch: {epoch}");

        Ok(response)
    }

    fn is_local(&self) -> bool {
        false
    }
}