espresso_node/api/

light_client.rs

use std::{cmp::Ordering, collections::HashMap, fmt::Display, sync::Arc, time::Duration};

use anyhow::{Result, bail, ensure};
use committable::{Commitment, Committable};
use disco_types::status::StatusCode;
use espresso_types::{BlockMerkleTree, Header, NsIndex, NsProof, SeqTypes};
use futures::{TryStreamExt, future::try_join, stream::StreamExt};
use hotshot_query_service::{
    Error, QueryError,
    availability::{AvailabilityDataSource, LeafQueryData, PayloadQueryData, VidCommonQueryData},
    data_source::{VersionedDataSource, storage::NodeStorage},
    merklized_state::{MerklizedStateDataSource, MerklizedStateHeightPersistence, Snapshot},
    node::BlockId,
    types::HeightIndexed,
};
use hotshot_types::simple_certificate::QuorumCertificate2;
use itertools::izip;
use jf_merkle_tree_compat::{
    DigestAlgorithm, MerkleTreeScheme,
    internal::MerkleNode,
    prelude::{Sha3Digest, Sha3Node},
};
use light_client::{
    client::NAMESPACES_PARAM_TAG,
    consensus::{header::HeaderProof, leaf::LeafProof, namespace::NamespaceProof},
};
use tagged_base64::TaggedBase64;
use tokio::spawn;
use versions::NEW_PROTOCOL_VERSION;

/// Construct a proof that the requested leaf is finalized.
///
/// The `finalized` hint is honored when it yields a bounded proof no longer than a direct finality
/// proof, since the client can verify hint-based proofs without signature checks.
pub(crate) async fn get_leaf_proof<State>(
    state: &State,
    requested_leaf: LeafQueryData<SeqTypes>,
    finalized: Option<usize>,
    fetch_timeout: Duration,
    chain_limit: usize,
) -> Result<LeafProof, Error>
where
    State: AvailabilityDataSource<SeqTypes> + VersionedDataSource + Clone + Send + Sync + 'static,
    for<'a> State::ReadOnly<'a>: NodeStorage<SeqTypes>,
{
    let requested = requested_leaf.height() as usize;
    let new_protocol = requested_leaf.header().version() >= NEW_PROTOCOL_VERSION;

    let mut hint = finalized.filter(|finalized| finalized.saturating_sub(requested) <= chain_limit);

    // New-protocol chains cannot terminate early, so a nearby cert2 may yield a shorter proof.
    if let Some(finalized) = hint
        && new_protocol
        && let Some(cert2_height) = earliest_cert2_height(state, requested as u64).await
        && cert2_height + 1 < finalized as u64
    {
        hint = None;
    }

    if let Some(finalized) = hint {
        get_leaf_proof_with_finalized_assumption(
            state,
            requested_leaf,
            finalized,
            fetch_timeout,
            chain_limit,
        )
        .await
    } else if new_protocol {
        get_leaf_proof_with_cert2(state, requested_leaf, fetch_timeout, chain_limit).await
    } else {
        get_leaf_proof_with_leaf_chain(state, requested_leaf, fetch_timeout, chain_limit).await
    }
}

/// The earliest stored cert2 height at or after `height`. Errors are treated as no cert2: this
/// only chooses between proof strategies, and the chosen path will surface any persistent error.
async fn earliest_cert2_height<State>(state: &State, height: u64) -> Option<u64>
where
    State: VersionedDataSource,
    for<'a> State::ReadOnly<'a>: NodeStorage<SeqTypes>,
{
    let cert2 = state
        .read()
        .await
        .ok()?
        .load_earliest_cert2(height)
        .await
        .ok()??;
    Some(cert2.data.block_number)
}

/// Build a leaf proof for a legacy leaf from the justify QCs of the leaves that follow it.
pub(crate) async fn get_leaf_proof_with_leaf_chain<State>(
    state: &State,
    requested_leaf: LeafQueryData<SeqTypes>,
    fetch_timeout: Duration,
    chain_limit: usize,
) -> Result<LeafProof, Error>
where
    State: AvailabilityDataSource<SeqTypes> + VersionedDataSource + Clone + Send + Sync + 'static,
    for<'a> State::ReadOnly<'a>: NodeStorage<SeqTypes>,
{
    let requested = requested_leaf.height() as usize;
    let mut tx = state.read().await.map_err(internal)?;
    let latest_height = NodeStorage::block_height(&mut tx).await.map_err(internal)?;
    drop(tx);

    let mut leaves = state.get_leaf_range(requested + 1..latest_height).await;
    let mut proof = LeafProof::default();
    let requested_leaf_qc = requested_leaf.qc().clone();
    proof.push(requested_leaf);

    let mut remaining = chain_limit;
    while let Some(leaf) = leaves.next().await {
        let leaf = leaf
            .with_timeout(fetch_timeout)
            .await
            .ok_or_else(|| not_found("missing leaves"))?;

        remaining = remaining
            .checked_sub(1)
            .ok_or_else(|| chain_too_long(requested, chain_limit))?;

        // HotStuff commit rules cannot terminate a chain of new-protocol leaves; switch to cert2.
        if leaf.header().version() >= NEW_PROTOCOL_VERSION {
            complete_proof_with_cert2(
                state,
                &mut proof,
                leaf,
                requested_leaf_qc.clone(),
                fetch_timeout,
                remaining,
            )
            .await?;
            return Ok(proof);
        }

        if proof.push(leaf) {
            return Ok(proof);
        }
    }

    // There is no stored legacy QC chain to fall back on: it could only close a chain whose tip was
    // legacy, and every network is past the 0.6 cutover, so running out of leaves is a 404.
    Err(not_found("not enough leaves to prove finality"))
}

/// Build a leaf proof for the new protocol using certificate2 finality.
///
/// Certificate2 directly commits the leaf at `cert2.data.block_number`. The indirect commit rule then
/// commits all of that leaf's uncommitted ancestors. If cert2 directly commits the requested leaf,
/// the proof contains only that leaf plus cert2. Otherwise, the proof includes the chain from the
/// requested leaf through the directly committed descendant; the verifier checks parent commitments
/// to prove the requested leaf is on that finalized chain.
pub(crate) async fn get_leaf_proof_with_cert2<State>(
    state: &State,
    requested_leaf: LeafQueryData<SeqTypes>,
    fetch_timeout: Duration,
    chain_limit: usize,
) -> Result<LeafProof, Error>
where
    State: AvailabilityDataSource<SeqTypes> + VersionedDataSource + Clone + Send + Sync + 'static,
    for<'a> State::ReadOnly<'a>: NodeStorage<SeqTypes>,
{
    let mut proof = LeafProof::default();
    let requested_leaf_qc = requested_leaf.qc().clone();
    complete_proof_with_cert2(
        state,
        &mut proof,
        requested_leaf,
        requested_leaf_qc,
        fetch_timeout,
        chain_limit,
    )
    .await?;
    Ok(proof)
}

/// Extend `proof` from `leaf` to the leaf directly committed by the earliest stored cert2, adding
/// at most `chain_limit` leaves.
async fn complete_proof_with_cert2<State>(
    state: &State,
    proof: &mut LeafProof,
    leaf: LeafQueryData<SeqTypes>,
    requested_leaf_qc: QuorumCertificate2<SeqTypes>,
    fetch_timeout: Duration,
    chain_limit: usize,
) -> Result<(), Error>
where
    State: AvailabilityDataSource<SeqTypes> + VersionedDataSource + Clone + Send + Sync + 'static,
    for<'a> State::ReadOnly<'a>: NodeStorage<SeqTypes>,
{
    let start_height = leaf.height();
    let start_commit = leaf.leaf().commit();
    let limit = chain_limit as u64;

    // The new leaf may already complete a HotStuff chain begun before the protocol cutover.
    if proof.push(leaf) {
        return Ok(());
    }

    let stored = state
        .read()
        .await
        .map_err(internal)?
        .load_earliest_cert2(start_height)
        .await
        .map_err(internal)?;

    let earliest = stored.as_ref().map(|cert2| cert2.data.block_number);
    let Some(cert2) = stored.filter(|cert2| cert2.data.block_number <= start_height + limit) else {
        // An indirectly finalized leaf has no cert2 of its own, so its descendants are asked, one
        // at a time. A cert2 no peer has never resolves, hence the timeout per step.
        let known_height = async {
            let mut tx = state.read().await?;
            anyhow::Ok(NodeStorage::block_height(&mut tx).await? as u64)
        }
        .await
        // A miss either way, and a failing store should not be asked once more per height.
        .inspect_err(|err| {
            tracing::warn!(
                %err,
                height = start_height,
                "failed to load block height, skipping cert2 fetch"
            )
        })
        .ok();
        if let Some(known_height) = known_height {
            let end = (start_height + limit + 1).min(known_height);
            let state = state.clone();
            spawn(async move {
                for height in start_height..end {
                    if state
                        .get_cert2(height)
                        .await
                        .with_timeout(fetch_timeout)
                        .await
                        .is_some()
                    {
                        break;
                    }
                }
            });
        }
        return Err(not_found(match earliest {
            None => format!("no cert2 finality proof available at or after height {start_height}"),
            Some(height) => format!(
                "earliest cert2 finality proof (height {height}) is more than {chain_limit} \
                 leaves past height {start_height}"
            ),
        }));
    };

    let cert2_height = cert2.data.block_number;
    if cert2_height < start_height {
        return Err(not_found(
            "cert2 finality proof is older than requested leaf",
        ));
    }

    if cert2_height == start_height {
        if start_commit != cert2.data.leaf_commit {
            return Err(internal("stored cert2 does not finalize the expected leaf"));
        }
        proof.add_certificate(Arc::new(cert2), requested_leaf_qc);
        return Ok(());
    }

    let mut leaves = state
        .get_leaf_range(start_height as usize + 1..cert2_height as usize + 1)
        .await;
    // Extend the proof chain until we reach the leaf directly committed by cert2. Once that leaf is
    // present and matches cert2's commitment, the requested leaf is finalized by the indirect
    // commit rule.
    while let Some(leaf) = leaves.next().await {
        let leaf = leaf
            .with_timeout(fetch_timeout)
            .await
            .ok_or_else(|| not_found("missing leaves"))?;

        if leaf.height() == cert2_height {
            if leaf.leaf().commit() != cert2.data.leaf_commit {
                return Err(internal("stored cert2 does not finalize the expected leaf"));
            }
            if !proof.push(leaf) {
                proof.add_certificate(Arc::new(cert2), requested_leaf_qc);
            }
            return Ok(());
        }
        if proof.push(leaf) {
            return Ok(());
        }
    }

    Err(not_found("missing cert2 leaf"))
}

pub(crate) async fn get_leaf_proof_with_finalized_assumption<State>(
    state: &State,
    requested_leaf: LeafQueryData<SeqTypes>,
    finalized: usize,
    fetch_timeout: Duration,
    chain_limit: usize,
) -> Result<LeafProof, Error>
where
    State: AvailabilityDataSource<SeqTypes>,
{
    let requested = requested_leaf.height() as usize;
    // If we have a known-finalized block, we will not need a final 2-chain of QCs to prove
    // the last leaf in the result finalized, since we will either terminate with a 3-chain
    // of leaves or at the `finalized` leaf.
    if finalized <= requested {
        return Err(Error::Custom {
            message: format!(
                "finalized leaf height ({finalized}) must be greater than requested ({requested})"
            ),
            status: StatusCode::BAD_REQUEST,
        });
    }
    if finalized - requested > chain_limit {
        return Err(Error::Custom {
            message: format!(
                "finalized leaf height ({finalized}) is more than {chain_limit} blocks past the \
                 requested leaf ({requested}); request a proof without the finalized parameter \
                 instead"
            ),
            status: StatusCode::BAD_REQUEST,
        });
    }

    let mut leaves = state.get_leaf_range(requested + 1..finalized).await;
    let mut proof = LeafProof::default();
    proof.push(requested_leaf);

    while let Some(leaf) = leaves.next().await {
        let leaf = leaf
            .with_timeout(fetch_timeout)
            .await
            .ok_or_else(|| not_found("missing leaves"))?;

        if proof.push(leaf) {
            return Ok(proof);
        }
    }

    Ok(proof)
}

pub(crate) async fn get_header_proof<State>(
    state: &State,
    root: u64,
    requested: BlockId<SeqTypes>,
    fetch_timeout: Duration,
) -> Result<HeaderProof, Error>
where
    State: AvailabilityDataSource<SeqTypes>
        + MerklizedStateDataSource<SeqTypes, BlockMerkleTree, { BlockMerkleTree::ARITY }>
        + MerklizedStateHeightPersistence
        + VersionedDataSource,
{
    let header = state
        .get_header(requested)
        .await
        .with_timeout(fetch_timeout)
        .await
        .ok_or_else(|| not_found(format!("unknown header {requested}")))?;
    if header.height() >= root {
        return Err(Error::Custom {
            message: format!(
                "height ({}) must be less than root ({root})",
                header.height()
            ),
            status: StatusCode::BAD_REQUEST,
        });
    }
    let path = MerklizedStateDataSource::<SeqTypes, BlockMerkleTree, _>::get_path(
        state,
        Snapshot::Index(root),
        header.height(),
    )
    .await;
    let path = match path {
        Ok(path) => path,
        Err(err) if matches!(err.status(), StatusCode::NOT_FOUND | StatusCode::GONE) => {
            // No usable snapshot at `root` (pruned or missing): rebuild from the newest one.
            let head = state.get_last_state_height().await.map_err(merkle_error)? as u64;
            if head <= root {
                return Err(merkle_error(err));
            }
            return rebuild_header_path(state, root, header, head, fetch_timeout).await;
        },
        Err(err) => return Err(merkle_error(err)),
    };

    Ok(HeaderProof::new(header, path))
}

fn merkle_error(source: QueryError) -> Error {
    Error::MerklizedState {
        source: source.into(),
    }
}

/// Prove `header` against `root` using the snapshot at `head`, when the state at `root` is pruned.
///
/// The result is verified against the commitment in header `root`.
async fn rebuild_header_path<State>(
    state: &State,
    root: u64,
    header: Header,
    head: u64,
    fetch_timeout: Duration,
) -> Result<HeaderProof, Error>
where
    State: AvailabilityDataSource<SeqTypes>
        + MerklizedStateDataSource<SeqTypes, BlockMerkleTree, { BlockMerkleTree::ARITY }>,
{
    let get_path = |height| {
        MerklizedStateDataSource::<SeqTypes, BlockMerkleTree, _>::get_path(
            state,
            Snapshot::Index(head),
            height,
        )
    };
    let frontier = get_path(root - 1).await.map_err(merkle_error)?;
    let target = get_path(header.height()).await.map_err(merkle_error)?;
    let root_header = state
        .get_header(BlockId::Number(root as usize))
        .await
        .with_timeout(fetch_timeout)
        .await
        .ok_or_else(|| not_found(format!("unknown header {root}")))?;

    let proof = rebuild_block_proof(root, &frontier, &target).map_err(internal)?;
    let proof = HeaderProof::new(header, proof);
    proof
        .verify_ref(root_header.block_merkle_tree_root())
        .map_err(|err| {
            internal(format!(
                "rebuilt block merkle path does not match root {root}: {err:#}"
            ))
        })?;
    Ok(proof)
}

type BlockMerkleProof = <BlockMerkleTree as MerkleTreeScheme>::MembershipProof;
type BlockMerkleNode = MerkleNode<Commitment<Header>, u64, Sha3Node>;

/// The path to `target.pos` in the block tree of the first `num_leaves` leaves.
///
/// Both proofs come from one snapshot with more than `num_leaves` leaves; `frontier` is for leaf
/// `num_leaves - 1`. Subtrees left of the frontier are full and unchanged since, subtrees right of
/// it are empty, and the frontier nodes are recomputed.
pub(crate) fn rebuild_block_proof(
    num_leaves: u64,
    frontier: &BlockMerkleProof,
    target: &BlockMerkleProof,
) -> Result<BlockMerkleProof> {
    ensure!(
        frontier.pos + 1 == num_leaves,
        "frontier proof is for leaf {}, expected {}",
        frontier.pos,
        num_leaves.wrapping_sub(1),
    );
    ensure!(target.pos < num_leaves, "target leaf is beyond the tree");
    ensure!(
        frontier.proof.len() == target.proof.len(),
        "proofs have different heights"
    );

    // Value of the frontier node at each level, starting with the leaf.
    let mut edge_values = vec![frontier.proof[0].value()];
    for (level, node) in frontier.proof.iter().enumerate().skip(1) {
        let edge = frontier.pos / ARITY.pow(level as u32 - 1);
        let data = child_values(node, edge / ARITY, edge, edge_values[level - 1])?;
        edge_values.push(<Sha3Digest as DigestAlgorithm<
            Commitment<Header>,
            u64,
            _,
        >>::digest(&data)?);
    }

    let mut proof = vec![target.proof[0].clone()];
    for (level, node) in target.proof.iter().enumerate().skip(1) {
        let divisor = ARITY.pow(level as u32 - 1);
        let edge = frontier.pos / divisor;
        let data = child_values(
            node,
            target.pos / divisor / ARITY,
            edge,
            edge_values[level - 1],
        )?;
        let children = data
            .iter()
            .enumerate()
            .map(|(n, value)| {
                let index = target.pos / divisor / ARITY * ARITY + n as u64;
                Arc::new(match index.cmp(&edge) {
                    Ordering::Greater => BlockMerkleNode::Empty,
                    _ => BlockMerkleNode::ForgettenSubtree { value: *value },
                })
            })
            .collect();
        proof.push(BlockMerkleNode::Branch {
            value: Sha3Node::default(),
            children,
        });
    }
    Ok(BlockMerkleProof {
        pos: target.pos,
        proof,
    })
}

const ARITY: u64 = BlockMerkleTree::ARITY as u64;

/// Child values of `node`, whose children have indices `parent * ARITY + n` at their level, in the
/// tree where the rightmost non-empty node of that level is `edge` with value `edge_value`.
fn child_values(
    node: &BlockMerkleNode,
    parent: u64,
    edge: u64,
    edge_value: Sha3Node,
) -> Result<Vec<Sha3Node>> {
    let BlockMerkleNode::Branch { children, .. } = node else {
        bail!("expected a branch node in the proof");
    };
    Ok(children
        .iter()
        .enumerate()
        .map(|(n, child)| match (parent * ARITY + n as u64).cmp(&edge) {
            Ordering::Less => child.value(),
            Ordering::Equal => edge_value,
            Ordering::Greater => Sha3Node::default(),
        })
        .collect())
}

pub(crate) async fn fetch_block_data_range<State>(
    state: &State,
    start: usize,
    end: usize,
    fetch_timeout: Duration,
    large_object_range_limit: usize,
) -> Result<
    Vec<(
        Header,
        PayloadQueryData<SeqTypes>,
        VidCommonQueryData<SeqTypes>,
    )>,
    Error,
>
where
    State: AvailabilityDataSource<SeqTypes>,
{
    if end <= start {
        return Err(Error::Custom {
            message: format!("requested empty interval [{start}, {end})"),
            status: StatusCode::BAD_REQUEST,
        });
    }
    if end - start > large_object_range_limit {
        return Err(Error::Custom {
            message: format!(
                "requested range [{start}, {end}) exceeds maximum size {large_object_range_limit}"
            ),
            status: StatusCode::BAD_REQUEST,
        });
    }

    let fetch_headers = async move {
        state
            .get_header_range(start..end)
            .await
            .enumerate()
            .then(|(i, fetch)| async move {
                fetch
                    .with_timeout(fetch_timeout)
                    .await
                    .ok_or_else(|| Error::Custom {
                        message: format!("missing header {}", start + i),
                        status: StatusCode::NOT_FOUND,
                    })
            })
            .try_collect::<Vec<_>>()
            .await
    };
    let fetch_payloads = async move {
        state
            .get_payload_range(start..end)
            .await
            .enumerate()
            .then(|(i, fetch)| async move {
                fetch
                    .with_timeout(fetch_timeout)
                    .await
                    .ok_or_else(|| Error::Custom {
                        message: format!("missing payload {}", start + i),
                        status: StatusCode::NOT_FOUND,
                    })
            })
            .try_collect::<Vec<_>>()
            .await
    };
    let fetch_vid_commons = async move {
        state
            .get_vid_common_range(start..end)
            .await
            .enumerate()
            .then(|(i, fetch)| async move {
                fetch
                    .with_timeout(fetch_timeout)
                    .await
                    .ok_or_else(|| Error::Custom {
                        message: format!("missing VID common {}", start + i),
                        status: StatusCode::NOT_FOUND,
                    })
            })
            .try_collect::<Vec<_>>()
            .await
    };
    let (headers, (payloads, vid_commons)) =
        try_join(fetch_headers, try_join(fetch_payloads, fetch_vid_commons)).await?;

    Ok(izip!(headers, payloads, vid_commons).collect())
}

/// Single-namespace version of [`get_namespaces_proof_range`].
pub(crate) async fn get_namespace_proof_range<State>(
    state: &State,
    start: usize,
    end: usize,
    namespace: u64,
    fetch_timeout: Duration,
    large_object_range_limit: usize,
) -> Result<Vec<NamespaceProof>, Error>
where
    State: AvailabilityDataSource<SeqTypes>,
{
    let blocks = get_namespaces_proof_range(
        state,
        start,
        end,
        &[namespace],
        fetch_timeout,
        large_object_range_limit,
    )
    .await?;
    Ok(blocks
        .into_iter()
        .map(|mut block| {
            block
                .remove(&namespace)
                .unwrap_or_else(NamespaceProof::not_present)
        })
        .collect())
}

pub(crate) async fn get_namespaces_proof_range<State>(
    state: &State,
    start: usize,
    end: usize,
    namespaces: &[u64],
    fetch_timeout: Duration,
    large_object_range_limit: usize,
) -> Result<Vec<HashMap<u64, NamespaceProof>>, Error>
where
    State: AvailabilityDataSource<SeqTypes>,
{
    fetch_block_data_range(state, start, end, fetch_timeout, large_object_range_limit)
        .await?
        .into_iter()
        .map(|(header, payload, vid_common)| {
            namespaces
                .iter()
                .filter_map(|&namespace| {
                    let ns_index = header.ns_table().find_ns_id(&namespace.into())?;
                    Some(
                        build_namespace_proof(&payload, &ns_index, &vid_common)
                            .map(|proof| (namespace, proof)),
                    )
                })
                .collect::<Result<HashMap<_, _>, _>>()
        })
        .collect()
}

/// Decode the `namespaces` path segment: a `TaggedBase64` string tagged `NS` wrapping a
/// JSON-encoded `Vec<u64>`.
pub(crate) fn parse_namespaces_str(encoded: &str) -> anyhow::Result<Vec<u64>> {
    let encoded: TaggedBase64 = encoded
        .parse()
        .map_err(|err| anyhow::anyhow!("invalid namespaces parameter: {err}"))?;
    if encoded.tag() != NAMESPACES_PARAM_TAG {
        anyhow::bail!(
            "invalid namespaces parameter tag: expected {NAMESPACES_PARAM_TAG}, got {}",
            encoded.tag()
        );
    }
    serde_json::from_slice(&encoded.value())
        .map_err(|err| anyhow::anyhow!("invalid namespaces parameter: {err}"))
}

/// Construct a [`NamespaceProof`] for the namespace at `ns_index` of the given block.
fn build_namespace_proof(
    payload: &PayloadQueryData<SeqTypes>,
    ns_index: &NsIndex,
    vid_common: &VidCommonQueryData<SeqTypes>,
) -> Result<NamespaceProof, Error> {
    let ns_proof =
        NsProof::new(payload.data(), ns_index, vid_common.common()).ok_or_else(|| {
            Error::Custom {
                message: "failed to construct namespace proof".into(),
                status: StatusCode::INTERNAL_SERVER_ERROR,
            }
        })?;
    Ok(NamespaceProof::new(ns_proof, vid_common.common().clone()))
}

fn internal(err: impl Display) -> Error {
    Error::Custom {
        message: err.to_string(),
        status: StatusCode::INTERNAL_SERVER_ERROR,
    }
}

fn not_found(msg: impl Into<String>) -> Error {
    Error::Custom {
        message: msg.into(),
        status: StatusCode::NOT_FOUND,
    }
}

fn chain_too_long(requested: usize, chain_limit: usize) -> Error {
    not_found(format!(
        "no finality proof found within {chain_limit} leaves of requested leaf {requested}"
    ))
}

#[cfg(test)]
mod test {
    use std::{marker::PhantomData, sync::Mutex};

    use axum::{Router, extract::Path, http::HeaderMap, routing::get};
    use committable::Committable;
    use disco_types::error::Error;
    use espresso_types::BLOCK_MERKLE_TREE_HEIGHT;
    use futures::future::join_all;
    use hotshot_query_service::{
        availability::{BlockQueryData, TransactionIndex, VidCommonQueryData},
        data_source::{
            Transaction,
            storage::{AvailabilityStorage, UpdateAvailabilityStorage, pruning::PrunerCfg},
        },
        fetching::provider::TrustedQueryServiceProvider,
        merklized_state::{MerklizedState, UpdateStateData},
    };
    use hotshot_types::{data::ViewNumber, simple_vote::Vote2Data};
    use http_wire::{ServerError, respond};
    use itertools::Itertools;
    use jf_merkle_tree_compat::{
        AppendableMerkleTreeScheme, ToTraversalPath, prelude::SHA3MerkleTree,
    };
    use light_client::{
        client::{Client as _, FallbackClient, QueryServiceClient},
        consensus::leaf::{FinalityProof, LeafProofHint},
        testing::{
            AlwaysTrueQuorum, LEAF_CHAIN_EPOCH_HEIGHT, TestClient, VersionCheckQuorum,
            custom_leaf_chain_with_upgrade, custom_leaf_chain_with_upgrades, leaf_chain,
            leaf_chain_with_upgrade,
        },
    };
    use rstest_reuse::apply;
    use tokio::{
        net::TcpListener,
        time::{sleep, timeout},
    };
    use vbs::version::Version;
    use versions::{
        DRB_AND_HEADER_UPGRADE_VERSION, EPOCH_REWARD_VERSION, EPOCH_VERSION, FEE_VERSION,
        NEW_PROTOCOL_VERSION, Upgrade,
    };

    use super::*;
    use crate::api::{
        data_source::{Provider, testing::TestableSqlDataSource},
        test::sql_data_sources,
    };

    const CHAIN_LIMIT: usize = 500;

    fn cert2_for_leaf(leaf: &LeafQueryData<SeqTypes>) -> espresso_types::Certificate2<SeqTypes> {
        let data = Vote2Data {
            leaf_commit: leaf.leaf().commit(),
            epoch: leaf.qc().data.epoch.unwrap(),
            block_number: leaf.height(),
        };
        espresso_types::Certificate2::new(
            data.clone(),
            data.commit(),
            leaf.leaf().view_number(),
            None,
            PhantomData,
        )
    }

    #[apply(sql_data_sources)]
    async fn test_two_chain<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Insert some leaves, forming a chain.
        let leaves = leaf_chain(1..=3, EPOCH_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.insert_leaf(&leaves[1]).await.unwrap();
            tx.insert_leaf(&leaves[2]).await.unwrap();
            tx.commit().await.unwrap();
        }

        // Ask for the first leaf; it is proved finalized by the chain formed along with the second.
        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_finalized<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Insert a single leaf. We will not be able to provide proofs ending in a leaf chain, but
        // we can return a leaf if the leaf after it is already known to be finalized.
        let leaves = leaf_chain(1..=2, EPOCH_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.commit().await.unwrap();
        }

        let proof = get_leaf_proof_with_finalized_assumption(
            &ds,
            leaves[0].clone(),
            2,
            Duration::MAX,
            CHAIN_LIMIT,
        )
        .await
        .unwrap();
        assert_eq!(
            proof
                .verify(LeafProofHint::assumption(leaves[1].leaf()))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_new_protocol_finalized_assumption<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        let leaves = leaf_chain(1..=2, NEW_PROTOCOL_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.commit().await.unwrap();
        }

        let proof = get_leaf_proof_with_finalized_assumption(
            &ds,
            leaves[0].clone(),
            2,
            Duration::MAX,
            CHAIN_LIMIT,
        )
        .await
        .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::Assumption));
        assert_eq!(
            proof
                .verify(LeafProofHint::assumption(leaves[1].leaf()))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_new_protocol_cert2<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        let leaves = leaf_chain(1..=2, NEW_PROTOCOL_VERSION).await;
        let cert2_leaf = &leaves[1];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.insert_leaf(cert2_leaf).await.unwrap();
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        let proof = get_leaf_proof_with_cert2(&ds, leaves[0].clone(), Duration::MAX, CHAIN_LIMIT)
            .await
            .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_new_protocol_cert2_chain_limit<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // A proof for the first leaf must extend to the nearest cert2, three leaves away.
        let leaves = leaf_chain(1..=4, NEW_PROTOCOL_VERSION).await;
        let cert2_leaf = &leaves[3];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_cert2(&ds, leaves[0].clone(), Duration::MAX, 2)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);

        let proof = get_leaf_proof_with_cert2(&ds, leaves[0].clone(), Duration::MAX, 3)
            .await
            .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_new_protocol_hint_vs_cert2<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        let leaves = leaf_chain(1..=5, NEW_PROTOCOL_VERSION).await;
        let cert2_leaf = &leaves[1];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        // The cert2 at height 2 beats the hint at height 5, so the hint is ignored.
        let proof = get_leaf_proof(&ds, leaves[0].clone(), Some(5), Duration::MAX, CHAIN_LIMIT)
            .await
            .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );

        // A hint at height 3 ties the cert2 proof length, so the hint wins.
        let proof = get_leaf_proof(&ds, leaves[0].clone(), Some(3), Duration::MAX, CHAIN_LIMIT)
            .await
            .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::Assumption));
        assert_eq!(
            proof
                .verify(LeafProofHint::assumption(leaves[2].leaf()))
                .await
                .unwrap(),
            leaves[0]
        );

        // With no cert2 at or after the requested height, the hint is honored.
        let proof = get_leaf_proof(&ds, leaves[2].clone(), Some(5), Duration::MAX, CHAIN_LIMIT)
            .await
            .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::Assumption));
        assert_eq!(
            proof
                .verify(LeafProofHint::assumption(leaves[4].leaf()))
                .await
                .unwrap(),
            leaves[2]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_distant_hint_falls_through<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Legacy leaves 1-3 upgrade to the new protocol at height 4, with a cert2 at height 5, so
        // both fall-through paths have a bounded proof available.
        let leaves = leaf_chain_with_upgrade(
            1..=5,
            4,
            Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, NEW_PROTOCOL_VERSION),
        )
        .await;
        let cert2_leaf = &leaves[4];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        // A hint more than `chain_limit` past a legacy leaf is ignored in favor of a QC chain.
        let proof = get_leaf_proof(
            &ds,
            leaves[0].clone(),
            Some(1000),
            Duration::MAX,
            CHAIN_LIMIT,
        )
        .await
        .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::HotStuff2 { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );

        // A hint more than `chain_limit` past a new-protocol leaf is ignored in favor of cert2.
        let proof = get_leaf_proof(
            &ds,
            leaves[3].clone(),
            Some(1000),
            Duration::MAX,
            CHAIN_LIMIT,
        )
        .await
        .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[3]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_qc_chain_chain_limit<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Proving the first leaf finalized requires walking the two subsequent leaves.
        let leaves = leaf_chain(1..=3, EPOCH_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, 1)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);

        let proof = get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, 2)
            .await
            .unwrap();
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_finalized_hint_too_far<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        let leaves = leaf_chain(1..=2, EPOCH_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_finalized_assumption(
            &ds,
            leaves[0].clone(),
            1000,
            Duration::MAX,
            10,
        )
        .await
        .unwrap_err();
        assert_eq!(err.status(), StatusCode::BAD_REQUEST);
    }

    #[apply(sql_data_sources)]
    async fn test_qc_chain_new_protocol_cutover<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Upgrade mid-chain and skew view numbers so no HotStuff 2-chain ever forms; the proof
        // for the first (pre-upgrade) leaf must fall through to cert2 finality.
        let leaves = custom_leaf_chain_with_upgrade(
            1..=4,
            2,
            Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, NEW_PROTOCOL_VERSION),
            |proposal| {
                proposal.view_number = ViewNumber::new(proposal.block_header.height() * 2);
            },
        )
        .await;
        assert_eq!(leaves[0].header().version(), DRB_AND_HEADER_UPGRADE_VERSION);
        assert_eq!(leaves[1].header().version(), NEW_PROTOCOL_VERSION);
        let cert2_leaf = &leaves[3];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_qc_chain_new_protocol_cutover_two_chain<D: TestableSqlDataSource>(
        _d: PhantomData<D>,
    ) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Upgrade at height 3: the HotStuff2 chain begun before the cutover completes without
        // needing a cert2 (none is stored).
        let leaves = leaf_chain_with_upgrade(
            1..=3,
            3,
            Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, NEW_PROTOCOL_VERSION),
        )
        .await;
        assert_eq!(leaves[1].header().version(), DRB_AND_HEADER_UPGRADE_VERSION);
        assert_eq!(leaves[2].header().version(), NEW_PROTOCOL_VERSION);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.commit().await.unwrap();
        }

        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::HotStuff2 { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&VersionCheckQuorum::new(
                    leaves.iter().map(|leaf| leaf.leaf().clone()),
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_qc_chain_last_legacy_leaf_uses_cert2<D: TestableSqlDataSource>(
        _d: PhantomData<D>,
    ) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Upgrade at height 2: any 2-chain proving the last legacy leaf spans the cutover, so the
        // prover must fall through to the cert2 at height 3 instead.
        let leaves = leaf_chain_with_upgrade(
            1..=4,
            2,
            Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, NEW_PROTOCOL_VERSION),
        )
        .await;
        assert_eq!(leaves[0].header().version(), DRB_AND_HEADER_UPGRADE_VERSION);
        assert_eq!(leaves[1].header().version(), NEW_PROTOCOL_VERSION);

        let cert2_leaf = &leaves[2];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&VersionCheckQuorum::new(
                    leaves.iter().map(|leaf| leaf.leaf().clone()),
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    /// With a gapless cutover, the leaf two before the cutover (`upgrade_height
    /// - 2`) is finalized by the old protocol: its HotStuff2 2-chain completes
    /// within pre-cutover leaves (deciding QC belongs to `upgrade_height - 1`,
    /// still V5), so no cert2 is needed.
    #[apply(sql_data_sources)]
    async fn test_leaf_before_cutover_uses_hotstuff2<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Cutover at height 4, so `upgrade_height - 2` is height 2 (leaves[1]).
        let leaves = leaf_chain_with_upgrade(
            1..=5,
            4,
            Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, NEW_PROTOCOL_VERSION),
        )
        .await;
        assert!(leaves[1].header().version() < NEW_PROTOCOL_VERSION);
        assert!(leaves[3].header().version() >= NEW_PROTOCOL_VERSION);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.commit().await.unwrap();
        }

        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[1].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::HotStuff2 { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&VersionCheckQuorum::new(
                    leaves.iter().map(|leaf| leaf.leaf().clone()),
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[1]
        );
    }

    /// When a view gap at the cutover prevents a HotStuff2 2-chain from
    /// completing within pre-cutover leaves, the leaf two before the cutover
    /// (`upgrade_height - 2`) instead falls through to cert2 finality. Skewing
    /// the view numbers breaks the consecutive-view requirement, so no 2-chain
    /// forms and the proof terminates at the first cert2.
    #[apply(sql_data_sources)]
    async fn test_leaf_before_cutover_uses_cert2<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Cutover at height 4, so `upgrade_height - 2` is height 2 (leaves[1]).
        // Skewed views mean no HotStuff2 2-chain ever forms.
        let leaves = custom_leaf_chain_with_upgrade(
            1..=5,
            4,
            Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, NEW_PROTOCOL_VERSION),
            |proposal| {
                proposal.view_number = ViewNumber::new(proposal.block_header.height() * 2);
            },
        )
        .await;
        assert!(leaves[1].header().version() < NEW_PROTOCOL_VERSION);
        assert!(leaves[3].header().version() >= NEW_PROTOCOL_VERSION);
        let cert2_leaf = &leaves[3];
        let cert2 = cert2_for_leaf(cert2_leaf);

        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.insert_cert2(cert2_leaf.height(), cert2).await.unwrap();
            tx.commit().await.unwrap();
        }

        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[1].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. }));
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[1]
        );
    }

    /// Every leaf of a history whose tip is past the 0.6 cutover is proved from stored leaves and
    /// cert2s alone, with no legacy QC chain stored.
    #[apply(sql_data_sources)]
    async fn test_leaf_proofs_across_history<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        // Every protocol version from mainnet's 0.2 genesis on, with the height it took effect.
        const UPGRADES: [(u64, Version); 4] = [
            (6, EPOCH_VERSION),
            (11, DRB_AND_HEADER_UPGRADE_VERSION),
            (16, EPOCH_REWARD_VERSION),
            (21, NEW_PROTOCOL_VERSION),
        ];
        // A view fails just before each of these heights, under HotStuff and under HotStuff2.
        const FAILED_VIEWS_BEFORE: [u64; 2] = [3, 13];
        const CERT2_HEIGHTS: [u64; 3] = [22, 24, 25];

        let view_of = |height: u64| {
            height + FAILED_VIEWS_BEFORE.iter().filter(|&&h| h <= height).count() as u64
        };
        let leaves = custom_leaf_chain_with_upgrades(1..=25, FEE_VERSION, &UPGRADES, |proposal| {
            proposal.view_number = ViewNumber::new(view_of(proposal.block_header.height()));
            // Upgrades take effect by view, so failed views move where each one starts.
            if let Some(cert) = &mut proposal.upgrade_certificate {
                let (height, _) = UPGRADES
                    .into_iter()
                    .find(|(_, version)| *version == cert.data.new_version)
                    .unwrap();
                let first_view = ViewNumber::new(view_of(height));
                cert.data.old_version_last_view = first_view - 1;
                cert.data.new_version_first_view = first_view;
                cert.data.decide_by = first_view;
            }
        })
        .await;
        let leaf_at = |height: u64| &leaves[height as usize - 1];

        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();
        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            for height in CERT2_HEIGHTS {
                tx.insert_cert2(height, cert2_for_leaf(leaf_at(height)))
                    .await
                    .unwrap();
            }
            tx.commit().await.unwrap();
        }

        let quorum = VersionCheckQuorum::new(
            leaves.iter().map(|leaf| leaf.leaf().clone()),
            LEAF_CHAIN_EPOCH_HEIGHT,
        );
        for leaf in &leaves {
            let height = leaf.height();
            let proof = get_leaf_proof(&ds, leaf.clone(), None, Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap_or_else(|err| panic!("no proof for height {height}: {err:#}"));
            match height {
                // A 3-chain, which past a failed view closes later on.
                1..=4 => assert!(matches!(proof.proof(), FinalityProof::HotStuff { .. })),
                // From 5, the first 2-chain closes before the 3-chain does.
                5..=19 => assert!(matches!(proof.proof(), FinalityProof::HotStuff2 { .. })),
                // The last legacy leaf's 2-chain would span the cutover.
                _ => assert!(matches!(proof.proof(), FinalityProof::NewProtocol { .. })),
            }
            assert_eq!(
                proof.verify(LeafProofHint::Quorum(&quorum)).await.unwrap(),
                *leaf,
                "height {height}"
            );
        }

        // The client may instead name a leaf it already trusts, here the next one.
        for (leaf, finalized) in leaves.iter().tuple_windows() {
            let height = leaf.height();
            let proof = get_leaf_proof(
                &ds,
                leaf.clone(),
                Some(finalized.height() as usize),
                Duration::MAX,
                CHAIN_LIMIT,
            )
            .await
            .unwrap_or_else(|err| panic!("no proof for height {height}: {err:#}"));
            assert!(
                matches!(proof.proof(), FinalityProof::Assumption),
                "height {height}"
            );
            assert_eq!(
                proof
                    .verify(LeafProofHint::<AlwaysTrueQuorum>::Assumption(
                        finalized.leaf()
                    ))
                    .await
                    .unwrap(),
                *leaf,
                "height {height}"
            );
        }
    }

    #[apply(sql_data_sources)]
    async fn test_bad_finalized<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Insert a single leaf. If we request this leaf but provide a finalized leaf which is
        // earlier, we should fail.
        let leaves = leaf_chain(1..2, EPOCH_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_finalized_assumption(
            &ds,
            leaves[0].clone(),
            0,
            Duration::MAX,
            CHAIN_LIMIT,
        )
        .await
        .unwrap_err();
        assert_eq!(err.status(), StatusCode::BAD_REQUEST);
    }

    #[apply(sql_data_sources)]
    async fn test_no_chain<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Insert multiple leaves that don't chain. We will not be able to prove these are
        // finalized.
        let leaves = leaf_chain(1..=4, EPOCH_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.insert_leaf(&leaves[2]).await.unwrap();
            tx.insert_leaf(&leaves[3]).await.unwrap();
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_leaf_chain(
            &ds,
            leaves[0].clone(),
            Duration::from_secs(1),
            CHAIN_LIMIT,
        )
        .await
        .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);

        // Even if we start from a finalized leave that extends one of the leaves we do have (4,
        // extends 3) we fail to generate a proof because we can't generate a chain from the
        // requested leaf (1) to the finalized leaf (4), since leaf 2 is missing.
        let err = get_leaf_proof_with_finalized_assumption(
            &ds,
            leaves[0].clone(),
            4,
            Duration::from_secs(1),
            CHAIN_LIMIT,
        )
        .await
        .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);
    }

    #[apply(sql_data_sources)]
    async fn test_chain_ends_before_finality<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // A legacy tip, as every replayed pre-0.6 chain has: the next leaf chains onto the
        // requested one, but no deciding QC has arrived and no cert2 ever will.
        let leaves = leaf_chain(1..=3, EPOCH_REWARD_VERSION).await;
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[0]).await.unwrap();
            tx.insert_leaf(&leaves[1]).await.unwrap();
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_leaf_chain(
            &ds,
            leaves[0].clone(),
            Duration::from_secs(1),
            CHAIN_LIMIT,
        )
        .await
        .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);
        assert!(
            matches!(&err, hotshot_query_service::Error::Custom { message, .. } if message.contains("not enough leaves")),
            "{err}"
        );

        // Once the leaf carrying the deciding QC is stored, the same request succeeds.
        {
            let mut tx = ds.write().await.unwrap();
            tx.insert_leaf(&leaves[2]).await.unwrap();
            tx.commit().await.unwrap();
        }
        let proof =
            get_leaf_proof_with_leaf_chain(&ds, leaves[0].clone(), Duration::MAX, CHAIN_LIMIT)
                .await
                .unwrap();
        assert_eq!(
            proof
                .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                    LEAF_CHAIN_EPOCH_HEIGHT
                )))
                .await
                .unwrap(),
            leaves[0]
        );
    }

    #[apply(sql_data_sources)]
    async fn test_header_proof<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Construct a chain of leaves, plus the corresponding block Merkle tree at each leaf.
        let leaves = leaf_chain(0..=2, EPOCH_VERSION).await;
        let mts = leaves
            .iter()
            .scan(
                BlockMerkleTree::new(BLOCK_MERKLE_TREE_HEIGHT),
                |mt, leaf| {
                    assert_eq!(mt.commitment(), leaf.header().block_merkle_tree_root());
                    let item = mt.clone();
                    mt.push(leaf.block_hash()).unwrap();
                    Some(item)
                },
            )
            .collect::<Vec<_>>();

        // Save all those objects in the DB.
        {
            let mut tx = ds.write().await.unwrap();
            for (leaf, mt) in leaves.iter().zip(&mts) {
                tx.insert_leaf(leaf).await.unwrap();

                if leaf.height() > 0 {
                    let merkle_path = mt.lookup(leaf.height() - 1).expect_ok().unwrap().1;
                    UpdateStateData::<SeqTypes, BlockMerkleTree, _>::insert_merkle_nodes(
                        &mut tx,
                        merkle_path,
                        ToTraversalPath::<{ BlockMerkleTree::ARITY }>::to_traversal_path(
                            &(leaf.height() - 1),
                            BLOCK_MERKLE_TREE_HEIGHT,
                        ),
                        leaf.height(),
                    )
                    .await
                    .unwrap();
                    UpdateStateData::<SeqTypes, BlockMerkleTree, _>::set_last_state_height(
                        &mut tx,
                        leaf.height() as usize,
                    )
                    .await
                    .unwrap();
                }
            }
            tx.commit().await.unwrap();
        }

        // Test happy path.
        for (root, mt) in mts.iter().enumerate().skip(1) {
            for (height, leaf) in leaves.iter().enumerate().take(root) {
                tracing::info!(root, height, "test happy path");
                let proof =
                    get_header_proof(&ds, root as u64, BlockId::Number(height), Duration::MAX)
                        .await
                        .unwrap();
                assert_eq!(proof.verify_ref(mt.commitment()).unwrap(), leaf.header());
            }
        }

        // Test unknown leaf.
        let err = get_header_proof(&ds, 5, BlockId::Number(4), Duration::from_secs(1))
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);

        // Test height >= root.
        let err = get_header_proof(&ds, 1, BlockId::Number(1), Duration::MAX)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::BAD_REQUEST);
    }

    // The light-weight block tree forgets all but its last leaf.
    type FullTree = SHA3MerkleTree<Commitment<Header>>;

    fn fake_tree(num_leaves: u64) -> FullTree {
        FullTree::from_elems(
            Some(BLOCK_MERKLE_TREE_HEIGHT),
            (0..num_leaves).map(|i| Commitment::<Header>::from_raw([i as u8 + 1; 32])),
        )
        .unwrap()
    }

    fn lookup(tree: &FullTree, pos: u64) -> BlockMerkleProof {
        tree.lookup(pos).expect_ok().unwrap().1
    }

    /// Rebuilding from `head_leaves` leaves matches the tree of `r` leaves for every `r` in `rs`.
    fn check_rebuild(head_leaves: u64, rs: impl IntoIterator<Item = u64>) {
        let head = fake_tree(head_leaves);
        for r in rs {
            let tree = fake_tree(r);
            let frontier = lookup(&head, r - 1);
            for i in 0..r {
                let proof = rebuild_block_proof(r, &frontier, &lookup(&head, i)).unwrap();
                assert_eq!(proof, lookup(&tree, i), "r={r} i={i}");
                assert!(
                    BlockMerkleTree::verify(tree.commitment(), i, &proof)
                        .unwrap()
                        .is_ok()
                );
            }
        }
    }

    #[test]
    fn test_rebuild_block_proof_matches_snapshot() {
        check_rebuild(100, [2, 4, 5, 10, 26, 28, 50, 99]);
    }

    #[test]
    fn test_rebuild_block_proof_last_leaf() {
        let head = fake_tree(40);
        let tree = fake_tree(13);
        let proof = rebuild_block_proof(13, &lookup(&head, 12), &lookup(&head, 12)).unwrap();
        assert_eq!(proof, lookup(&tree, 12));
    }

    #[test]
    fn test_rebuild_block_proof_full_subtree() {
        check_rebuild(100, [3, 9, 27, 81]);
    }

    #[test]
    fn test_rebuild_block_proof_single_leaf() {
        check_rebuild(5, [1]);
    }

    #[test]
    fn test_rebuild_block_proof_wrong_frontier() {
        let head = fake_tree(40);
        let err = rebuild_block_proof(13, &lookup(&head, 11), &lookup(&head, 3)).unwrap_err();
        assert!(err.to_string().contains("frontier"), "{err:#}");
    }

    /// A data source with `num_leaves` headers, state at every height, and state below `pruned`
    /// collected. Returns the data source and the block tree at each height. The storage guard
    /// keeps the database alive.
    async fn pruned_data_source<D: TestableSqlDataSource>(
        num_leaves: usize,
        pruned: u64,
    ) -> (impl Sized, D, Vec<BlockMerkleTree>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        let leaves = leaf_chain(0..num_leaves as u64, EPOCH_VERSION).await;
        let mts = leaves
            .iter()
            .scan(
                BlockMerkleTree::new(BLOCK_MERKLE_TREE_HEIGHT),
                |mt, leaf| {
                    let item = mt.clone();
                    mt.push(leaf.block_hash()).unwrap();
                    Some(item)
                },
            )
            .collect::<Vec<_>>();

        let mut tx = ds.write().await.unwrap();
        for (leaf, mt) in leaves.iter().zip(&mts) {
            tx.insert_leaf(leaf).await.unwrap();
            if leaf.height() > 0 {
                let merkle_path = mt.lookup(leaf.height() - 1).expect_ok().unwrap().1;
                UpdateStateData::<SeqTypes, BlockMerkleTree, _>::insert_merkle_nodes(
                    &mut tx,
                    merkle_path,
                    ToTraversalPath::<{ BlockMerkleTree::ARITY }>::to_traversal_path(
                        &(leaf.height() - 1),
                        BLOCK_MERKLE_TREE_HEIGHT,
                    ),
                    leaf.height(),
                )
                .await
                .unwrap();
                UpdateStateData::<SeqTypes, BlockMerkleTree, _>::set_last_state_height(
                    &mut tx,
                    leaf.height() as usize,
                )
                .await
                .unwrap();
            }
        }
        tx.commit().await.unwrap();

        let cfg =
            PrunerCfg::default().with_state_tables(vec![BlockMerkleTree::state_type().into()]);
        D::prune_state_below(&ds, pruned, &cfg).await.unwrap();
        (storage, ds, mts)
    }

    #[apply(sql_data_sources)]
    async fn test_header_proof_pruned_root<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let num_leaves = 40;
        let (_storage, ds, mts) = pruned_data_source::<D>(num_leaves, 30).await;

        // Roots below the pruned height are rebuilt from the newest snapshot.
        for root in [1, 2, 4, 10, 13, 27, 28] {
            for height in 0..root {
                let proof = get_header_proof(
                    &ds,
                    root as u64,
                    BlockId::Number(height),
                    Duration::from_secs(10),
                )
                .await
                .unwrap();
                proof.verify(mts[root].commitment()).unwrap();
            }
        }
    }

    #[apply(sql_data_sources)]
    async fn test_header_proof_root_beyond_head_fails<D: TestableSqlDataSource>(
        _d: PhantomData<D>,
    ) {
        let (_storage, ds, _) = pruned_data_source::<D>(10, 5).await;
        // The state at the head height is stored; beyond it there is no snapshot to rebuild from.
        let err = get_header_proof(&ds, 20, BlockId::Number(3), Duration::from_secs(10))
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);
    }

    #[apply(sql_data_sources)]
    async fn test_namespace_proof<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let storage = D::create_storage().await;
        let ds = D::create(D::persistence_options(&storage), Default::default(), false)
            .await
            .unwrap();

        // Construct a chain of blocks.
        let client = TestClient::default();
        let leaves = join_all((0..=2).map(|i| client.leaf(i))).await;
        let payloads = join_all((0..=2).map(|i| client.payload(i))).await;
        let vid_commons = join_all((0..=2).map(|i| client.vid_common(i))).await;

        // Save all those objects in the DB.
        {
            let mut tx = ds.write().await.unwrap();
            for (leaf, payload, vid_common) in izip!(&leaves, &payloads, &vid_commons) {
                tx.insert_leaf(leaf).await.unwrap();
                tx.insert_block(&BlockQueryData::<SeqTypes>::new(
                    leaf.header().clone(),
                    payload.clone(),
                ))
                .await
                .unwrap();
                tx.insert_vid(
                    &VidCommonQueryData::<SeqTypes>::new(leaf.header().clone(), vid_common.clone()),
                    None,
                )
                .await
                .unwrap();
            }
            tx.commit().await.unwrap();
        }

        // Test happy path: all blocks.
        let ns = payloads[0]
            .transaction(&TransactionIndex {
                ns_index: 0.into(),
                position: 0,
            })
            .unwrap()
            .namespace();
        let proofs = get_namespace_proof_range(&ds, 0, 3, ns.into(), Duration::MAX, 100)
            .await
            .unwrap();
        assert_eq!(proofs.len(), 3);
        for (leaf, proof) in leaves.iter().zip(proofs) {
            proof.verify(leaf.header(), ns).unwrap();
        }

        // Test happy path: subset.
        let tx = payloads[1]
            .transaction(&TransactionIndex {
                ns_index: 0.into(),
                position: 0,
            })
            .unwrap();
        let ns = tx.namespace();
        let proofs = get_namespace_proof_range(&ds, 1, 2, ns.into(), Duration::MAX, 100)
            .await
            .unwrap();
        assert_eq!(proofs.len(), 1);
        assert_eq!(proofs[0].verify(leaves[1].header(), ns).unwrap(), [tx]);

        // Test missing data in range.
        let err = get_namespace_proof_range(&ds, 0, 4, ns.into(), Duration::from_secs(1), 100)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);

        // Test invalid range.
        let err = get_namespace_proof_range(&ds, 1, 0, ns.into(), Duration::from_secs(1), 100)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::BAD_REQUEST);
        assert!(
            err.to_string().contains("requested empty interval"),
            "{err:#}"
        );

        // Test large range.
        let err = get_namespace_proof_range(&ds, 0, 10_000, ns.into(), Duration::from_secs(1), 100)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::BAD_REQUEST);
        assert!(err.to_string().contains("exceeds maximum size"), "{err:#}");
    }

    /// A peer with one cert2, 404ing every other height as a node does for a leaf finalized
    /// indirectly. Records the heights it is asked for.
    async fn serve_cert2(
        cert2: espresso_types::Certificate2<SeqTypes>,
    ) -> (TrustedQueryServiceProvider, Arc<Mutex<Vec<u64>>>) {
        let requested = Arc::new(Mutex::new(Vec::new()));
        let seen = requested.clone();
        let router = Router::new().route(
            "/availability/cert2/{height}",
            get(move |headers: HeaderMap, Path(height): Path<u64>| {
                seen.lock().unwrap().push(height);
                let served = (height == cert2.data.block_number).then(|| cert2.clone());
                async move {
                    respond(
                        &headers,
                        served.ok_or_else(|| ServerError {
                            status: axum::http::StatusCode::NOT_FOUND,
                            message: format!("no cert2 at height {height}"),
                        }),
                    )
                }
            }),
        );
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
        let provider = TrustedQueryServiceProvider::new(format!("http://{addr}").parse().unwrap());
        (provider, requested)
    }

    /// A peer serving `cert2` over the catchup route at every height, or 404ing when it has none.
    /// Records the heights it is asked for.
    async fn serve_catchup_cert2(
        cert2: Option<espresso_types::Certificate2<SeqTypes>>,
    ) -> (QueryServiceClient, Arc<Mutex<Vec<u64>>>) {
        let requested = Arc::new(Mutex::new(Vec::new()));
        let seen = requested.clone();
        let router = Router::new().route(
            "/catchup/{height}/cert2",
            get(move |headers: HeaderMap, Path(height): Path<u64>| {
                seen.lock().unwrap().push(height);
                let served = cert2.clone();
                async move {
                    respond(
                        &headers,
                        served.ok_or_else(|| ServerError {
                            status: axum::http::StatusCode::NOT_FOUND,
                            message: format!("no cert2 available for height {height}"),
                        }),
                    )
                }
            }),
        );
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
        let client = QueryServiceClient::new(format!("http://{addr}").parse().unwrap());
        (client, requested)
    }

    /// A proof no stored cert2 can complete asks the peers for one. The peer 404s the requested
    /// height itself, so recovery has to reach a descendant.
    #[apply(sql_data_sources)]
    async fn test_cert2_recovery_from_peer<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let leaves = leaf_chain(1..=4, NEW_PROTOCOL_VERSION).await;
        let (peer, requested) = serve_cert2(cert2_for_leaf(&leaves[1])).await;

        // Once with no cert2 stored, once with one stored too far past the requested leaf.
        for stored_out_of_limit in [false, true] {
            requested.lock().unwrap().clear();
            let storage = D::create_storage().await;
            let ds = D::on_request_data_source(
                &storage,
                Provider::default().with_cert2_provider(peer.clone()),
            )
            .await
            .unwrap();
            {
                let mut tx = ds.write().await.unwrap();
                for leaf in &leaves {
                    tx.insert_leaf(leaf).await.unwrap();
                }
                if stored_out_of_limit {
                    tx.insert_cert2(leaves[3].height(), cert2_for_leaf(&leaves[3]))
                        .await
                        .unwrap();
                }
                tx.commit().await.unwrap();
            }

            let err = get_leaf_proof_with_cert2(&ds, leaves[0].clone(), Duration::from_secs(1), 2)
                .await
                .unwrap_err();
            assert_eq!(err.status(), StatusCode::NOT_FOUND);

            timeout(Duration::from_secs(30), async {
                loop {
                    let mut tx = ds.read().await.unwrap();
                    let stored: Option<espresso_types::Certificate2<SeqTypes>> =
                        tx.load_cert2(leaves[1].height()).await.unwrap();
                    if stored.is_some() {
                        break;
                    }
                    drop(tx);
                    sleep(Duration::from_millis(100)).await;
                }
            })
            .await
            .expect("descendant cert2 was not recovered from the peer");

            let asked = requested.lock().unwrap().clone();
            assert_eq!(
                asked,
                [1, 2],
                "the walk did not stop at the first cert2 it found"
            );

            let proof =
                get_leaf_proof_with_cert2(&ds, leaves[0].clone(), Duration::from_secs(1), 2)
                    .await
                    .unwrap();
            assert_eq!(
                proof
                    .verify(LeafProofHint::Quorum(&AlwaysTrueQuorum::new(
                        LEAF_CHAIN_EPOCH_HEIGHT
                    )))
                    .await
                    .unwrap(),
                leaves[0]
            );
        }
    }

    /// Recovery stops at `chain_limit` past the requested leaf, even with a cert2 just beyond.
    #[apply(sql_data_sources)]
    async fn test_cert2_recovery_stops_at_the_window<D: TestableSqlDataSource>(_d: PhantomData<D>) {
        let leaves = leaf_chain(1..=4, NEW_PROTOCOL_VERSION).await;
        // The only cert2 the peer has is one leaf beyond what a `chain_limit` of 2 may reach.
        let (peer, requested) = serve_cert2(cert2_for_leaf(&leaves[3])).await;

        let storage = D::create_storage().await;
        let ds = D::on_request_data_source(&storage, Provider::default().with_cert2_provider(peer))
            .await
            .unwrap();
        {
            let mut tx = ds.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            tx.commit().await.unwrap();
        }

        let err = get_leaf_proof_with_cert2(&ds, leaves[0].clone(), Duration::from_secs(1), 2)
            .await
            .unwrap_err();
        assert_eq!(err.status(), StatusCode::NOT_FOUND);

        // Nothing answers, so the walk runs the whole window and then stops.
        timeout(Duration::from_secs(30), async {
            while !(1..=3).all(|h| requested.lock().unwrap().contains(&h)) {
                sleep(Duration::from_millis(100)).await;
            }
        })
        .await
        .expect("recovery did not reach every height within the limit");
        sleep(Duration::from_secs(1)).await;
        let asked = requested.lock().unwrap().clone();
        assert!(
            asked.iter().all(|h| (1..=3).contains(h)),
            "recovery reached past the limit: {asked:?}"
        );
    }

    #[tokio::test]
    async fn test_fallback_client_asks_every_peer_for_cert2() {
        let leaves = leaf_chain(1..=2, NEW_PROTOCOL_VERSION).await;
        let cert2 = cert2_for_leaf(&leaves[1]);
        let height = leaves[1].height();

        let mut peers = Vec::new();
        let mut requests = Vec::new();
        for holds_cert2 in [false, false, false, false, true] {
            let (peer, requested) = serve_catchup_cert2(holds_cert2.then(|| cert2.clone())).await;
            peers.push(peer);
            requests.push(requested);
        }
        let client = FallbackClient::new(peers).unwrap();

        assert_eq!(client.cert2(height).await.unwrap(), Some(cert2));
        for (i, requested) in requests.iter().enumerate() {
            assert_eq!(
                *requested.lock().unwrap(),
                [height],
                "peer {i} was not asked exactly once"
            );
        }
    }
}