Add receive_unstaked_tokens to lst staker

cosmwasm/lst-staker/src/lib.rs

@@ -4,14 +4,17 @@ use std::{collections::BTreeMap, num::NonZeroU32};
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
     Addr, Binary, Coin, DecCoin, Decimal256, DelegatorReward, Deps, DepsMut, DistributionMsg, Env,
-    Int256, MessageInfo, Response, StakingMsg, StdError, Uint128, Uint256, to_json_binary,
+    Int256, MessageInfo, Response, StakingMsg, StdError, Uint128, Uint256, ensure, to_json_binary,
     wasm_execute,
 };
 use cw_account::ensure_local_admin_or_self;
 use cw_utils::{PaymentError, must_pay};
 use depolama::StorageExt;
 use frissitheto::{InitStateVersionError, UpgradeError, UpgradeMsg};
-use lst::msg::{ConfigResponse, StakerExecuteMsg};
+use lst::{
+    msg::{Batch, ConfigResponse, StakerExecuteMsg},
+    types::BatchId,
+};
 
 use crate::{
     event::{Rebase, SetLstHubAddress, SetValidators, Stake, Unstake, ValidatorConfigured},
@@ -87,6 +90,11 @@ pub fn execute(
 
             rebase(deps.as_ref(), &env)
         }
+        ExecuteMsg::Staker(StakerExecuteMsg::ReceiveUnstakedTokens { batch_id }) => {
+            ensure_lst_hub(deps.as_ref(), &info)?;
+
+            receive_unstaked_tokens(deps.as_ref(), &env, batch_id)
+        }
     }
 }
 
@@ -419,6 +427,53 @@ pub fn migrate(
     )
 }
 
+// Call receive unstaked tokens to LstHub
+//
+// This only works if batch status is submitted and current time is over the receive time to make sure unbonding is complete already
+fn receive_unstaked_tokens(
+    deps: Deps,
+    env: &Env,
+    batch_id: BatchId,
+) -> Result<Response, ContractError> {
+    // query the lst hub to get batch detail
+    let lst_hub = deps.storage.read_item::<LstHub>()?.to_string();
+    let batch = deps
+        .querier
+        .query_wasm_smart::<Batch>(lst_hub.clone(), &lst::msg::QueryMsg::Batch { batch_id })?;
+
+    // get expected_native_unstaked from the submitted batch
+    let unstaked_amount = match batch {
+        Batch::Submitted(submitted_batch) => {
+            // make sure the current time is over the batch receive time
+            ensure!(
+                submitted_batch.receive_time <= env.block.time.seconds(),
+                ContractError::BatchNotReady {
+                    now: env.block.time.seconds(),
+                    ready_at: submitted_batch.receive_time,
+                }
+            );
+            submitted_batch.expected_native_unstaked
+        }
+        Batch::Pending(_) => {
+            return Err(ContractError::BatchStillPending { batch_id });
+        }
+        Batch::Received(_) => {
+            return Err(ContractError::BatchAlreadyReceived { batch_id });
+        }
+    };
+
+    Ok(Response::new()
+        // call receive unstaked tokens with the unstaked tokens to lst hub
+        .add_message(wasm_execute(
+            lst_hub,
+            &lst::msg::ExecuteMsg::ReceiveUnstakedTokens { batch_id },
+            vec![Coin {
+                denom: query_native_token_denom(deps)?,
+                amount: Uint128::new(unstaked_amount),
+            }],
+        )?))
+}
+
 #[derive(Debug, PartialEq, thiserror::Error)]
 pub enum ContractError {
     #[error(transparent)]
@@ -447,4 +502,16 @@ pub enum ContractError {
 
     #[error("sender {sender} is not the lst hub")]
     OnlyLstHub { sender: Addr },
+
+    #[error("sender {sender} is not the lst hub")]
+    InvalidBatch { sender: Addr },
+
+    #[error("batch {batch_id} is still pending")]
+    BatchStillPending { batch_id: BatchId },
+
+    #[error("batch {batch_id} has already been received")]
+    BatchAlreadyReceived { batch_id: BatchId },
+
+    #[error("batch is not ready to be submitted/received (now={now}, ready_at={ready_at})")]
+    BatchNotReady { now: u64, ready_at: u64 },
 }

cosmwasm/lst-staker/src/tests.rs

@@ -11,7 +11,10 @@ use cw_account::{
     types::{Admin, LocalAdmin},
 };
 use depolama::StorageExt;
-use lst::{msg::ConfigResponse, types::ProtocolFeeConfig};
+use lst::{
+    msg::ConfigResponse,
+    types::{BatchId, ProtocolFeeConfig},
+};
 
 use crate::{
     ContractError, execute, msg::ExecuteMsg, redisribute_delegations, withdraw_all_rewards,
@@ -162,6 +165,21 @@ fn lst_ops_require_lst() {
             sender: non_admin.clone(),
         },
     );
+
+    assert_eq!(
+        execute(
+            deps.as_mut(),
+            mock_env(),
+            message_info(&non_admin, &[]),
+            ExecuteMsg::Staker(lst::msg::StakerExecuteMsg::ReceiveUnstakedTokens {
+                batch_id: BatchId::ONE
+            }),
+        )
+        .unwrap_err(),
+        ContractError::OnlyLstHub {
+            sender: non_admin.clone(),
+        },
+    );
 }
 
 #[test]

cosmwasm/lst/src/contract.rs

@@ -69,9 +69,9 @@ use crate::{
     error::ContractError,
     event::Init,
     execute::{
-        FEE_RATE_DENOMINATOR, accept_ownership, bond, circuit_breaker, rebase, receive_rewards,
-        receive_unstaked_tokens, resume_contract, revoke_ownership_transfer, slash_batches,
-        submit_batch, transfer_ownership, unbond, update_config, withdraw,
+        FEE_RATE_DENOMINATOR, accept_ownership, bond, circuit_breaker, rebase, receive_batch,
+        receive_rewards, receive_unstaked_tokens, resume_contract, revoke_ownership_transfer,
+        slash_batches, submit_batch, transfer_ownership, unbond, update_config, withdraw,
     },
     msg::{ExecuteMsg, InitMsg, QueryMsg},
     query::{
@@ -226,6 +226,7 @@ pub fn execute(
             },
         ),
         ExecuteMsg::SlashBatches { new_amounts } => slash_batches(deps, info, new_amounts),
+        ExecuteMsg::ReceiveBatch { batch_id } => receive_batch(deps, env, info, batch_id),
     }
 }
 

cosmwasm/lst/src/event.rs

@@ -170,3 +170,10 @@ pub struct SlashBatch {
     pub batch_id: BatchId,
     pub amount: u128,
 }
+
+#[derive(Event)]
+#[event("receive_batch")]
+pub struct ReceiveBatch {
+    pub batch_id: BatchId,
+    pub amount: u128,
+}

cosmwasm/lst/src/execute.rs

@@ -69,9 +69,9 @@ use depolama::StorageExt;
 use crate::{
     error::{ContractError, ContractResult},
     event::{
-        AcceptOwnership, Bond, CircuitBreaker, Rebase, ReceiveRewards, ReceiveUnstakedTokens,
-        ResumeContract, RevokeOwnershipTransfer, SlashBatch, SubmitBatch, TransferOwnership,
-        Unbond, Withdraw,
+        AcceptOwnership, Bond, CircuitBreaker, Rebase, ReceiveBatch, ReceiveRewards,
+        ReceiveUnstakedTokens, ResumeContract, RevokeOwnershipTransfer, SlashBatch, SubmitBatch,
+        TransferOwnership, Unbond, Withdraw,
     },
     helpers::{assets_to_shares, shares_to_assets, total_assets},
     msg::StakerExecuteMsg,
@@ -858,3 +858,45 @@ pub fn slash_batches(
          }| SlashBatch { batch_id, amount },
     )))
 }
+
+/// Receive batch to receive unstaked tokens from Staker
+/// This is permissionless as it only can be used to transfer received unstaked tokens of batch from staker to lst hub
+pub fn receive_batch(
+    deps: DepsMut,
+    env: Env,
+    _info: MessageInfo,
+    batch_id: BatchId,
+) -> ContractResult<Response> {
+    ensure_not_stopped(deps.as_ref())?;
+
+    let SubmittedBatch {
+        total_lst_to_burn: _,
+        unstake_requests_count: _,
+        receive_time,
+        expected_native_unstaked,
+    } = deps
+        .storage
+        .read::<SubmittedBatches>(&batch_id)
+        .map_err(|_| ContractError::BatchNotFound { batch_id })?;
+
+    ensure!(
+        receive_time <= env.block.time.seconds(),
+        ContractError::BatchNotReady {
+            now: env.block.time.seconds(),
+            ready_at: receive_time,
+        }
+    );
+
+    let response = Response::new()
+        .add_message(wasm_execute(
+            deps.storage.read_item::<StakerAddress>()?.to_string(),
+            &StakerExecuteMsg::ReceiveUnstakedTokens { batch_id },
+            vec![],
+        )?)
+        .add_event(ReceiveBatch {
+            batch_id,
+            amount: expected_native_unstaked,
+        });
+
+    Ok(response)
+}

cosmwasm/lst/src/msg.rs

@@ -188,6 +188,9 @@ pub enum ExecuteMsg {
     SlashBatches {
         new_amounts: Vec<BatchExpectedAmount>,
     },
+
+    /// Call Staker to received unstaked tokens for specific batch
+    ReceiveBatch { batch_id: BatchId },
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -338,4 +341,8 @@ pub enum StakerExecuteMsg {
     ///
     /// This must only be callable by the LST hub itself.
     Rebase {},
+    /// Receive unstaked tokens to mark batch as received
+    ///
+    /// This must only be callable by the LST hub itself.
+    ReceiveUnstakedTokens { batch_id: BatchId },
 }

cosmwasm/lst/src/tests/mod.rs

@@ -78,6 +78,7 @@ mod instantiate_tests;
 mod ownership_tests;
 mod query_tests;
 mod rebase_tests;
+mod receive_batch_tests;
 mod reward_tests;
 mod submit_batch_tests;
 mod test_helper;