solana-ledger 4.2.0-beta.1

//! Relevant structures and filtering logic to be applied to ingested and recovered shreds.
//! as well as shred level feature flag detection

use {
    super::{
        DATA_SHREDS_PER_FEC_BLOCK, Error, Payload, ReedSolomonCache, Shred, ShredFetchStats,
        ShredFlags, ShredType, ShredVariant, layout, merkle,
    },
    crate::blockstore,
    agave_feature_set::discard_unexpected_data_complete_shreds,
    assert_matches::debug_assert_matches,
    solana_clock::Slot,
    solana_epoch_schedule::EpochSchedule,
    solana_perf::packet::PacketRef,
    solana_pubkey::Pubkey,
    solana_runtime::bank::Bank,
    solana_streamer::{evicting_sender::EvictingSender, streamer::ChannelSend},
    std::{
        sync::{
            Arc,
            atomic::{AtomicU8, Ordering},
        },
        time::{Duration, Instant},
    },
};

/// Per-slot shred index limits used by bank-aware shred filtering.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct ShredLimits {
    /// Exclusive upper bound for data shred indexes in a slot.
    max_data_shreds_per_slot: u32,
    /// Exclusive upper bound for code shred indexes in a slot.
    max_code_shreds_per_slot: u32,
}

impl ShredLimits {
    /// Slot-independent maximum cap used by tests.
    #[cfg(test)]
    const DEFAULT: Self = Self {
        max_data_shreds_per_slot: super::MAX_DATA_SHREDS_PER_SLOT as u32,
        max_code_shreds_per_slot: super::MAX_CODE_SHREDS_PER_SLOT as u32,
    };

    /// Creates shred limits with exclusive upper bounds for data and code indexes.
    const fn new(max_data_shreds_per_slot: u32, max_code_shreds_per_slot: u32) -> Self {
        Self {
            max_data_shreds_per_slot,
            max_code_shreds_per_slot,
        }
    }

    /// Returns true if `index` is below the data shred limit.
    #[inline]
    const fn is_data_index_in_bounds(self, index: u32) -> bool {
        index < self.max_data_shreds_per_slot
    }

    /// Returns true if `index` is below the code shred limit.
    #[inline]
    const fn is_code_index_in_bounds(self, index: u32) -> bool {
        index < self.max_code_shreds_per_slot
    }
}

/// Controls turbine and repair behavior for testing network partitions.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum TurbineModeKind {
    /// Normal operation - turbine and repair both enabled.
    #[default]
    Enabled = 0,
    /// Turbine disabled but repair still works.
    TurbineDisabled = 1,
    /// Both turbine and repair disabled.
    TurbineAndRepairDisabled = 2,
}

impl TurbineModeKind {
    fn should_discard_packet(self, is_repair: bool) -> bool {
        match self {
            Self::Enabled => false,
            Self::TurbineDisabled => !is_repair,
            Self::TurbineAndRepairDisabled => true,
        }
    }
}

impl From<u8> for TurbineModeKind {
    fn from(value: u8) -> Self {
        match value {
            0 => Self::Enabled,
            1 => Self::TurbineDisabled,
            2 => Self::TurbineAndRepairDisabled,
            _ => Self::Enabled,
        }
    }
}

impl From<TurbineModeKind> for u8 {
    fn from(mode: TurbineModeKind) -> Self {
        mode as u8
    }
}

/// Thread-safe wrapper around [`TurbineModeKind`] for dynamically controlling
/// turbine and repair behavior at runtime.
#[derive(Clone, Debug)]
pub struct TurbineMode(Arc<AtomicU8>);

impl TurbineMode {
    pub fn new(kind: TurbineModeKind) -> Self {
        Self(Arc::new(AtomicU8::new(kind as u8)))
    }

    pub fn get(&self) -> TurbineModeKind {
        TurbineModeKind::from(self.0.load(Ordering::Relaxed))
    }

    pub fn set(&self, kind: TurbineModeKind) {
        self.0.store(kind as u8, Ordering::Relaxed);
    }
}

impl Default for TurbineMode {
    fn default() -> Self {
        Self::new(TurbineModeKind::default())
    }
}

/// Bank-backed shred limit lookup shared by fetch and recovery paths.
struct ShredLimitContext {
    root_bank: Arc<Bank>,
    #[cfg(test)]
    shred_limits_override: Option<ShredLimits>,
}

impl ShredLimitContext {
    fn new(root_bank: Arc<Bank>) -> Self {
        Self {
            root_bank,
            #[cfg(test)]
            shred_limits_override: None,
        }
    }

    /// Returns the shred limits for `slot` as derived from the current root bank.
    fn shred_limits(&self, slot: Slot) -> ShredLimits {
        #[cfg(test)]
        if let Some(shred_limits) = self.shred_limits_override {
            return shred_limits;
        }

        ShredLimits::new(
            self.root_bank.max_data_shreds_per_slot_for_slot(slot),
            self.root_bank.max_code_shreds_per_slot_for_slot(slot),
        )
    }
}

pub struct ShredFilterContext {
    last_updated: Instant,

    // Fields for slot range filtering updated via the root bank
    root: Slot,
    max_slot: Slot,
    epoch_schedule: EpochSchedule,

    // Static from startup, filter out shreds of invalid version
    shred_version: u16,

    // Filter on turbine / repair packets
    turbine_mode: Option<TurbineMode>,
    // Cache of above filter so that we don't have to incur an
    // atomic load for each packet. This is recached during `maybe_update`
    cached_turbine_mode: TurbineModeKind,

    // Whether to discard shreds that are not end of FEC set
    // but specify DATA_COMPLETE
    discard_unexpected_data_complete_shreds_feature_slot: Option<Slot>,

    // The shred limits per slot, derived from the root bank for each shred slot.
    shred_limit_context: ShredLimitContext,

    pub stats: ShredFetchStats,
}

impl ShredFilterContext {
    pub fn new(root_bank: Arc<Bank>, shred_version: u16) -> Self {
        Self::new_with_turbine_mode(root_bank, shred_version, None)
    }

    pub fn new_with_turbine_mode(
        root_bank: Arc<Bank>,
        shred_version: u16,
        turbine_mode: Option<TurbineMode>,
    ) -> Self {
        let root = root_bank.slot();
        let max_slot = max_shred_slot(root, root_bank.get_slots_in_epoch(root_bank.epoch()));
        let discard_unexpected_data_complete_shreds_feature_slot = root_bank
            .feature_set
            .activated_slot(&discard_unexpected_data_complete_shreds::id());
        let cached_turbine_mode = turbine_mode
            .as_ref()
            .map(TurbineMode::get)
            .unwrap_or_default();

        Self {
            last_updated: Instant::now(),
            root,
            max_slot,
            epoch_schedule: root_bank.epoch_schedule().clone(),
            shred_version,
            turbine_mode,
            cached_turbine_mode,
            discard_unexpected_data_complete_shreds_feature_slot,
            shred_limit_context: ShredLimitContext::new(root_bank),
            stats: ShredFetchStats::default(),
        }
    }

    /// Periodically update filtering context based on the root bank (no more than once per slot)
    /// This is done to amortize the cost over multiple shreds, and delay in updating is completely
    /// acceptable as max shred / feature flag filtering has tolerance on the order of an epoch.
    pub fn maybe_update(&mut self, root_bank: Arc<Bank>) {
        if let Some(turbine_mode) = self.turbine_mode.as_ref() {
            self.cached_turbine_mode = turbine_mode.get();
        }

        if self.last_updated.elapsed().as_nanos() > root_bank.ns_per_slot {
            self.last_updated = Instant::now();
            self.root = root_bank.slot();
            self.max_slot =
                max_shred_slot(self.root, root_bank.get_slots_in_epoch(root_bank.epoch()));
            debug_assert!(self.root < self.max_slot);

            self.epoch_schedule = root_bank.epoch_schedule().clone();
            self.discard_unexpected_data_complete_shreds_feature_slot = root_bank
                .feature_set
                .activated_slot(&discard_unexpected_data_complete_shreds::id());
            self.shred_limit_context = ShredLimitContext::new(root_bank);
        }
    }

    fn shred_limits(&self, slot: Slot) -> ShredLimits {
        self.shred_limit_context.shred_limits(slot)
    }

    #[cfg(test)]
    fn set_shred_limits_for_tests(&mut self, shred_limits: ShredLimits) {
        self.shred_limit_context.shred_limits_override = Some(shred_limits);
    }

    pub fn maybe_submit_stats(&mut self, metric_name: &'static str, cadence: Duration) -> bool {
        self.stats.maybe_submit(metric_name, cadence)
    }

    #[must_use]
    pub fn should_discard_packet<'a, P>(&mut self, packet: P) -> bool
    where
        P: Into<PacketRef<'a>>,
    {
        let packet = packet.into();
        if self
            .cached_turbine_mode
            .should_discard_packet(packet.meta().repair())
        {
            return true;
        }
        let Some(shred) = layout::get_shred(packet) else {
            self.stats.index_overrun += 1;
            return true;
        };
        self.should_discard_shred(shred)
    }

    #[must_use]
    pub fn should_discard_shred(&mut self, shred: &[u8]) -> bool {
        match layout::get_version(shred) {
            None => {
                self.stats.index_overrun += 1;
                return true;
            }
            Some(version) => {
                if version != self.shred_version {
                    self.stats.shred_version_mismatch += 1;
                    return true;
                }
            }
        }
        let Ok(shred_variant) = layout::get_shred_variant(shred) else {
            self.stats.bad_shred_type += 1;
            return true;
        };
        let slot = match layout::get_slot(shred) {
            Some(slot) => {
                if slot > self.max_slot {
                    self.stats.slot_out_of_range += 1;
                    return true;
                }
                slot
            }
            None => {
                self.stats.slot_bad_deserialize += 1;
                return true;
            }
        };
        let Some(index) = layout::get_index(shred) else {
            self.stats.index_bad_deserialize += 1;
            return true;
        };
        let Some(fec_set_index) = layout::get_fec_set_index(shred) else {
            self.stats.fec_set_index_bad_deserialize += 1;
            return true;
        };
        let shred_limits = self.shred_limits(slot);

        match ShredType::from(shred_variant) {
            ShredType::Code => {
                if !shred_limits.is_code_index_in_bounds(index) {
                    self.stats.index_out_of_bounds += 1;
                    return true;
                }
                if slot <= self.root {
                    self.stats.slot_out_of_range += 1;
                    return true;
                }

                let Ok(erasure_config) = layout::get_erasure_config(shred) else {
                    self.stats.erasure_config_bad_deserialize += 1;
                    return true;
                };

                if !erasure_config.is_fixed() {
                    self.stats.misaligned_erasure_config += 1;
                    return true;
                }
            }
            ShredType::Data => {
                if !shred_limits.is_data_index_in_bounds(index) {
                    self.stats.index_out_of_bounds += 1;
                    return true;
                }
                let Some(parent_offset) = layout::get_parent_offset(shred) else {
                    self.stats.bad_parent_offset += 1;
                    return true;
                };
                let Some(parent) = slot.checked_sub(Slot::from(parent_offset)) else {
                    self.stats.bad_parent_offset += 1;
                    return true;
                };
                if !blockstore::verify_shred_slots(slot, parent, self.root) {
                    self.stats.slot_out_of_range += 1;
                    return true;
                }

                let Ok(shred_flags) = layout::get_flags(shred) else {
                    self.stats.shred_flags_bad_deserialize += 1;
                    return true;
                };

                let expected_data_complete_index = fec_set_index
                    .checked_add(DATA_SHREDS_PER_FEC_BLOCK as u32)
                    .and_then(|index| index.checked_sub(1));
                if shred_flags.contains(ShredFlags::DATA_COMPLETE_SHRED)
                    && (expected_data_complete_index != Some(index))
                {
                    self.stats.unexpected_data_complete_shred += 1;

                    if check_feature_activation(
                        self.discard_unexpected_data_complete_shreds_feature_slot,
                        slot,
                        &self.epoch_schedule,
                    ) {
                        return true;
                    }
                }

                if shred_flags.contains(ShredFlags::LAST_SHRED_IN_SLOT)
                    && !check_last_data_shred_index(index)
                {
                    self.stats.misaligned_last_data_index += 1;
                    return true;
                }
            }
        }

        if !check_fixed_fec_set(index, fec_set_index) {
            self.stats.misaligned_fec_set += 1;
            return true;
        }

        match shred_variant {
            ShredVariant::MerkleCode { .. } => {
                self.stats.num_shreds_merkle_code_chained =
                    self.stats.num_shreds_merkle_code_chained.saturating_add(1);
            }
            ShredVariant::MerkleData { .. } => {
                self.stats.num_shreds_merkle_data_chained =
                    self.stats.num_shreds_merkle_data_chained.saturating_add(1);
            }
        }
        false
    }
}

/// Returns true if the feature is effective for the shred slot.
/// Note: unlike normal feature flags, shred feature flags take effect 1 epoch after activation.
#[must_use]
pub fn check_feature_activation_from_bank(
    feature: &Pubkey,
    shred_slot: Slot,
    root_bank: &Bank,
) -> bool {
    check_feature_activation(
        root_bank.feature_set.activated_slot(feature),
        shred_slot,
        root_bank.epoch_schedule(),
    )
}

/// Returns true if the feature is effective for the shred slot.
/// Note: unlike normal feature flags, shred feature flags take effect 1 epoch after activation.
#[must_use]
pub fn check_feature_activation(
    feature_slot: Option<Slot>,
    shred_slot: Slot,
    epoch_schedule: &EpochSchedule,
) -> bool {
    let Some(feature_slot) = feature_slot else {
        return false;
    };

    let feature_epoch = epoch_schedule.get_epoch(feature_slot);
    let shred_epoch = epoch_schedule.get_epoch(shred_slot);
    feature_epoch < shred_epoch
}

/// The maximum shred slot we allow for ingest given a current `root` slot.
fn max_shred_slot(root: Slot, slots_per_epoch: Slot) -> Slot {
    // When running with very short epochs (e.g. for testing), we want to avoid
    // filtering out shreds that we actually need. This value was chosen empirically
    // because it's large enough to protect against observed short epoch problems
    // while being small enough to keep the overhead small on deduper, blockstore,
    // etc.
    const MAX_SHRED_DISTANCE_MINIMUM: Slot = 500;
    // Allow shreds up to half an epoch into the future to support catching up to the tip of the cluster.
    root.saturating_add(MAX_SHRED_DISTANCE_MINIMUM.max(slots_per_epoch / 2))
}

/// Returns true if `index` and `fec_set_index` are valid under the assumption that
/// all erasure sets contain exactly `DATA_SHREDS_PER_FEC_BLOCK` data and coding shreds:
/// - `index` is between `fec_set_index` and `fec_set_index + DATA_SHREDS_PER_FEC_BLOCK`
/// - `fec_set_index` is a multiple of `DATA_SHREDS_PER_FEC_BLOCK`
fn check_fixed_fec_set(index: u32, fec_set_index: u32) -> bool {
    let Some(fec_set_end_exclusive) = fec_set_index.checked_add(DATA_SHREDS_PER_FEC_BLOCK as u32)
    else {
        return false;
    };
    index >= fec_set_index
        && index < fec_set_end_exclusive
        && fec_set_index.is_multiple_of(DATA_SHREDS_PER_FEC_BLOCK as u32)
}

/// Returns true if `index` of the last data shred is valid under the assumption that
/// all erasure sets contain exactly `DATA_SHREDS_PER_FEC_BLOCK` data and coding shreds:
/// - `index + 1` must be a multiple of `DATA_SHREDS_PER_FEC_BLOCK`
///
/// Note: this check is critical to verify that the last fec set is sufficiently sized.
/// This is checked during shred ingest with `enforce_fixed_fec_set` active.
fn check_last_data_shred_index(index: u32) -> bool {
    (index + 1).is_multiple_of(DATA_SHREDS_PER_FEC_BLOCK as u32)
}

/// Holds the context to perform filtering on shred recovery
pub struct ShredRecoveryContext {
    /// Used to perform RS erasure code recovery
    pub reed_solomon_cache: ReedSolomonCache,
    /// Sender to retransmit the recovered shreds
    retransmit_sender: EvictingSender<Vec<Payload>>,
    /// Used for filtering recovered shreds
    shred_filter_ctx: ShredFilterContext,
}

impl ShredRecoveryContext {
    pub fn new(
        reed_solomon_cache: ReedSolomonCache,
        retransmit_sender: EvictingSender<Vec<Payload>>,
        root_bank: Arc<Bank>,
        shred_version: u16,
    ) -> Self {
        let shred_filter_ctx = ShredFilterContext::new(root_bank, shred_version);
        Self {
            reed_solomon_cache,
            retransmit_sender,
            shred_filter_ctx,
        }
    }

    /// Periodically (no more than once per slot) update the context used for filtering
    /// recovered shreds. This is not latency sensitive as feature flags and slot filtering windows
    /// have an epoch order tolerance
    pub fn maybe_update(&mut self, root_bank: Arc<Bank>) {
        self.shred_filter_ctx.maybe_update(root_bank);
    }

    /// Submit stats at most every two seconds
    pub fn maybe_submit_stats(&mut self) {
        self.shred_filter_ctx
            .maybe_submit_stats("shred-recovery", Duration::from_secs(2));
    }

    /// Recover shreds and apply filtering to the recovered shreds
    pub fn recover<T: IntoIterator<Item = Shred>>(
        &mut self,
        shreds: T,
        recovered_shreds: &mut Vec<Payload>,
        recovered_data_shreds: &mut Vec<Shred>,
    ) -> Result<(), Error> {
        let shreds = shreds
            .into_iter()
            .filter(|shred| !self.shred_filter_ctx.should_discard_shred(shred.payload()))
            .map(|shred| {
                debug_assert_matches!(
                    shred.common_header().shred_variant,
                    ShredVariant::MerkleCode { .. } | ShredVariant::MerkleData { .. }
                );
                merkle::Shred::try_from(shred)
            })
            .collect::<Result<_, _>>()?;

        // With Merkle shreds, leader signs the Merkle root of the erasure batch
        // and all shreds within the same erasure batch have the same signature.
        // For recovered shreds, the (unique) signature is copied from shreds which
        // were received from turbine (or repair) and are already sig-verified.
        // The same signature also verifies for recovered shreds because when
        // reconstructing the Merkle tree for the erasure batch, we will obtain the
        // same Merkle root.
        let shreds = merkle::recover(shreds, &self.reed_solomon_cache)?;
        shreds
            .filter_map(|shred| shred.ok().map(Shred::from))
            .filter(|shred| !self.should_discard_shred(shred))
            .for_each(|shred| {
                // All shreds should be retransmitted, but because there are no
                // more missing data shreds in the erasure batch, coding shreds
                // are not stored in blockstore.
                match shred.shred_type() {
                    ShredType::Code => {
                        // Don't need Arc overhead here!
                        recovered_shreds.push(shred.into_payload());
                    }
                    ShredType::Data => {
                        // Verify that the cloning is cheap here.
                        recovered_shreds.push(shred.payload().clone());
                        recovered_data_shreds.push(shred);
                    }
                }
            });
        Ok(())
    }
    /// Send recovered shreds for retransmit
    pub fn try_retransmit_shreds(&self, recovered_shreds: Vec<Payload>) {
        if !recovered_shreds.is_empty() {
            let _ = self.retransmit_sender.try_send(recovered_shreds);
        }
    }

    /// If SIMD-0337 is active, apply filtering rules to recovered shreds
    pub fn should_discard_shred(&mut self, shred: &Shred) -> bool {
        if !check_feature_activation(
            self.shred_filter_ctx
                .discard_unexpected_data_complete_shreds_feature_slot,
            shred.slot(),
            &self.shred_filter_ctx.epoch_schedule,
        ) {
            // Only apply filtering rules to recovered shreds when SIMD-0337 is active
            return false;
        }

        self.shred_filter_ctx.should_discard_shred(shred.payload())
    }
}

#[cfg(test)]
mod tests {
    use {
        super::{
            super::{Shred, make_merkle_shreds_for_tests},
            *,
        },
        crate::{
            genesis_utils::create_genesis_config,
            shred::{MAX_CODE_SHREDS_PER_SLOT, tests::*},
        },
        assert_matches::assert_matches,
        itertools::Itertools,
        solana_leader_schedule::SlotLeader,
        solana_perf::packet::{Packet, PacketFlags},
        solana_runtime::{
            bank::Bank,
            slot_params::{slot_time_feature_gates, slot_time_feature_ids},
        },
        std::{
            io::{Cursor, Seek, SeekFrom, Write},
            sync::Arc,
            time::Duration,
        },
        test_case::test_case,
    };

    fn new_test_bank(slot: Slot, discard_unexpected_data_complete_shreds: bool) -> Arc<Bank> {
        let genesis_config = create_genesis_config(1).genesis_config;
        let mut bank = Bank::new_for_tests(&genesis_config);
        let feature_id = agave_feature_set::discard_unexpected_data_complete_shreds::id();
        bank.deactivate_feature(&feature_id);
        if discard_unexpected_data_complete_shreds {
            bank.activate_feature(&feature_id);
        }
        let (bank, bank_forks) = bank.wrap_with_bank_forks_for_tests();
        if slot == 0 {
            bank
        } else {
            Bank::new_from_parent_with_bank_forks(&bank_forks, bank, SlotLeader::default(), slot)
        }
    }

    fn deactivate_slot_time_features(bank: &mut Bank) {
        for feature_id in slot_time_feature_ids() {
            bank.deactivate_feature(&feature_id);
        }
    }

    fn shred_filter_for_tests(
        feature_ids: impl IntoIterator<Item = Pubkey>,
    ) -> (ShredFilterContext, Slot) {
        let genesis_config = create_genesis_config(1).genesis_config;
        let mut root_bank = Bank::new_for_tests(&genesis_config);
        deactivate_slot_time_features(&mut root_bank);
        for feature_id in feature_ids {
            root_bank.activate_feature(&feature_id);
        }
        let first_epoch_slot = root_bank.epoch_schedule().get_first_slot_in_epoch(1);
        let (root_bank, _) = root_bank.wrap_with_bank_forks_for_tests();
        (ShredFilterContext::new(root_bank, 0), first_epoch_slot)
    }

    #[test_case(true ; "last_in_slot")]
    #[test_case(false ; "not_last_in_slot")]
    fn test_should_discard_shred(is_last_in_slot: bool) {
        agave_logger::setup();
        let mut rng = rand::rng();
        let slot = 200;
        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            slot,
            1200 * 5, // data_size
            is_last_in_slot,
        )
        .unwrap();
        let shreds: Vec<_> = shreds.into_iter().map(Shred::from).collect();
        assert_eq!(shreds.iter().map(Shred::fec_set_index).dedup().count(), 1);

        assert_matches!(shreds[0].shred_type(), ShredType::Data);
        let parent_slot = shreds[0].parent().unwrap();
        let shred_version = shreds[0].common_header().version;

        let root_bank = new_test_bank(0, false);
        let parent_exceeded_root_bank = new_test_bank(parent_slot + 1, false);
        let slot_root_bank = new_test_bank(slot, false);
        let mut packet = Packet::default();

        {
            let shred = shreds.first().unwrap();
            assert_eq!(shred.shred_type(), ShredType::Data);
            shred.copy_to_packet(&mut packet);
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(!shred_filter_context.should_discard_packet(&packet));
        }
        {
            let mut packet = packet.clone();
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            packet.meta_mut().size = OFFSET_OF_SHRED_VARIANT;
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_overrun, 1);

            packet.meta_mut().size = OFFSET_OF_SHRED_INDEX;
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_overrun, 2);

            packet.meta_mut().size = OFFSET_OF_SHRED_INDEX + 1;
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_overrun, 3);

            packet.meta_mut().size = OFFSET_OF_SHRED_INDEX + SIZE_OF_SHRED_INDEX - 1;
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_overrun, 4);

            packet.meta_mut().size = OFFSET_OF_SHRED_INDEX + SIZE_OF_SHRED_INDEX + 2;
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_overrun, 5);
        }
        {
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version.wrapping_add(1));
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.shred_version_mismatch, 1);
        }
        {
            let mut shred_filter_context =
                ShredFilterContext::new(parent_exceeded_root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.slot_out_of_range, 1);
        }
        {
            let parent_offset = 0u16;
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor.seek(SeekFrom::Start(83)).unwrap();
                cursor.write_all(&parent_offset.to_le_bytes()).unwrap();
            }
            assert_eq!(
                layout::get_parent_offset(packet.data(..).unwrap()),
                Some(parent_offset)
            );
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.slot_out_of_range, 1);
        }
        {
            let parent_offset = u16::try_from(slot + 1).unwrap();
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor.seek(SeekFrom::Start(83)).unwrap();
                cursor.write_all(&parent_offset.to_le_bytes()).unwrap();
            }
            assert_eq!(
                layout::get_parent_offset(packet.data(..).unwrap()),
                Some(parent_offset)
            );
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.bad_parent_offset, 1);
        }
        {
            let index = u32::MAX - 10;
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor
                    .seek(SeekFrom::Start(OFFSET_OF_SHRED_INDEX as u64))
                    .unwrap();
                cursor.write_all(&index.to_le_bytes()).unwrap();
            }
            assert_eq!(layout::get_index(packet.data(..).unwrap()), Some(index));
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_out_of_bounds, 1);
        }

        {
            let shred = shreds.last().unwrap();
            assert_eq!(shred.shred_type(), ShredType::Code);
            shreds.last().unwrap().copy_to_packet(&mut packet);
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(!shred_filter_context.should_discard_packet(&packet));
        }
        {
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version.wrapping_add(1));
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.shred_version_mismatch, 1);
        }
        {
            let mut shred_filter_context =
                ShredFilterContext::new(slot_root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.slot_out_of_range, 1);
        }
        {
            let index = u32::try_from(MAX_CODE_SHREDS_PER_SLOT).unwrap();
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor
                    .seek(SeekFrom::Start(OFFSET_OF_SHRED_INDEX as u64))
                    .unwrap();
                cursor.write_all(&index.to_le_bytes()).unwrap();
            }
            assert_eq!(layout::get_index(packet.data(..).unwrap()), Some(index));
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.index_out_of_bounds, 1);
        }
    }

    #[test]
    fn test_recovery_shred_limits() {
        agave_logger::setup();
        let mut rng = rand::rng();
        let slot = 200;
        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            slot,
            1200 * 5, // data_size
            false,    // is_last_in_slot
        )
        .unwrap();
        let shreds: Vec<_> = shreds.into_iter().map(Shred::from).collect();
        let coding_shreds: Vec<_> = shreds
            .into_iter()
            .filter(|shred| shred.shred_type() == ShredType::Code)
            .collect();
        let shred_version = coding_shreds[0].common_header().version;

        // Feed recovery only coding shreds. Without the custom limit below, this
        // is enough parity to recover the missing data shreds.
        let max_code_shreds_per_slot = coding_shreds[0].index();
        let (dummy_retransmit_sender, _) = EvictingSender::new_bounded(0);
        let mut shred_recovery_context = ShredRecoveryContext::new(
            ReedSolomonCache::default(),
            dummy_retransmit_sender,
            new_test_bank(0, false),
            shred_version,
        );
        shred_recovery_context
            .shred_filter_ctx
            .set_shred_limits_for_tests(ShredLimits::new(
                ShredLimits::DEFAULT.max_data_shreds_per_slot,
                // Setting the limit to the first coding index makes every
                // coding shred in this batch invalid.
                max_code_shreds_per_slot,
            ));
        let mut recovered_shreds = Vec::new();
        let mut recovered_data_shreds = Vec::new();

        // Recovery sees insufficient parity shreds and cannot recover the FEC.
        assert_matches!(
            shred_recovery_context.recover(
                coding_shreds,
                &mut recovered_shreds,
                &mut recovered_data_shreds,
            ),
            Err(Error::Erasure(
                reed_solomon_erasure::Error::TooFewParityShards
            ))
        );
        assert!(recovered_shreds.is_empty());
        assert!(recovered_data_shreds.is_empty());
    }

    #[test]
    fn test_should_discard_packet_with_turbine_mode() {
        agave_logger::setup();
        let mut rng = rand::rng();
        let root_bank = new_test_bank(0, false);
        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            42,
            1200 * 5, // data_size
            false,    // is_last_in_slot
        )
        .unwrap();
        let shred = Shred::from(shreds[0].clone());
        let shred_version = shred.common_header().version;
        let turbine_mode = TurbineMode::new(TurbineModeKind::TurbineDisabled);
        let mut packet = Packet::default();
        shred.copy_to_packet(&mut packet);

        let mut shred_filter_context = ShredFilterContext::new_with_turbine_mode(
            root_bank.clone(),
            shred_version,
            Some(turbine_mode.clone()),
        );
        assert!(shred_filter_context.should_discard_packet(&packet));

        packet.meta_mut().flags.insert(PacketFlags::REPAIR);
        assert!(!shred_filter_context.should_discard_packet(&packet));

        turbine_mode.set(TurbineModeKind::TurbineAndRepairDisabled);
        shred_filter_context.last_updated = Instant::now() - Duration::from_secs(1);
        shred_filter_context.maybe_update(root_bank.clone());
        assert!(shred_filter_context.should_discard_packet(&packet));

        packet.meta_mut().flags.remove(PacketFlags::REPAIR);
        turbine_mode.set(TurbineModeKind::Enabled);
        shred_filter_context.last_updated = Instant::now() - Duration::from_secs(1);
        shred_filter_context.maybe_update(root_bank.clone());
        assert!(!shred_filter_context.should_discard_packet(&packet));
    }

    #[test]
    fn test_should_discard_shred_fec_set_checks() {
        agave_logger::setup();
        let mut rng = rand::rng();
        let slot = 200;
        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            slot,
            1200 * 5, // data_size
            false,    // is_last_in_slot
        )
        .unwrap();
        let shreds: Vec<_> = shreds.into_iter().map(Shred::from).collect();
        assert_eq!(shreds.iter().map(Shred::fec_set_index).dedup().count(), 1);

        assert_matches!(shreds[0].shred_type(), ShredType::Data);
        let shred_version = shreds[0].common_header().version;
        let root_bank = new_test_bank(0, false);

        {
            let mut packet = Packet::default();
            shreds[0].copy_to_packet(&mut packet);

            let bad_fec_set_index = 5u32;
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor
                    .seek(SeekFrom::Start(OFFSET_OF_FEC_SET_INDEX as u64))
                    .unwrap();
                cursor.write_all(&bad_fec_set_index.to_le_bytes()).unwrap();
            }

            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.misaligned_fec_set, 1);
        }

        {
            let mut packet = Packet::default();
            shreds[0].copy_to_packet(&mut packet);

            let fec_set_index = 64u32;
            let bad_index = 100u32;
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor
                    .seek(SeekFrom::Start(OFFSET_OF_SHRED_INDEX as u64))
                    .unwrap();
                cursor.write_all(&bad_index.to_le_bytes()).unwrap();
                cursor
                    .seek(SeekFrom::Start(OFFSET_OF_FEC_SET_INDEX as u64))
                    .unwrap();
                cursor.write_all(&fec_set_index.to_le_bytes()).unwrap();
            }

            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.misaligned_fec_set, 1);
        }

        {
            let code_shred = shreds
                .iter()
                .find(|s| s.shred_type() == ShredType::Code)
                .unwrap();
            let mut packet = Packet::default();
            code_shred.copy_to_packet(&mut packet);

            let bad_num_data = 16u16;
            {
                let mut cursor = Cursor::new(packet.buffer_mut());
                cursor
                    .seek(SeekFrom::Start(OFFSET_OF_NUM_DATA as u64))
                    .unwrap();
                cursor.write_all(&bad_num_data.to_le_bytes()).unwrap();
            }

            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            assert!(shred_filter_context.should_discard_packet(&packet));
            assert_eq!(shred_filter_context.stats.misaligned_erasure_config, 1);
        }

        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            slot,
            1200 * 5, // data_size
            true,     // is_last_in_slot
        )
        .unwrap();
        let shreds: Vec<_> = shreds.into_iter().map(Shred::from).collect();
        let shred_version = shreds[0].common_header().version;
        let data_shreds: Vec<_> = shreds
            .iter()
            .filter(|s| s.shred_type() == ShredType::Data)
            .collect();
        let last_data_shred = data_shreds.last().unwrap();
        assert!(last_data_shred.last_in_slot());
        let mut packet = Packet::default();
        last_data_shred.copy_to_packet(&mut packet);

        let bad_last_index = 30u32;
        let fec_set_index = 0u32;
        {
            let mut cursor = Cursor::new(packet.buffer_mut());
            cursor
                .seek(SeekFrom::Start(OFFSET_OF_SHRED_INDEX as u64))
                .unwrap();
            cursor.write_all(&bad_last_index.to_le_bytes()).unwrap();
            cursor
                .seek(SeekFrom::Start(OFFSET_OF_FEC_SET_INDEX as u64))
                .unwrap();
            cursor.write_all(&fec_set_index.to_le_bytes()).unwrap();
        }

        let mut shred_filter_context = ShredFilterContext::new(root_bank.clone(), shred_version);
        assert!(shred_filter_context.should_discard_packet(&packet));
        assert_eq!(shred_filter_context.stats.misaligned_last_data_index, 1);
    }

    #[test]
    fn test_should_discard_shred_with_custom_shred_limits() {
        agave_logger::setup();
        let mut rng = rand::rng();
        let slot = 200;
        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            slot,
            1200 * 5, // data_size
            false,    // is_last_in_slot
        )
        .unwrap();
        let shreds: Vec<_> = shreds.into_iter().map(Shred::from).collect();
        let shred_version = shreds[0].common_header().version;
        let root_bank = new_test_bank(0, false);

        for shred_type in [ShredType::Data, ShredType::Code] {
            let shred = shreds
                .iter()
                .find(|shred| shred.shred_type() == shred_type)
                .unwrap();
            let index = shred.index();
            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            match shred_type {
                ShredType::Data => shred_filter_context.set_shred_limits_for_tests(
                    ShredLimits::new(index + 1, ShredLimits::DEFAULT.max_code_shreds_per_slot),
                ),
                ShredType::Code => shred_filter_context.set_shred_limits_for_tests(
                    ShredLimits::new(ShredLimits::DEFAULT.max_data_shreds_per_slot, index + 1),
                ),
            }
            assert!(
                !shred_filter_context.should_discard_shred(shred.payload()),
                "{shred_type:?} shred should be within the custom limit"
            );

            let mut shred_filter_context =
                ShredFilterContext::new(root_bank.clone(), shred_version);
            match shred_type {
                ShredType::Data => shred_filter_context.set_shred_limits_for_tests(
                    ShredLimits::new(index, ShredLimits::DEFAULT.max_code_shreds_per_slot),
                ),
                ShredType::Code => shred_filter_context.set_shred_limits_for_tests(
                    ShredLimits::new(ShredLimits::DEFAULT.max_data_shreds_per_slot, index),
                ),
            }
            assert!(
                shred_filter_context.should_discard_shred(shred.payload()),
                "{shred_type:?} shred should be discarded at the custom limit"
            );
            assert_eq!(shred_filter_context.stats.index_out_of_bounds, 1);
        }
    }

    #[test]
    fn test_shred_limit_for_slot_times() {
        for (features, shred_limits) in std::iter::once((vec![], ShredLimits::DEFAULT)).chain(
            slot_time_feature_gates().map(|(feature_id, params)| {
                (
                    vec![feature_id],
                    ShredLimits::new(
                        params.max_data_shreds_per_slot(),
                        params.max_code_shreds_per_slot(),
                    ),
                )
            }),
        ) {
            let (shred_filter, effective_slot) = shred_filter_for_tests(features);
            assert_eq!(
                shred_filter.shred_limits(effective_slot.saturating_sub(1)),
                ShredLimits::DEFAULT
            );
            assert_eq!(shred_filter.shred_limits(effective_slot), shred_limits);
        }

        let [reduce_to_350ms, _, _, reduce_to_200ms] = slot_time_feature_ids();
        let (shred_filter, effective_slot) =
            shred_filter_for_tests([reduce_to_350ms, reduce_to_200ms]);
        assert_eq!(
            shred_filter.shred_limits(effective_slot),
            ShredLimits::new(16_384, 16_384)
        );
    }

    #[test]
    fn test_data_complete_shred_index_validation() {
        agave_logger::setup();
        let mut rng = rand::rng();
        let root_bank = new_test_bank(0, true);
        let slot = root_bank.get_slots_in_epoch(root_bank.epoch());
        let shreds = make_merkle_shreds_for_tests(
            &mut rng,
            slot,
            1200 * 5, // data_size
            false,    // is_last_in_slot
        )
        .unwrap();
        let shreds: Vec<_> = shreds.into_iter().map(Shred::from).collect();

        let data_shred = shreds
            .iter()
            .find(|s| s.shred_type() == ShredType::Data)
            .unwrap();

        let shred_version = data_shred.common_header().version;

        let mut packet = Packet::default();
        data_shred.copy_to_packet(&mut packet);

        let fec_set_index = 64u32;
        let wrong_index = fec_set_index + 10;

        {
            let mut cursor = Cursor::new(packet.buffer_mut());
            cursor
                .seek(SeekFrom::Start(OFFSET_OF_SHRED_INDEX as u64))
                .unwrap();
            cursor.write_all(&wrong_index.to_le_bytes()).unwrap();
            cursor
                .seek(SeekFrom::Start(OFFSET_OF_FEC_SET_INDEX as u64))
                .unwrap();
            cursor.write_all(&fec_set_index.to_le_bytes()).unwrap();
            cursor
                .seek(SeekFrom::Start(OFFSET_OF_SHRED_FLAGS as u64))
                .unwrap();
            cursor
                .write_all(&[ShredFlags::DATA_COMPLETE_SHRED.bits()])
                .unwrap();
        }

        let mut shred_filter_context = ShredFilterContext::new(root_bank.clone(), shred_version);
        assert!(shred_filter_context.should_discard_packet(&packet));
        assert_eq!(shred_filter_context.stats.unexpected_data_complete_shred, 1);

        let correct_index = fec_set_index + DATA_SHREDS_PER_FEC_BLOCK as u32 - 1;
        {
            let mut cursor = Cursor::new(packet.buffer_mut());
            cursor
                .seek(SeekFrom::Start(OFFSET_OF_SHRED_INDEX as u64))
                .unwrap();
            cursor.write_all(&correct_index.to_le_bytes()).unwrap();
        }

        let mut shred_filter_context = ShredFilterContext::new(root_bank, shred_version);
        assert!(!shred_filter_context.should_discard_packet(&packet));
        assert_eq!(shred_filter_context.stats.unexpected_data_complete_shred, 0);
    }
}