Send rewards during claiming to a Consumer or Babylon address

Depends on ICS-20 transfer info being set

contracts/babylon/src/contract.rs

@@ -1,6 +1,6 @@
 use cosmwasm_std::{
-    to_json_binary, to_json_string, Addr, Binary, Deps, DepsMut, Empty, Env, IbcMsg, MessageInfo,
-    QueryResponse, Reply, Response, SubMsg, SubMsgResponse, WasmMsg,
+    to_json_binary, Addr, Binary, Deps, DepsMut, Empty, Env, IbcMsg, MessageInfo, QueryResponse,
+    Reply, Response, SubMsg, SubMsgResponse, WasmMsg,
 };
 use cw2::set_contract_version;
 use cw_utils::{must_pay, ParseReplyError};
@@ -290,30 +290,26 @@ pub fn execute(
             // TODO: Add events
             Ok(res)
         }
-        ExecuteMsg::SendRewards { fp_distribution } => {
+        ExecuteMsg::SendRewards { to_address } => {
             let cfg = CONFIG.load(deps.storage)?;
             // Assert the funds are there
             must_pay(&info, &cfg.denom)?;
             // Assert the sender is right
-            let btc_finality = cfg
-                .btc_finality
-                .ok_or(ContractError::BtcFinalityNotSet {})?;
-            if info.sender != btc_finality {
+            let btc_staking = cfg.btc_staking.ok_or(ContractError::BtcStakingNotSet {})?;
+            if info.sender != btc_staking {
                 return Err(ContractError::Unauthorized {});
             }
             // Route to babylon over IBC, if available
             let transfer_info = IBC_TRANSFER.may_load(deps.storage)?;
             match transfer_info {
                 Some(transfer_info) => {
-                    // Build the payload
-                    let payload_msg = to_json_string(&fp_distribution)?;
                     // Construct the transfer message
                     let ibc_msg = IbcMsg::Transfer {
                         channel_id: transfer_info.channel_id,
-                        to_address: transfer_info.to_address,
+                        to_address,
                         amount: info.funds[0].clone(),
                         timeout: packet_timeout(&env),
-                        memo: Some(payload_msg),
+                        memo: None,
                     };
 
                     // Send packet only if we are IBC enabled

contracts/babylon/src/msg/contract.rs

@@ -4,7 +4,6 @@ use cosmwasm_std::{Binary, StdError, StdResult};
 use babylon_apis::finality_api::Evidence;
 
 use crate::msg::btc_header::BtcHeader;
-use babylon_apis::btc_staking_api::RewardInfo;
 #[cfg(not(target_arch = "wasm32"))]
 use {
     crate::msg::btc_header::{BtcHeaderResponse, BtcHeadersResponse},
@@ -116,10 +115,11 @@ pub enum ExecuteMsg {
     /// This will be forwarded over IBC to the Babylon side for propagation to other Consumers, and
     /// Babylon itself
     Slashing { evidence: Evidence },
-    /// `SendRewards` is a message sent by the finality contract, to send rewards to Babylon
+    /// `SendRewards` is a message sent by the staking contract, to send rewards to the Babylon
+    /// chain
     SendRewards {
-        /// `fp_distribution` is the list of finality providers and their rewards
-        fp_distribution: Vec<RewardInfo>,
+        /// `to_address` is the address on the Babylon chain to send the rewards to
+        to_address: String,
     },
 }
 

contracts/btc-staking/src/staking.rs

@@ -3,12 +3,13 @@ use bitcoin::consensus::deserialize;
 use bitcoin::hashes::Hash;
 use bitcoin::{Transaction, Txid};
 use cosmwasm_std::{
-    coin, BankMsg, DepsMut, Env, Event, MessageInfo, Response, StdResult, Storage, Uint128, Uint256,
+    coin, to_json_binary, BankMsg, CanonicalAddr, CosmosMsg, DepsMut, Env, Event, MessageInfo,
+    Response, StdResult, Storage, Uint128, Uint256, WasmMsg,
 };
 use hex::ToHex;
 
 use crate::error::ContractError;
-use crate::state::config::{ADMIN, CONFIG, PARAMS};
+use crate::state::config::{Config, ADMIN, CONFIG, PARAMS};
 use crate::state::delegations::{delegations, DelegationDistribution};
 use crate::state::staking::{
     fps, BtcDelegation, DelegatorUnbondingInfo, FinalityProviderState, ACTIVATED_HEIGHT,
@@ -25,6 +26,7 @@ use babylon_apis::{to_canonical_addr, Validate};
 use babylon_bindings::BabylonMsg;
 use babylon_contract::msg::btc_header::BtcHeaderResponse;
 use babylon_contract::msg::contract::QueryMsg as BabylonQueryMsg;
+use babylon_contract::msg::ibc::TransferInfoResponse;
 use cosmwasm_std::Order::Ascending;
 use cw_utils::{must_pay, nonpayable};
 use std::str::FromStr;
@@ -467,25 +469,64 @@ pub fn handle_withdraw_rewards(
         return Err(ContractError::NoRewards);
     }
 
-    // Create the bank packet.
-    // Sends to the staker address on the Consumer.
-    // TODO: Send to the staker address on Babylon over IBC (ICS-020)
-    let recipient = deps.api.addr_humanize(&staker_canonical_addr)?;
-    let rewards_denom = cfg.denom;
-    let msg = BankMsg::Send {
-        to_address: recipient.to_string(),
-        amount: vec![coin(amount.u128(), rewards_denom)],
-    };
+    let (recipient, wasm_msg) =
+        send_rewards_msg(&deps, &staker_addr, &staker_canonical_addr, cfg, amount)?;
     let resp = Response::new()
-        .add_message(msg)
+        .add_message(wasm_msg)
         .add_attribute("action", "withdraw_rewards")
         .add_attribute("staker", staker_addr)
         .add_attribute("fp", fp_pubkey_hex)
-        .add_attribute("recipient", &recipient)
+        .add_attribute("recipient", recipient)
         .add_attribute("amount", amount.to_string());
     Ok(resp)
 }
 
+fn send_rewards_msg(
+    deps: &DepsMut,
+    staker_addr: &str,
+    staker_canonical_addr: &CanonicalAddr,
+    cfg: Config,
+    amount: Uint128,
+) -> Result<(String, CosmosMsg<BabylonMsg>), ContractError> {
+    // Query the babylon contract for transfer info
+    // TODO: Turn into a parameter set during instantiation to avoid query
+    let transfer_info: TransferInfoResponse = deps.querier.query_wasm_smart(
+        cfg.babylon.to_string(),
+        &babylon_contract::msg::contract::QueryMsg::TransferInfo {},
+    )?;
+
+    // Create the bank / routing packet
+    let (recipient, wasm_msg) = match transfer_info {
+        None => {
+            // Consumer withdrawal.
+            // Send rewards to the staker address on the Consumer
+            let recipient = deps.api.addr_humanize(staker_canonical_addr)?.to_string();
+            let msg = BankMsg::Send {
+                to_address: recipient.clone(),
+                amount: vec![coin(amount.u128(), cfg.denom)],
+            };
+            let wasm_msg = CosmosMsg::Bank(msg);
+            (recipient, wasm_msg)
+        }
+        Some(_) => {
+            // Babylon withdrawal.
+            // Send rewards to the staker address on Babylon, through the babylon contract (ICS-020
+            // transfer)
+            let recipient = staker_addr.to_string();
+            let msg = babylon_contract::msg::contract::ExecuteMsg::SendRewards {
+                to_address: recipient.clone(),
+            };
+            let wasm_msg = CosmosMsg::Wasm(WasmMsg::Execute {
+                contract_addr: cfg.babylon.to_string(),
+                msg: to_json_binary(&msg)?,
+                funds: vec![coin(amount.u128(), cfg.denom)],
+            });
+            (recipient, wasm_msg)
+        }
+    };
+    Ok((recipient, wasm_msg))
+}
+
 pub fn withdraw_delegation_reward(
     deps: DepsMut,
     delegation: &mut DelegationDistribution,