From dba8a4ee4ea2c86ed7ccfb03cab9a097b1a1b20b Mon Sep 17 00:00:00 2001 From: kingsleydon <10992364+kingsleydon@users.noreply.github.com> Date: Wed, 11 Dec 2024 18:57:32 +0800 Subject: [PATCH] ensure invest pools when vault withdraw from stake pools --- pallets/phala/src/compute/stake_pool_v2.rs | 29 ++++++++++++++++------ 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/pallets/phala/src/compute/stake_pool_v2.rs b/pallets/phala/src/compute/stake_pool_v2.rs index f0e4fbbae..5d2cbc1ab 100644 --- a/pallets/phala/src/compute/stake_pool_v2.rs +++ b/pallets/phala/src/compute/stake_pool_v2.rs @@ -780,13 +780,6 @@ pub mod pallet { who.clone(), pool_info.basepool.pid, )?; - let nft_id = maybe_nft_id.ok_or(Error::::NoNftToWithdraw)?; - // The nft instance must be wrote to Nft storage at the end of the function - // this nft's property shouldn't be accessed or wrote again from storage before set_nft_attr - // is called. Or the property of the nft will be overwrote incorrectly. - let mut nft_guard = - base_pool::Pallet::::get_nft_attr_guard(pool_info.basepool.cid, nft_id)?; - let nft = &mut nft_guard.attr; let in_queue_shares = match pool_info .basepool .withdraw_queue @@ -803,10 +796,32 @@ pub mod pallet { } None => Zero::zero(), }; + ensure!(maybe_nft_id.is_some() || in_queue_shares > Zero::zero(), Error::::NoNftToWithdraw); + let nft_id = match maybe_nft_id { + Some(nft_id) => nft_id, + None => base_pool::Pallet::::mint_nft( + pool_info.basepool.cid, + who.clone(), + Zero::zero(), + pool_info.basepool.pid, + )?, + }; + // The nft instance must be wrote to Nft storage at the end of the function + // this nft's property shouldn't be accessed or wrote again from storage before set_nft_attr + // is called. Or the property of the nft will be overwrote incorrectly. + let mut nft_guard = + base_pool::Pallet::::get_nft_attr_guard(pool_info.basepool.cid, nft_id)?; + let nft = &mut nft_guard.attr; ensure!( base_pool::is_nondust_balance(shares) && (shares <= nft.shares + in_queue_shares), Error::::InvalidWithdrawalAmount ); + if let Some(vault_pid) = as_vault { + let vault_info = ensure_vault::(vault_pid)?; + if !vault_info.invest_pools.contains(&pid) { + vault_info.invest_pools.push(pid); + } + } base_pool::Pallet::::try_withdraw( &mut pool_info.basepool, nft,