PoS: allow to set a receiver target for claim reward tx

Author: tzemanovic

crates/node/src/shell/finalize_block.rs

@@ -2444,15 +2444,18 @@ mod test_finalize_block {
             .unwrap();
 
         // Claim the rewards from the initial epoch
-        let reward_1 =
-            proof_of_stake::claim_reward_tokens::<
-                _,
-                governance::Store<_>,
-                token::Store<_>,
-            >(
-                &mut shell.state, None, &validator.address, current_epoch
-            )
-            .unwrap();
+        let reward_1 = proof_of_stake::claim_reward_tokens::<
+            _,
+            governance::Store<_>,
+            token::Store<_>,
+        >(
+            &mut shell.state,
+            None,
+            &validator.address,
+            None,
+            current_epoch,
+        )
+        .unwrap();
         total_claimed += reward_1;
         assert_eq!(reward_1, query_rewards);
         assert!(is_reward_equal_enough(total_rewards, total_claimed, 1));
@@ -2481,7 +2484,11 @@ mod test_finalize_block {
             governance::Store<_>,
             token::Store<_>,
         >(
-            &mut shell.state, None, &validator.address, current_epoch
+            &mut shell.state,
+            None,
+            &validator.address,
+            None,
+            current_epoch,
         )
         .unwrap();
         assert_eq!(att, token::Amount::zero());
@@ -2520,7 +2527,11 @@ mod test_finalize_block {
             governance::Store<_>,
             token::Store<_>,
         >(
-            &mut shell.state, None, &validator.address, current_epoch
+            &mut shell.state,
+            None,
+            &validator.address,
+            None,
+            current_epoch,
         )
         .unwrap();
         total_claimed += rew;
@@ -2601,15 +2612,18 @@ mod test_finalize_block {
             .unwrap();
 
         // Claim tokens
-        let reward_2 =
-            proof_of_stake::claim_reward_tokens::<
-                _,
-                governance::Store<_>,
-                token::Store<_>,
-            >(
-                &mut shell.state, None, &validator.address, current_epoch
-            )
-            .unwrap();
+        let reward_2 = proof_of_stake::claim_reward_tokens::<
+            _,
+            governance::Store<_>,
+            token::Store<_>,
+        >(
+            &mut shell.state,
+            None,
+            &validator.address,
+            None,
+            current_epoch,
+        )
+        .unwrap();
         total_claimed += reward_2;
         assert_eq!(query_rewards, reward_2);
 
@@ -2730,15 +2744,18 @@ mod test_finalize_block {
         }
 
         // Claim the rewards for the validator for the first two epochs
-        let val_reward_1 =
-            proof_of_stake::claim_reward_tokens::<
-                _,
-                governance::Store<_>,
-                token::Store<_>,
-            >(
-                &mut shell.state, None, &validator.address, current_epoch
-            )
-            .unwrap();
+        let val_reward_1 = proof_of_stake::claim_reward_tokens::<
+            _,
+            governance::Store<_>,
+            token::Store<_>,
+        >(
+            &mut shell.state,
+            None,
+            &validator.address,
+            None,
+            current_epoch,
+        )
+        .unwrap();
         total_claimed += val_reward_1;
         assert!(is_reward_equal_enough(
             total_rewards,
@@ -2758,15 +2775,18 @@ mod test_finalize_block {
         total_rewards += inflation_3;
 
         // Claim again for the validator
-        let val_reward_2 =
-            proof_of_stake::claim_reward_tokens::<
-                _,
-                governance::Store<_>,
-                token::Store<_>,
-            >(
-                &mut shell.state, None, &validator.address, current_epoch
-            )
-            .unwrap();
+        let val_reward_2 = proof_of_stake::claim_reward_tokens::<
+            _,
+            governance::Store<_>,
+            token::Store<_>,
+        >(
+            &mut shell.state,
+            None,
+            &validator.address,
+            None,
+            current_epoch,
+        )
+        .unwrap();
 
         // Claim for the delegator
         let del_reward_1 = proof_of_stake::claim_reward_tokens::<
@@ -2777,6 +2797,7 @@ mod test_finalize_block {
             &mut shell.state,
             Some(&delegator),
             &validator.address,
+            None,
             current_epoch,
         )
         .unwrap();

crates/proof_of_stake/src/lib.rs

@@ -2805,6 +2805,7 @@ pub fn claim_reward_tokens<S, Gov, Token>(
     storage: &mut S,
     source: Option<&Address>,
     validator: &Address,
+    receiver: Option<&Address>,
     current_epoch: Epoch,
 ) -> Result<token::Amount>
 where
@@ -2814,8 +2815,14 @@ where
 {
     tracing::debug!("Claiming rewards in epoch {current_epoch}");
 
-    let source = source.cloned().unwrap_or_else(|| validator.clone());
-    tracing::debug!("Source {} --> Validator {}", source, validator);
+    let source = source.unwrap_or(validator).clone();
+    let receiver = receiver.unwrap_or(&source).clone();
+    tracing::debug!(
+        "Source {} --> Validator {}, Receiver {}",
+        source,
+        validator,
+        receiver
+    );
 
     let mut reward_tokens = compute_current_rewards_from_bonds::<S, Gov>(
         storage,
@@ -2832,17 +2839,23 @@ where
     // Update the last claim epoch in storage
     write_last_reward_claim_epoch(storage, &source, validator, current_epoch)?;
 
-    // Transfer the bonded tokens from PoS to the source
+    // Transfer the bonded tokens from PoS to the receiver
     let staking_token = staking_token_address(storage);
-    Token::transfer(storage, &staking_token, &ADDRESS, &source, reward_tokens)?;
+    Token::transfer(
+        storage,
+        &staking_token,
+        &ADDRESS,
+        &receiver,
+        reward_tokens,
+    )?;
     Token::emit_transfer_event(
         storage,
         CLAIM_REWARDS_EVENT_DESC.into(),
         trans_token::EventLevel::Tx,
         &staking_token,
         reward_tokens,
         trans_token::UserAccount::Internal(ADDRESS),
-        trans_token::UserAccount::Internal(source),
+        trans_token::UserAccount::Internal(receiver),
     )?;
 
     Ok(reward_tokens)

crates/sdk/src/queries/vp/pos.rs

@@ -1248,14 +1248,15 @@ mod test {
             .expect("Test failed");
         client.state.in_mem_mut().block.height = height + 1;
 
-        let claimed = namada_proof_of_stake::claim_reward_tokens::<
-            _,
-            governance::Store<_>,
-            namada_token::Store<_>,
-        >(
-            &mut client.state, Some(&delegator), &validator, epoch
-        )
-        .expect("Claiming rewards failed");
+        let claimed =
+            namada_proof_of_stake::claim_reward_tokens::<
+                _,
+                governance::Store<_>,
+                namada_token::Store<_>,
+            >(
+                &mut client.state, Some(&delegator), &validator, None, epoch
+            )
+            .expect("Claiming rewards failed");
 
         assert_eq!(claimed, del_reward_epoch_3 + del_reward_epoch_2);
 
@@ -1264,7 +1265,7 @@ mod test {
                 _,
                 governance::Store<_>,
                 namada_token::Store<_>,
-            >(&mut client.state, None, &validator, epoch)
+            >(&mut client.state, None, &validator, None, epoch)
             .expect("Claiming validator rewards failed");
 
         assert_eq!(claimed_validator, val_reward_epoch_3 + val_reward_epoch_2);

crates/sdk/src/tx.rs

@@ -1757,7 +1757,7 @@ pub async fn build_claim_rewards(
         None => Ok(source.clone()),
     }?;
 
-    let data = pos::ClaimRewards { validator, source };
+    let data = pos::ClaimRewardsCompat { validator, source };
 
     build(
         context,

crates/tx/src/action.rs

@@ -15,7 +15,7 @@ use namada_core::storage::KeySeg;
 use namada_core::{address, storage};
 
 pub use crate::data::pos::{
-    Bond, ClaimRewards, Redelegation, Unbond, Withdraw,
+    Bond, ClaimRewardsCompat as ClaimRewards, Redelegation, Unbond, Withdraw,
 };
 
 /// Actions applied from txs.

crates/tx/src/data/pos.rs

@@ -128,6 +128,23 @@ pub struct ClaimRewards {
     /// Source address for claiming rewards from a bond. For self-bonds, the
     /// validator is also the source
     pub source: Option<Address>,
+    /// Optional rewards receiver address
+    pub receiver: Option<Address>,
+}
+
+/// Compatibility data definition from previous version (before addition of a
+/// receiver field).
+///
+/// TODO: Temporarily replaces `crate::data::pos::ClaimRewards` that now
+/// contains an additional `receiver` field to maintain consensus compatibility.
+/// PoS VP reads this action and will need to be updated to use the new field.
+#[derive(Clone, Debug, BorshSerialize, BorshDeserialize, PartialEq)]
+pub struct ClaimRewardsCompat {
+    /// Validator address
+    pub validator: Address,
+    /// Source address for claiming rewards from a bond. For self-bonds, the
+    /// validator is also the source
+    pub source: Option<Address>,
 }
 
 /// A redelegation of bonded tokens from one validator to another.

crates/tx_prelude/src/proof_of_stake.rs

@@ -203,6 +203,7 @@ impl Ctx {
         &mut self,
         source: Option<&Address>,
         validator: &Address,
+        receiver: Option<&Address>,
     ) -> Result<token::Amount> {
         // The tx must be authorized by the source address
         let verifier = source.as_ref().unwrap_or(&validator);
@@ -218,6 +219,7 @@ impl Ctx {
             self,
             source,
             validator,
+            receiver,
             current_epoch,
         )
     }

wasm/tx_claim_rewards/src/lib.rs

@@ -1,15 +1,27 @@
 //! A tx for a user to claim PoS inflationary rewards due to bonds used as
 //! voting power in consensus.
 
+use namada_tx_prelude::transaction::pos;
 use namada_tx_prelude::*;
 
 #[transaction]
 fn apply_tx(ctx: &mut Ctx, tx_data: BatchedTx) -> TxResult {
     let data = ctx.get_tx_data(&tx_data)?;
-    let withdraw = transaction::pos::Withdraw::try_from_slice(&data[..])
-        .wrap_err("Failed to decode Withdraw value")?;
-
-    ctx.claim_reward_tokens(withdraw.source.as_ref(), &withdraw.validator)
+    let (validator, source, receiver) = if let Ok(pos::ClaimRewards {
+        validator,
+        source,
+        receiver,
+    }) =
+        pos::ClaimRewards::try_from_slice(&data[..])
+    {
+        (validator, source, receiver)
+    } else {
+        let pos::ClaimRewardsCompat { validator, source } =
+            pos::ClaimRewardsCompat::try_from_slice(&data[..])
+                .wrap_err("Failed to decode ClaimRewards value")?;
+        (validator, source, None)
+    };
+    ctx.claim_reward_tokens(source.as_ref(), &validator, receiver.as_ref())
         .wrap_err("Failed to claim rewards")?;
 
     Ok(())