staking: handle param change and unbonding edge case (#4056)

This fix a bug in the claim logic that would cause the node to crash if
a parameter change had occurred, or if an undelegation was marked to
have taken place after the unbonding height specified in `Unbonding {
unbonds_at_height }`. The latter can happen if a validator gets ejected
from the consensus set for inactivity, and is equivalent to an
`Unbonded` state.

crates/bin/pcli/src/command/tx.rs

@@ -57,7 +57,7 @@ use penumbra_shielded_pool::Ics20Withdrawal;
 use penumbra_stake::rate::RateData;
 use penumbra_stake::{DelegationToken, IdentityKey, Penalty, UnbondingToken, UndelegateClaimPlan};
 use penumbra_transaction::{gas::swap_claim_gas_cost, memo::MemoPlaintext};
-use penumbra_view::ViewClient;
+use penumbra_view::{SpendableNoteRecord, ViewClient};
 use penumbra_wallet::plan::{self, Planner};
 use proposal::ProposalCmd;
 
@@ -660,20 +660,36 @@ impl TxCmd {
                 // We want to claim them into the same address index that currently holds the tokens.
                 let notes = view.unspent_notes_by_address_and_asset().await?;
 
+                let notes: Vec<(
+                    AddressIndex,
+                    Vec<(UnbondingToken, Vec<SpendableNoteRecord>)>,
+                )> = notes
+                    .into_iter()
+                    .map(|(address_index, notes_by_asset)| {
+                        let mut filtered_notes: Vec<(UnbondingToken, Vec<SpendableNoteRecord>)> =
+                            notes_by_asset
+                                .into_iter()
+                                .filter_map(|(asset_id, notes)| {
+                                    // Filter for notes that are unbonding tokens.
+                                    let denom = asset_cache
+                                        .get(&asset_id)
+                                        .expect("asset ID should exist in asset cache")
+                                        .clone();
+                                    match UnbondingToken::try_from(denom) {
+                                        Ok(token) => Some((token, notes)),
+                                        Err(_) => None,
+                                    }
+                                })
+                                .collect();
+
+                        filtered_notes.sort_by_key(|(token, _)| token.unbonding_start_height());
+
+                        (address_index, filtered_notes)
+                    })
+                    .collect();
+
                 for (address_index, notes_by_asset) in notes.into_iter() {
-                    for (token, notes) in
-                        notes_by_asset.into_iter().filter_map(|(asset_id, notes)| {
-                            // Filter for notes that are unbonding tokens.
-                            let denom = asset_cache
-                                .get(&asset_id)
-                                .expect("asset ID should exist in asset cache")
-                                .clone();
-                            match UnbondingToken::try_from(denom) {
-                                Ok(token) => Some((token, notes)),
-                                Err(_) => None,
-                            }
-                        })
-                    {
+                    for (token, notes) in notes_by_asset.into_iter() {
                         println!("claiming {}", token.denom().default_unit());
 
                         let validator_identity = token.validator();
@@ -685,7 +701,8 @@ impl TxCmd {
                             .epoch_by_height(EpochByHeightRequest {
                                 height: unbonding_start_height,
                             })
-                            .await?
+                            .await
+                            .expect("can get epoch by height")
                             .into_inner()
                             .epoch
                             .context("unable to get epoch for unbonding start height")?;

crates/core/component/stake/src/component/action_handler/undelegate_claim.rs

@@ -37,17 +37,17 @@ impl ActionHandler for UndelegateClaim {
         // if profiling shows that they cause a bottleneck we could (CAREFULLY)
         // move some of them back.
 
-        // If the validator delegation pool is bonded, or unbonding, we must
-        // check that the unbonding delay (measured in blocks) has elapsed.
         let current_height = state.get_block_height().await?;
+        let unbonding_start_height = self.body.unbonding_start_height;
+        ensure!(
+            current_height >= unbonding_start_height,
+            "the unbonding start height must be less than or equal to the current height"
+        );
 
-        // Check if the unbonding height has been reached, it will always be the case if
-        // the validator delegation pool is unbonded.
+        // Compute the unbonding height for the claim, and check that it is less than or equal to the current height.
+        // If the pool is `Unbonded` or unbonding at an already elapsed height, we default to the current height.
         let allowed_unbonding_height = state
-            .compute_unbonding_height(
-                &self.body.validator_identity,
-                self.body.unbonding_start_height,
-            )
+            .compute_unbonding_height(&self.body.validator_identity, unbonding_start_height)
             .await?
             .unwrap_or(current_height);
 
@@ -66,6 +66,13 @@ impl ActionHandler for UndelegateClaim {
             .await?;
         let unbonding_epoch_end = state.get_epoch_by_height(allowed_unbonding_height).await?;
 
+        // This should never happen, but if it did we want to make sure that it wouldn't
+        // crash the mempool.
+        ensure!(
+            unbonding_epoch_end.index >= unbonding_epoch_start.index,
+            "unbonding epoch end must be greater than or equal to unbonding epoch start"
+        );
+
         // Compute the penalty for the epoch range [unbonding_epoch_start, unbonding_epoch_end], and check
         // that it matches the penalty in the claim.
         let expected_penalty = state
@@ -78,7 +85,9 @@ impl ActionHandler for UndelegateClaim {
 
         ensure!(
             self.body.penalty == expected_penalty,
-            "penalty does not match expected penalty"
+            "penalty (kept_rate: {}) does not match expected penalty (kept_rate: {})",
+            self.body.penalty.kept_rate(),
+            expected_penalty.kept_rate(),
         );
 
         /* ---------- execution ----------- */

crates/core/component/stake/src/component/stake.rs

@@ -388,6 +388,9 @@ pub trait SlashingData: StateRead {
         epoch_index_start: u64,
         epoch_index_end: u64,
     ) -> Result<Penalty> {
+        if epoch_index_start > epoch_index_end {
+            anyhow::bail!("invalid penalty window")
+        }
         let range = self
             .get_penalty_for_range(id, epoch_index_start, epoch_index_end)
             .await;

crates/core/component/stake/src/component/validator_handler/validator_store.rs

@@ -164,14 +164,30 @@ pub trait ValidatorDataRead: StateRead {
         };
 
         let min_block_delay = self.get_stake_params().await?.unbonding_delay;
-
         let upper_bound_height = start_height.saturating_add(min_block_delay);
 
         let unbonding_height = match val_bonding_state {
+            // The pool is bonded, so the unbonding height is the start height plus the delay.
             Bonded => Some(upper_bound_height),
-            // When the minimum delay parameter changes, an unbonding validator may
-            // have a delay that is larger than the new minimum delay. In this case,
-            Unbonding { unbonds_at_height } => Some(unbonds_at_height.min(upper_bound_height)),
+            // The pool is unbonding at a specific height, so we can use that.
+            Unbonding { unbonds_at_height } => {
+                if unbonds_at_height > start_height {
+                    // The unbonding height is the minimum of the unbonding height and the upper bound.
+                    // There are a couple reasons:
+                    // - The unbonding delay parameter can change, and in particular, it can decrease.
+                    // - We might be processing an undelegation that was initiated before the validator
+                    //   began unbonding, and the unbonding height is in the past.
+                    Some(unbonds_at_height.min(upper_bound_height))
+                } else {
+                    // In some cases, the allowed unbonding height can be smaller than
+                    // undelgation start height, for example if the unbonding delay has
+                    // changed in a parameter update, or if the unbonding has finished
+                    // and the validator is not indexed by the staking module anymore.
+                    // This is functionally equivalent to dealing with an `Unbonded` pool.
+                    None
+                }
+            }
+            // The pool is unbonded, so the unbonding height can be decided by the caller.
             Unbonded => None,
         };