hotshot_new_protocol/vid/

reconstruct.rs

use std::{
    collections::{BTreeMap, BTreeSet, HashSet, btree_map::Entry},
    ops::Range,
};

use committable::Commitment;
use hotshot_types::{
    data::{
        EpochNumber, VidCommitment2, VidDisperseShare2, ViewNumber,
        ns_table::{ns_table_payload_byte_len, parse_ns_table},
    },
    traits::{BlockPayload, block_contents::EncodeBytes, node_implementation::NodeType},
    vid::avidm_gf2::{AvidmGf2Common, AvidmGf2Param, AvidmGf2Scheme, AvidmGf2Share},
};
use tokio::task::{AbortHandle, JoinSet};
use tracing::{error, warn};

use super::ns_lens_match_metadata;

pub(crate) type Metadata<T> = <<T as NodeType>::BlockPayload as BlockPayload<T>>::Metadata;

type ReconstructResult<T> =
    Result<ObtainedPayload<T>, VidReconstructError<<T as NodeType>::SignatureKey>>;

/// A block's payload in hand, with what the node needs to file it.
///
/// Produced by decoding a view's shares and, on the recovery path, by a
/// payload a peer sent whole. Either way the bytes are bound to
/// `payload_commitment` before this is built.
pub struct ObtainedPayload<T: NodeType> {
    pub view: ViewNumber,
    pub epoch: EpochNumber,
    pub payload_commitment: VidCommitment2,
    pub payload: T::BlockPayload,
    pub metadata: <T::BlockPayload as BlockPayload<T>>::Metadata,
    pub tx_commitments: Vec<Commitment<T::Transaction>>,
}

/// Why a reconstruction attempt failed.
#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
pub enum VidReconstructErrorKind {
    /// Unverifiable shares were weeded out; reconstruction retries once the
    /// remaining shares cover the recovery threshold again.
    #[error("awaiting more shares after weeding out unverifiable ones")]
    AwaitingShares,
    /// Every share verified yet the payload still does not re-commit: the
    /// disperser committed to a non-codeword, so no subset can ever recover it.
    #[error("unrecoverable: verified shares cannot decode to a payload matching the commitment")]
    Unrecoverable,
}

/// A failed reconstruction attempt for one view and claimed commitment.
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
#[error("VID reconstruction failed for view {view}: {kind}")]
pub struct VidReconstructError<K> {
    pub view: ViewNumber,
    pub payload_commitment: VidCommitment2,
    pub kind: VidReconstructErrorKind,
    /// Voters whose shares failed verification against the commitment.
    /// Provably bad and attributable: each share arrived in a message
    /// signed by its voter.
    pub bad_share_keys: Vec<K>,
}

/// What vouches for the proposal a view is pinned to.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PinnedBy {
    /// The leader alone: the proposal validated and paired with this node's
    /// share. An equivocating leader may have given the quorum another one.
    Leader,
    /// A certificate: the quorum voted for it, so the view decides on it if
    /// it decides at all.
    Certificate,
}

#[derive(Default)]
pub struct VidReconstructor<T: NodeType> {
    /// Shares held for a commitment no proposal has pinned their view to,
    /// one per voter: they arrived before the proposal, or while the leader's
    /// pin holds another commitment that a certificate may yet replace.
    pending: BTreeMap<ViewNumber, BTreeMap<T::SignatureKey, VidDisperseShare2<T>>>,
    /// One accumulator per view, created when a proposal pins it:
    /// reconstruction needs the proposal's metadata, so only the pinned
    /// commitment is worth accumulating.
    accumulators: BTreeMap<ViewNumber, VidShareAccumulator<T>>,
    reconstructed: BTreeSet<ViewNumber>,
    tasks: JoinSet<ReconstructResult<T>>,
    calculations: BTreeMap<ViewNumber, AbortHandle>,
    tracer: Option<crate::leader_trace::LeaderTracerHandle>,
}

pub(crate) struct VidShareAccumulator<T: NodeType> {
    /// The payload commitment claimed by the proposal the view is pinned to.
    payload_commitment: VidCommitment2,
    /// What vouches for that proposal; only a certificate's pin is final.
    pinned_by: PinnedBy,
    metadata: Metadata<T>,
    epoch: EpochNumber,
    /// The VID erasure parameters the committee fixes for this view, used to
    /// reject shares carrying a forged `common.param` (see [`Self::accept`]).
    /// `None` if the committee could not be resolved; the param check is then
    /// skipped, matching the previously unchecked path.
    expected_param: Option<AvidmGf2Param>,
    /// Common data pinned by the first admitted share; hash-bound to
    /// `payload_commitment` (see [`Self::accept`]).
    common: Option<AvidmGf2Common>,
    /// Admitted shares by voter; their shard ranges are pairwise disjoint
    /// (see [`Self::accept`]).
    shares: BTreeMap<T::SignatureKey, AvidmGf2Share>,
    /// Every voter whose share was admitted, including weeded ones: a voter
    /// whose share failed verification doesn't get a second submission.
    seen_keys: HashSet<T::SignatureKey>,
    /// Set when a fully verified share set failed to decode to a payload
    /// matching the commitment: the disperser is provably faulty and further
    /// attempts are pointless.
    exhausted: bool,
}

impl<T: NodeType> VidReconstructor<T> {
    pub fn new() -> Self {
        Self {
            pending: BTreeMap::new(),
            accumulators: BTreeMap::new(),
            reconstructed: BTreeSet::new(),
            tasks: JoinSet::new(),
            calculations: BTreeMap::new(),
            tracer: None,
        }
    }

    pub fn set_tracer(&mut self, tracer: Option<crate::leader_trace::LeaderTracerHandle>) {
        self.tracer = tracer;
    }

    /// Pin `view` to a proposal's payload commitment and metadata, and admit
    /// the held shares for that commitment.
    ///
    /// The first pin holds against later ones, except that a certificate's
    /// replaces the leader's: an equivocating leader may have paired this
    /// node with a proposal the quorum did not certify.
    pub(crate) fn handle_proposal(
        &mut self,
        view: ViewNumber,
        payload_commitment: VidCommitment2,
        metadata: Metadata<T>,
        epoch: EpochNumber,
        expected_param: Option<AvidmGf2Param>,
        pinned_by: PinnedBy,
    ) {
        if self.reconstructed.contains(&view) {
            return;
        }
        let accumulator = match self.accumulators.entry(view) {
            Entry::Occupied(mut existing) => {
                let pinned = existing.get_mut();
                if pinned.payload_commitment == payload_commitment {
                    if pinned_by == PinnedBy::Certificate {
                        pinned.pinned_by = PinnedBy::Certificate;
                    }
                    return;
                }
                if pinned_by == PinnedBy::Leader || pinned.pinned_by == PinnedBy::Certificate {
                    warn!(%view, "conflicting proposal for a view pinned to another commitment");
                    return;
                }
                warn!(%view, "leader equivocated; re-pinning to the certified proposal");
                if let Some(task) = self.calculations.remove(&view) {
                    task.abort();
                }
                existing.insert(VidShareAccumulator::new(
                    payload_commitment,
                    metadata,
                    epoch,
                    expected_param,
                    pinned_by,
                ));
                existing.into_mut()
            },
            Entry::Vacant(slot) => slot.insert(VidShareAccumulator::new(
                payload_commitment,
                metadata,
                epoch,
                expected_param,
                pinned_by,
            )),
        };
        // Shares for other commitments stay held: a certificate may yet re-pin
        // the view to one of them.
        let held = self.pending.remove(&view).unwrap_or_default();
        let (admitted, held): (BTreeMap<_, _>, BTreeMap<_, _>) = held
            .into_iter()
            .partition(|(_, share)| share.payload_commitment == payload_commitment);
        if !held.is_empty() {
            self.pending.insert(view, held);
        }
        for (sender, share) in admitted {
            accumulator.accept(view, sender, share);
        }
        self.try_reconstruct(view);
    }

    pub(crate) fn handle_vid_share(
        &mut self,
        sender: T::SignatureKey,
        share: VidDisperseShare2<T>,
    ) {
        let view = share.view_number;
        // A share carries the voter it belongs to; only the authenticated
        // sender may contribute its own. This cheap check bounds each node to
        // one slot and guards the `pending` window, where a share cannot yet
        // be verified against a commitment.
        if share.recipient_key != sender {
            warn!(%view, ?sender, "VID share recipient key does not match its sender");
            return;
        }
        if self.reconstructed.contains(&view) {
            return;
        }
        let Some(accumulator) = self.accumulators.get_mut(&view) else {
            self.hold(view, sender, share);
            return;
        };
        if share.payload_commitment == accumulator.payload_commitment {
            accumulator.accept(view, sender, share);
            self.try_reconstruct(view);
            return;
        }
        match accumulator.pinned_by {
            // A certificate may yet re-pin the view to the share's commitment.
            PinnedBy::Leader => self.hold(view, sender, share),
            PinnedBy::Certificate => {
                warn!(%view, ?sender, "VID share commitment differs from the certified proposal's");
            },
        }
    }

    /// Hold a voter's share until a proposal pins the view to its commitment.
    /// One slot per voter: the first share it sends for the view stays.
    fn hold(&mut self, view: ViewNumber, sender: T::SignatureKey, share: VidDisperseShare2<T>) {
        self.pending
            .entry(view)
            .or_default()
            .entry(sender)
            .or_insert(share);
    }

    pub async fn next(&mut self) -> Option<ReconstructResult<T>> {
        loop {
            match self.tasks.join_next().await? {
                Ok(Ok(out)) => {
                    if self.is_superseded(out.view, out.payload_commitment) {
                        continue;
                    }
                    self.calculations.remove(&out.view);
                    self.accumulators.remove(&out.view);
                    self.reconstructed.insert(out.view);
                    return Some(Ok(out));
                },
                Ok(Err(err)) => {
                    if self.is_superseded(err.view, err.payload_commitment) {
                        continue;
                    }
                    self.calculations.remove(&err.view);
                    self.handle_failed_attempt(&err);
                    return Some(Err(err));
                },
                Err(err) => {
                    if err.is_panic() {
                        error!(%err, "VID reconstruction task panicked");
                    }
                    self.calculations.retain(|_, task| task.id() != err.id());
                },
            }
        }
    }

    /// Whether an attempt's result is for a pin a certificate has since
    /// replaced. The attempt was aborted, but a blocking task that had already
    /// started still runs to completion.
    fn is_superseded(&self, view: ViewNumber, payload_commitment: VidCommitment2) -> bool {
        self.accumulators
            .get(&view)
            .is_some_and(|accumulator| accumulator.payload_commitment != payload_commitment)
    }

    /// Apply the outcome of a failed attempt: weed the bad shares out of the
    /// accumulator, then either mark the payload as unrecoverable or retry
    /// (`try_reconstruct` re-checks coverage, which shares that arrived while
    /// the attempt ran may already restore).
    fn handle_failed_attempt(&mut self, err: &VidReconstructError<T::SignatureKey>) {
        let Some(accumulator) = self.accumulators.get_mut(&err.view) else {
            return;
        };
        for key in &err.bad_share_keys {
            accumulator.shares.remove(key);
        }
        match err.kind {
            VidReconstructErrorKind::Unrecoverable => accumulator.exhausted = true,
            VidReconstructErrorKind::AwaitingShares => self.try_reconstruct(err.view),
        }
    }

    fn try_reconstruct(&mut self, view: ViewNumber) {
        if self.calculations.contains_key(&view) {
            return;
        }
        let Some(accumulator) = self.accumulators.get(&view) else {
            return;
        };
        if accumulator.exhausted || !accumulator.has_enough_shares() {
            return;
        }
        // Enough shares implies an admitted share, which pinned the common.
        let Some(common) = accumulator.common.clone() else {
            return;
        };
        let payload_commitment = accumulator.payload_commitment;
        let metadata = accumulator.metadata.clone();
        let shares: Vec<(T::SignatureKey, AvidmGf2Share)> = accumulator
            .shares
            .iter()
            .map(|(key, share)| (key.clone(), share.clone()))
            .collect();
        let epoch = accumulator.epoch;
        crate::trace_leader_event!(
            self.tracer,
            view,
            crate::leader_trace::LeaderEvent::ThresholdShareReachedVMinus1
        );
        let tracer = self.tracer.clone();
        let task = self.tasks.spawn_blocking(move || {
            crate::trace_leader_event!(
                tracer,
                view,
                crate::leader_trace::LeaderEvent::RecoverVMinus1Start
            );
            let result = reconstruct::<T>(
                view,
                epoch,
                payload_commitment,
                common,
                shares,
                metadata,
                tracer.clone(),
            );
            if result.is_ok() {
                crate::trace_leader_event!(
                    tracer,
                    view,
                    crate::leader_trace::LeaderEvent::RecoverVMinus1End
                );
            }
            result
        });
        self.calculations.insert(view, task);
    }

    pub fn gc(&mut self, view_number: ViewNumber) {
        let keep = self.calculations.split_off(&view_number);
        for handle in self.calculations.values_mut() {
            handle.abort();
        }
        self.calculations = keep;
        self.pending = self.pending.split_off(&view_number);
        self.accumulators = self.accumulators.split_off(&view_number);
        self.reconstructed = self.reconstructed.split_off(&view_number);
    }

    /// Stop tracking `view`.
    ///
    /// Either because its payload was reconstructed (or obtained elsewhere)
    /// or because it timed out and will never be decided: record it so
    /// `handle_vid_share` ignores later shares, drop its accumulator and
    /// pending shares, and abort any in-flight reconstruction task.
    pub fn retire_view(&mut self, view: ViewNumber) {
        self.reconstructed.insert(view);
        self.pending.remove(&view);
        self.accumulators.remove(&view);
        if let Some(handle) = self.calculations.remove(&view) {
            handle.abort();
        }
    }
}

impl<T: NodeType> VidShareAccumulator<T> {
    fn new(
        payload_commitment: VidCommitment2,
        metadata: Metadata<T>,
        epoch: EpochNumber,
        expected_param: Option<AvidmGf2Param>,
        pinned_by: PinnedBy,
    ) -> Self {
        Self {
            payload_commitment,
            pinned_by,
            metadata,
            epoch,
            expected_param,
            common: None,
            shares: BTreeMap::new(),
            seen_keys: HashSet::new(),
            exhausted: false,
        }
    }

    /// Admit `share`, which claims this accumulator's commitment, from the
    /// authenticated `sender`, dropping it if it fails any intake check. The
    /// non-overlapping path stays crypto-free; a shard-range overlap is the
    /// only case that triggers per-share verification, via
    /// [`Self::resolve_conflict`].
    fn accept(&mut self, view: ViewNumber, sender: T::SignatureKey, share: VidDisperseShare2<T>) {
        debug_assert_eq!(share.payload_commitment, self.payload_commitment);
        if self.exhausted {
            return;
        }
        // The commitment binds a share's `ns_commits` but not its `param`, so a
        // Byzantine voter can pair real `ns_commits` with a forged `param` (e.g.
        // an inflated `recovery_threshold`). Pinning that common as the
        // verification oracle would reject every honest share, so reject it now.
        if let Some(expected) = &self.expected_param
            && share.common.param != *expected
        {
            warn!(%view, ?sender, "VID share common param differs from the committee's");
            return;
        }
        // The commitment hash-binds the common's `ns_commits` but not its
        // `ns_lens`, which a forwarded share could shuffle against the
        // proposal's table. Trust a common as the verification oracle only
        // after both checks; later shares must carry the same common.
        if let Some(common) = &self.common {
            if share.common != *common {
                warn!(%view, ?sender, "VID share common differs from the accumulator's");
                return;
            }
        } else if !AvidmGf2Scheme::is_consistent(&self.payload_commitment, &share.common) {
            warn!(%view, ?sender, "VID share common is inconsistent with its commitment");
            return;
        } else if !ns_lens_match_metadata(&share.common, &self.metadata.encode()) {
            warn!(
                %view, ?sender,
                "VID share namespace lengths disagree with the proposal's namespace table"
            );
            return;
        } else {
            self.common = Some(share.common.clone());
        }
        // A share whose namespaces disagree on the shard range is malformed.
        let Some(range) = share.share.range() else {
            warn!(%view, ?sender, "VID share has an inconsistent shard range");
            return;
        };
        // An empty range contributes nothing; positions past the end of the
        // encoded payload would inflate coverage without aiding decoding.
        if range.is_empty() || range.end > share.common.param.total_weights {
            warn!(%view, ?sender, ?range, "VID share has an empty or out-of-bounds shard range");
            return;
        }
        if self.seen_keys.contains(&sender) {
            return;
        }
        // Honest dispersal assigns disjoint ranges, so an overlap with an
        // admitted share proves a squat; resolve it (needs verification) below.
        let conflicts: Vec<T::SignatureKey> = self
            .shares
            .iter()
            .filter(|(_, admitted)| {
                admitted
                    .range()
                    .is_some_and(|covered| covered.start < range.end && range.start < covered.end)
            })
            .map(|(key, _)| key.clone())
            .collect();
        if conflicts.is_empty() {
            self.seen_keys.insert(sender.clone());
            self.shares.insert(sender, share.share);
            return;
        }
        self.resolve_conflict(view, sender, share, conflicts);
    }

    /// Resolve a shard-range collision between the incoming `share` and the
    /// already-admitted `conflicts`: verify each against the commitment-bound
    /// common, evict those that fail, and admit the newcomer only if it
    /// verifies and no surviving share still covers its range.
    fn resolve_conflict(
        &mut self,
        view: ViewNumber,
        sender: T::SignatureKey,
        share: VidDisperseShare2<T>,
        conflicts: Vec<T::SignatureKey>,
    ) {
        // A conflict implies a prior admission, which pinned the common.
        let Some(common) = self.common.clone() else {
            return;
        };
        // The sender has used its one slot regardless of the outcome.
        self.seen_keys.insert(sender.clone());
        let mut survivor = false;
        for key in conflicts {
            let verified = self
                .shares
                .get(&key)
                .is_some_and(|admitted| share_verifies(&common, admitted));
            if verified {
                survivor = true;
            } else {
                warn!(%view, ?key, "evicting unverifiable VID share squatting a shard range");
                self.shares.remove(&key);
            }
        }
        // A verified share still covers the contested range: the newcomer would
        // double-cover it, so drop it.
        if survivor {
            return;
        }
        if share_verifies(&common, &share.share) {
            self.shares.insert(sender, share.share);
        } else {
            warn!(%view, ?sender, "dropping unverifiable VID share at intake conflict");
        }
    }

    fn ranges(&self) -> impl Iterator<Item = &Range<usize>> {
        // Admitted shares always have a consistent range (checked in `accept`).
        self.shares.values().filter_map(AvidmGf2Share::range)
    }

    /// Number of shard positions covered by the admitted shares; exact
    /// because their ranges are disjoint.
    fn coverage(&self) -> usize {
        self.ranges().map(ExactSizeIterator::len).sum()
    }

    fn has_enough_shares(&self) -> bool {
        self.common
            .as_ref()
            .is_some_and(|common| self.coverage() >= common.param.recovery_threshold)
    }
}

/// Decode the shares and accept the result only if it re-commits to
/// `payload_commitment`. On failure, report the shares that fail
/// verification against the commitment (each share is self-authenticating
/// via its merkle proofs) so they can be weeded out. If every share
/// verifies, the payload is unrecoverable: the shares cover the recovery
/// threshold with disjoint ranges, so the disperser committed to a
/// non-codeword and no share subset can ever succeed.
fn reconstruct<T: NodeType>(
    view: ViewNumber,
    epoch: EpochNumber,
    payload_commitment: VidCommitment2,
    common: AvidmGf2Common,
    shares: Vec<(T::SignatureKey, AvidmGf2Share)>,
    metadata: Metadata<T>,
    tracer: Option<crate::leader_trace::LeaderTracerHandle>,
) -> ReconstructResult<T> {
    let (keys, shares): (Vec<_>, Vec<_>) = shares.into_iter().unzip();
    if let Some(bytes) =
        decode_and_recommit::<T>(view, &common, &shares, &payload_commitment, &metadata)
    {
        crate::trace_leader_event!(
            tracer,
            view,
            crate::leader_trace::LeaderEvent::RecoverVMinus1DecodeEnd
        );
        let payload = T::BlockPayload::from_bytes(&bytes, &metadata);
        let tx_commitments = payload.transaction_commitments(&metadata);
        let output = ObtainedPayload {
            view,
            epoch,
            payload_commitment,
            payload,
            metadata,
            tx_commitments,
        };
        return Ok(output);
    }

    let bad_share_keys: Vec<_> = keys
        .into_iter()
        .zip(&shares)
        .filter(|(_, share)| !share_verifies(&common, share))
        .map(|(key, _)| key)
        .collect();

    let kind = if bad_share_keys.is_empty() {
        warn!(
            %view,
            %payload_commitment,
            "verified shares cannot decode to a payload matching the commitment"
        );
        VidReconstructErrorKind::Unrecoverable
    } else {
        warn!(
            %view,
            %payload_commitment,
            ?bad_share_keys,
            "weeded out VID shares that failed verification"
        );
        VidReconstructErrorKind::AwaitingShares
    };
    Err(VidReconstructError {
        view,
        payload_commitment,
        kind,
        bad_share_keys,
    })
}

/// Decode `shares` and return the payload bytes only if they re-commit to
/// `payload_commitment`. Recovery alone does not bind the decoded bytes
/// to the commitment: a Byzantine disperser can commit to a non-codeword,
/// and a bad share poisons the erasure decoding.
fn decode_and_recommit<T: NodeType>(
    view: ViewNumber,
    common: &AvidmGf2Common,
    shares: &[AvidmGf2Share],
    payload_commitment: &VidCommitment2,
    metadata: &Metadata<T>,
) -> Option<Vec<u8>> {
    let bytes = match AvidmGf2Scheme::recover(common, shares) {
        Ok(bytes) => bytes,
        Err(err) => {
            warn!(%view, %err, "VID recovery failed");
            return None;
        },
    };
    matches_commitment(
        view,
        &common.param,
        &metadata.encode(),
        &bytes,
        payload_commitment,
    )
    .then_some(bytes)
}

/// Whether `payload` is what `payload_commitment` commits to, sliced by the
/// header's namespace table `ns_table`.
///
/// The commitment covers each namespace's bytes but not the table, and
/// `parse_ns_table` substitutes one namespace spanning the payload for a table
/// sized for another payload. For a one-namespace block that substitute is the
/// honest table, so a header claiming a shorter or longer payload would
/// recommit to the same bytes and decide with a table that misdescribes them.
/// Refuse such a table instead: it must be sized for `payload`, or be no
/// table at all, which the disperser also reads as one namespace.
pub(crate) fn matches_commitment(
    view: ViewNumber,
    param: &AvidmGf2Param,
    ns_table: &[u8],
    payload: &[u8],
    payload_commitment: &VidCommitment2,
) -> bool {
    if let Some(claimed) = ns_table_payload_byte_len(ns_table)
        && claimed != payload.len()
    {
        warn!(
            %view,
            claimed,
            actual = payload.len(),
            "namespace table is sized for another payload"
        );
        return false;
    }
    let ns_table = parse_ns_table(payload.len(), ns_table);
    match AvidmGf2Scheme::commit(param, payload, ns_table) {
        Ok((recomputed, _)) if recomputed == *payload_commitment => true,
        Ok((recomputed, _)) => {
            warn!(
                %view,
                expected = %payload_commitment,
                %recomputed,
                "payload does not match the payload commitment"
            );
            false
        },
        Err(err) => {
            warn!(%view, %err, "failed to recommit VID payload");
            false
        },
    }
}

/// Whether `share` verifies against a `common` already known to be hash-bound
/// to the commitment (a `verify_with_verified_common` success).
fn share_verifies(common: &AvidmGf2Common, share: &AvidmGf2Share) -> bool {
    matches!(
        AvidmGf2Scheme::verify_share_with_verified_common(common, share),
        Ok(Ok(()))
    )
}