Merge pull request #203 from osmosis-labs/trinity/on-channel-close

feat: on_close_channel logic

Author: vuong177

contracts/consumer/converter/src/contract.rs

@@ -75,7 +75,7 @@ impl ConverterContract<'_> {
         virtual_staking_code_id: u64,
         tombstoned_unbond_enable: bool,
         admin: Option<String>,
-        max_retrieve: u16,
+        max_retrieve: u32,
     ) -> Result<custom::Response, ContractError> {
         nonpayable(&ctx.info)?;
         // validate args

contracts/consumer/converter/src/error.rs

@@ -26,6 +26,12 @@ pub enum ContractError {
     #[error("You must start the channel handshake on this side, it doesn't support OpenTry")]
     IbcOpenTryDisallowed,
 
+    #[error("IBC channels not match")]
+    IbcChannelNotMatch,
+
+    #[error("You must start the channel close on provider side")]
+    IbcChannelCloseInitDisallowed,
+
     #[error("Sent wrong denom over IBC: {sent}, expected {expected}")]
     WrongDenom { sent: String, expected: String },
 

contracts/consumer/converter/src/ibc.rs

@@ -188,11 +188,18 @@ pub(crate) fn valset_update_msg(
 /// On closed channel, we take all tokens from reflect contract to this contract.
 /// We also delete the channel entry from accounts.
 pub fn ibc_channel_close(
-    _deps: DepsMut,
+    deps: DepsMut,
     _env: Env,
     _msg: IbcChannelCloseMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
-    todo!();
+    let contract = ConverterContract::new();
+    let msg = virtual_staking_api::sv::ExecMsg::HandleCloseChannel {};
+    let msg = WasmMsg::Execute {
+        contract_addr: contract.virtual_stake.load(deps.storage)?.into(),
+        msg: to_json_binary(&msg)?,
+        funds: vec![],
+    };
+    Ok(IbcBasicResponse::new().add_message(msg))
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
@@ -228,7 +235,7 @@ pub fn ibc_packet_receive(
             tx_id: _,
         } => {
             let response = contract.unstake(deps, delegator, validator, unstake)?;
-            let ack = ack_success(&UnstakeAck {})?;
+            let ack: cosmwasm_std::Binary = ack_success(&UnstakeAck {})?;
             IbcReceiveResponse::new()
                 .set_ack(ack)
                 .add_submessages(response.messages)
@@ -247,9 +254,16 @@ pub fn ibc_packet_receive(
         ProviderPacket::TransferRewards {
             rewards, recipient, ..
         } => {
-            let msg = contract.transfer_rewards(deps.as_ref(), recipient, rewards)?;
+            let msg =
+                contract.transfer_rewards(deps.as_ref(), recipient.clone(), rewards.clone())?;
+            let event = Event::new("mesh-transfer-rewards")
+                .add_attribute("recipient", &recipient)
+                .add_attribute("rewards", &rewards.amount.to_string());
             let ack = ack_success(&TransferRewardsAck {})?;
-            IbcReceiveResponse::new().set_ack(ack).add_message(msg)
+            IbcReceiveResponse::new()
+                .set_ack(ack)
+                .add_message(msg)
+                .add_event(event)
         }
     };
     Ok(res)

contracts/consumer/virtual-staking/src/contract.rs

@@ -1,5 +1,6 @@
 use std::cmp::Ordering;
 use std::collections::{BTreeMap, HashMap, HashSet};
+use std::u32;
 
 use cosmwasm_std::{
     coin, ensure_eq, to_json_binary, Coin, CosmosMsg, CustomQuery, DepsMut, DistributionMsg, Env,
@@ -74,7 +75,7 @@ impl VirtualStakingContract<'_> {
     pub fn instantiate(
         &self,
         ctx: InstantiateCtx<VirtualStakeCustomQuery>,
-        max_retrieve: u16,
+        max_retrieve: u32,
         tombstoned_unbond_enable: bool,
     ) -> Result<Response<VirtualStakeCustomMsg>, ContractError> {
         nonpayable(&ctx.info)?;
@@ -622,6 +623,52 @@ impl VirtualStakingApi for VirtualStakingContract<'_> {
         Ok(Response::new().add_messages(msgs))
     }
 
+    fn handle_close_channel(
+        &self,
+        ctx: ExecCtx<Self::QueryC>,
+    ) -> Result<Response<VirtualStakeCustomMsg>, Self::Error> {
+        nonpayable(&ctx.info)?;
+        let ExecCtx { deps, env, info } = ctx;
+        let config = self.config.load(deps.storage)?;
+        ensure_eq!(info.sender, config.converter, ContractError::Unauthorized); // only the converter can call this
+
+        let all_delegations = TokenQuerier::new(&deps.querier)
+            .all_delegations(env.contract.address.to_string(), u32::MAX)?;
+
+        let mut msgs = vec![VirtualStakeMsg::DeleteAllScheduledTasks {}];
+        for delegation in all_delegations.delegations.iter() {
+            let amount = Coin {
+                denom: config.denom.clone(),
+                amount: delegation.amount,
+            };
+            msgs.push(VirtualStakeMsg::UpdateDelegation {
+                amount: amount.clone(),
+                is_deduct: true,
+                delegator: delegation.delegator.clone(),
+                validator: delegation.validator.clone(),
+            });
+            msgs.push(VirtualStakeMsg::Unbond {
+                amount,
+                validator: delegation.validator.clone(),
+            });
+            self.bond_requests
+                .save(deps.storage, &delegation.validator, &Uint128::zero())?;
+        }
+
+        let requests: Vec<(String, Uint128)> = self
+            .bond_requests
+            .range(
+                deps.as_ref().storage,
+                None,
+                None,
+                cosmwasm_std::Order::Ascending,
+            )
+            .collect::<Result<_, _>>()?;
+        self.bonded.save(deps.storage, &requests)?;
+
+        Ok(Response::new().add_messages(msgs))
+    }
+
     // FIXME: need to handle custom message types and queries
     /**
      * This is called once per epoch to withdraw all rewards and rebalance the bonded tokens.

contracts/consumer/virtual-staking/src/state.rs

@@ -9,9 +9,9 @@ pub struct Config {
     /// The address of the converter contract (that is authorized to bond/unbond and will receive rewards)
     pub converter: Addr,
 
+    /// Maximum delegations per query
+    pub max_retrieve: u32,
+
     /// If it enable, tombstoned validators will be unbond automatically
     pub tombstoned_unbond_enable: bool,
-
-    /// Maximum delegations per query
-    pub max_retrieve: u16,
 }

contracts/provider/external-staking/src/contract.rs

@@ -1,6 +1,6 @@
 use cosmwasm_std::{
-    coin, ensure, ensure_eq, to_json_binary, Coin, Decimal, DepsMut, Env, Event, IbcMsg, Order,
-    Response, StdResult, Storage, Uint128, Uint256, WasmMsg,
+    coin, ensure, ensure_eq, to_json_binary, Addr, Coin, Decimal, DepsMut, Env, Event, IbcMsg,
+    Order, Response, StdResult, Storage, Uint128, Uint256, WasmMsg,
 };
 use cw2::set_contract_version;
 use cw_storage_plus::{Bounder, Item, Map};
@@ -497,6 +497,55 @@ impl ExternalStakingContract<'_> {
         Ok(event)
     }
 
+    pub(crate) fn handle_close_channel(
+        &self,
+        deps: DepsMut,
+        env: Env,
+    ) -> Result<(), ContractError> {
+        let stakes: Vec<((Addr, String), Stake)> = self
+            .stakes
+            .stake
+            .range(
+                deps.as_ref().storage,
+                None,
+                None,
+                cosmwasm_std::Order::Ascending,
+            )
+            .collect::<Result<_, _>>()?;
+
+        for ((user, validator), stake) in stakes.iter() {
+            let mut new_stake = stake.clone();
+            let amount = new_stake.stake.low().clone();
+            let unbond = PendingUnbond {
+                amount,
+                release_at: env.block.time,
+            };
+            new_stake.pending_unbonds.push(unbond);
+            new_stake.stake = ValueRange::new_val(Uint128::zero());
+
+            let mut distribution = self
+                .distribution
+                .may_load(deps.storage, &validator)?
+                .unwrap_or_default();
+            new_stake
+                .points_alignment
+                .stake_decreased(amount, distribution.points_per_stake);
+
+            distribution.total_stake -= amount;
+
+            // Save stake
+            self.stakes
+                .stake
+                .save(deps.storage, (user, validator), &new_stake)?;
+
+            // Save distribution
+            self.distribution
+                .save(deps.storage, validator, &distribution)?;
+        }
+
+        Ok(())
+    }
+
     /// In non-test code, this is called from `ibc_packet_ack`
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn valset_update(

contracts/provider/external-staking/src/error.rs

@@ -45,6 +45,12 @@ pub enum ContractError {
     #[error("You must start the channel handshake on the other side, it doesn't support OpenInit")]
     IbcOpenInitDisallowed,
 
+    #[error("IBC channels not match")]
+    IbcChannelNotMatch,
+
+    #[error("You must start the channel close on this side")]
+    IbcChannelCloseConfirmDisallowed,
+
     #[error("Invalid authorized endpoint: {0}")]
     InvalidEndpoint(String),
 

contracts/provider/external-staking/src/ibc.rs

@@ -104,11 +104,25 @@ pub fn ibc_channel_connect(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_channel_close(
-    _deps: DepsMut,
-    _env: Env,
-    _msg: IbcChannelCloseMsg,
+    deps: DepsMut,
+    env: Env,
+    msg: IbcChannelCloseMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
-    todo!();
+    match msg {
+        IbcChannelCloseMsg::CloseInit { channel } => {
+            if channel.ne(&IBC_CHANNEL.load(deps.storage)?) {
+                return Err(ContractError::IbcChannelNotMatch);
+            }
+        }
+        IbcChannelCloseMsg::CloseConfirm { .. } => {
+            return Err(ContractError::IbcChannelCloseConfirmDisallowed)
+        }
+    };
+
+    let contract = ExternalStakingContract::new();
+    contract.handle_close_channel(deps, env)?;
+
+    Ok(IbcBasicResponse::new())
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]

packages/apis/src/virtual_staking_api.rs

@@ -65,6 +65,14 @@ pub trait VirtualStakingApi {
         amount: Coin,
     ) -> Result<Response<Self::ExecC>, Self::Error>;
 
+    /// Handle the close channel process.
+    /// Unbond all tokens from contract and delete scheduled tasks.
+    #[sv::msg(exec)]
+    fn handle_close_channel(
+        &self,
+        ctx: ExecCtx<Self::QueryC>,
+    ) -> Result<Response<Self::ExecC>, Self::Error>;
+
     /// SudoMsg::HandleEpoch{} should be called once per epoch by the sdk (in EndBlock).
     /// It allows the virtual staking contract to bond or unbond any pending requests, as well
     /// as to perform a rebalance if needed (over the max cap).

packages/bindings/src/query.rs

@@ -28,7 +28,7 @@ pub enum VirtualStakeQuery {
 
     /// Returns a max retrieve amount of delegations for the given contract
     #[returns(AllDelegationsResponse)]
-    AllDelegations { contract: String, max_retrieve: u16 },
+    AllDelegations { contract: String, max_retrieve: u32 },
 }
 
 /// Bookkeeping info in the virtual staking sdk module
@@ -108,7 +108,7 @@ impl<'a> TokenQuerier<'a> {
     pub fn all_delegations(
         &self,
         contract: String,
-        max_retrieve: u16,
+        max_retrieve: u32,
     ) -> StdResult<AllDelegationsResponse> {
         let all_delegations_query = VirtualStakeQuery::AllDelegations {
             contract,