diff --git a/pallets/phala/src/compute/stake_pool_v2.rs b/pallets/phala/src/compute/stake_pool_v2.rs index f0e4fbbae..b1dc4fd2e 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,33 @@ 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, + // An nft is necessary to initiate a smaller withdrawal + 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 mut 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, diff --git a/pallets/phala/src/compute/vault.rs b/pallets/phala/src/compute/vault.rs index bf0003af4..a8fca5cca 100644 --- a/pallets/phala/src/compute/vault.rs +++ b/pallets/phala/src/compute/vault.rs @@ -429,7 +429,7 @@ pub mod pallet { if base_pool::is_nondust_balance(property.shares) { true } else { - if (!base_pool::is_nondust_balance(total_shares)) { + if !base_pool::is_nondust_balance(total_shares) { let _ = base_pool::Pallet::::burn_nft( &base_pool::pallet_id::(), stake_pool.basepool.cid, diff --git a/pallets/phala/src/snapshots/phala_pallets__test__vault_force_withdraw_with_dust_investment.snap b/pallets/phala/src/snapshots/phala_pallets__test__vault_force_withdraw_with_dust_investment.snap index 254c0c8cb..8aa66072a 100644 --- a/pallets/phala/src/snapshots/phala_pallets__test__vault_force_withdraw_with_dust_investment.snap +++ b/pallets/phala/src/snapshots/phala_pallets__test__vault_force_withdraw_with_dust_investment.snap @@ -3,51 +3,6 @@ source: pallets/phala/src/test.rs expression: take_events() --- [ - RuntimeEvent::RmrkCore( - Event::PropertiesRemoved { - collection_id: 10000, - maybe_nft_id: Some( - 0, - ), - }, - ), - RuntimeEvent::Uniques( - Event::Burned { - collection: 10000, - item: 0, - owner: 13009150994509951074, - }, - ), - RuntimeEvent::RmrkCore( - Event::NFTBurned { - owner: 7813586407040180578, - collection_id: 10000, - nft_id: 0, - }, - ), - RuntimeEvent::RmrkCore( - Event::PropertyRemoved { - collection_id: 10000, - maybe_nft_id: Some( - 0, - ), - key: BoundedVec( - [ - 99, - 114, - 101, - 97, - 116, - 101, - 116, - 105, - 109, - 101, - ], - 32000, - ), - }, - ), RuntimeEvent::PhalaVault( Event::ForceShutdown { pid: 1, diff --git a/pallets/phala/src/test.rs b/pallets/phala/src/test.rs index 394b564f4..e1e70492c 100644 --- a/pallets/phala/src/test.rs +++ b/pallets/phala/src/test.rs @@ -2363,10 +2363,6 @@ fn vault_force_withdraw_with_dust_investment() { // Verify the vault is locked assert!(vault::VaultLocks::::contains_key(vault1)); - - // After force withdrawal, the invest_pools should be empty since the only pool had no shares - let vault = ensure_vault::(vault1).unwrap(); - assert!(vault.invest_pools.is_empty(), "invest_pools should be cleaned up"); }); }