[Staking] Fixes and logic adjustments

pallets/ajuna-nft-staking/src/lib.rs

@@ -251,8 +251,6 @@ pub mod pallet {
 		Inactive,
 		/// The contract is staking hence cannot be claimed or sniped.
 		Staking,
-		/// The contract is expired hence cannot be claimed.
-		Expired,
 		/// The contract is claimable, so it cannot be cancelled or sniped.
 		Claimable,
 		/// The contract is available, or not yet accepted.
@@ -275,6 +273,8 @@ pub mod pallet {
 		NotSniper,
 		/// The given account attempts to snipe its own contract.
 		CannotSnipeOwnContract,
+		/// Cannot claim contract with unknown owner, only sniping is possible.
+		CannotClaimUnknownContract,
 	}
 
 	#[pallet::call]
@@ -362,7 +362,7 @@ pub mod pallet {
 			let _ = Self::ensure_creator(origin)?;
 			Self::ensure_pallet_unlocked()?;
 			Self::ensure_removable(&contract_id)?;
-			Self::remove_contract(contract_id)
+			Self::remove_non_staked_contract(contract_id)
 		}
 
 		/// Accept an available staking contract.
@@ -512,8 +512,9 @@ pub mod pallet {
 			Ok(())
 		}
 
-		fn remove_contract(contract_id: T::ItemId) -> DispatchResult {
+		fn remove_non_staked_contract(contract_id: T::ItemId) -> DispatchResult {
 			Contracts::<T>::remove(contract_id);
+			ContractsMetadata::<T>::remove(contract_id);
 			let collection_id = Self::contract_collection_id()?;
 			T::NftHelper::burn(&collection_id, &contract_id, None)?;
 			Self::deposit_event(Event::<T>::Removed { contract_id });
@@ -572,27 +573,38 @@ pub mod pallet {
 					T::NftHelper::transfer(&collection_id, &item_id, beneficiary),
 			}?;
 
-			// Return staked items.
-			let contract_owner = Self::contract_owner(&contract_id)?;
-			Self::transfer_items(&Self::staked_items(&contract_id)?, &contract_owner)?;
+			// Return staked items, if the owner is not present the staked items may go to a
+			// different party.
+			let stake_beneficiary = if let Ok(contract_owner) = Self::contract_owner(&contract_id) {
+				// Burn contract NFT.
+				let collection_id = Self::contract_collection_id()?;
+				T::NftHelper::burn(&collection_id, &contract_id, Some(&contract_owner))?;
+
+				// Retain contract IDs held.
+				let mut contract_ids = Self::contract_ids(&contract_owner)?;
+				contract_ids.retain(|id| id != &contract_id);
+				if contract_ids.is_empty() {
+					ContractIds::<T>::remove(&contract_owner);
+				} else {
+					ContractIds::<T>::insert(&contract_owner, contract_ids);
+				}
 
-			// Burn contract NFT.
-			let collection_id = Self::contract_collection_id()?;
-			T::NftHelper::burn(&collection_id, &contract_id, Some(&contract_owner))?;
+				contract_owner
+			} else {
+				match op {
+					Operation::Claim => return Err(Error::<T>::CannotClaimUnknownContract.into()),
+					Operation::Cancel => creator,
+					Operation::Snipe => who.clone(),
+				}
+			};
+
+			Self::transfer_items(&Self::staked_items(&contract_id)?, &stake_beneficiary)?;
 
 			// Clean up storage.
 			ContractStakedItems::<T>::remove(contract_id);
 			ContractAccepted::<T>::remove(contract_id);
 			Contracts::<T>::remove(contract_id);
-
-			// Retain contract IDs held.
-			let mut contract_ids = Self::contract_ids(&contract_owner)?;
-			contract_ids.retain(|id| id != &contract_id);
-			if contract_ids.is_empty() {
-				ContractIds::<T>::remove(contract_owner);
-			} else {
-				ContractIds::<T>::insert(contract_owner, contract_ids);
-			}
+			ContractsMetadata::<T>::remove(contract_id);
 
 			// Emit events.
 			match op {
@@ -654,12 +666,16 @@ pub mod pallet {
 			who: &T::AccountId,
 			is_snipe: bool,
 		) -> Result<ContractOf<T>, DispatchError> {
-			let owner = Self::contract_owner(contract_id)?;
+			let maybe_owner = Self::contract_owner(contract_id);
 			if is_snipe {
-				ensure!(owner != Self::account_id(), Error::<T>::Available);
-				ensure!(who != &owner, Error::<T>::CannotSnipeOwnContract);
+				// If it's a snipe and the contract has no owner, we don't apply any checks
+				// since it's possible that the original contract NFT was deleted maliciously
+				if let Ok(owner) = maybe_owner {
+					ensure!(owner != Self::account_id(), Error::<T>::Available);
+					ensure!(who != &owner, Error::<T>::CannotSnipeOwnContract);
+				}
 			} else {
-				ensure!(who == &owner, Error::<T>::ContractOwnership);
+				ensure!(who == &maybe_owner?, Error::<T>::ContractOwnership);
 			}
 			let contract = Self::contract(contract_id)?;
 			Ok(contract)
@@ -755,16 +771,13 @@ pub mod pallet {
 
 			match op {
 				Operation::Claim => {
-					ensure!(now <= expiry, Error::<T>::Expired);
 					ensure!(now >= end, Error::<T>::Staking);
 				},
 				Operation::Cancel => {
-					ensure!(now <= expiry, Error::<T>::Expired);
 					ensure!(now < end, Error::<T>::Claimable);
 				},
 				Operation::Snipe => {
-					ensure!(now >= end, Error::<T>::Staking);
-					ensure!(now > expiry, Error::<T>::Claimable);
+					ensure!(now >= expiry, Error::<T>::Claimable);
 				},
 			}
 			Ok(())

pallets/ajuna-nft-staking/src/tests/ext/cancel.rs

@@ -240,7 +240,7 @@ fn rejects_when_contract_is_expired() {
 			run_to_block(accepted_at + stake_duration + claim_duration + 2);
 			assert_noop!(
 				NftStake::cancel(RuntimeOrigin::signed(BOB), contract_id),
-				Error::<Test>::Expired
+				Error::<Test>::Claimable
 			);
 
 			// Regression for happy case, before reaching claim phase.

pallets/ajuna-nft-staking/src/tests/ext/claim.rs

@@ -25,10 +25,12 @@ fn works_with_token_reward() {
 	];
 	let fee_clauses = vec![(0, Clause::HasAttribute(RESERVED_COLLECTION_2, bounded_vec![33]))];
 	let stake_duration = 4;
+	let claim_duration = 10;
 	let reward_amount = 135;
 	let contract = Contract::default()
 		.reward(Reward::Tokens(reward_amount))
 		.stake_duration(stake_duration)
+		.claim_duration(claim_duration)
 		.stake_clauses(AttributeNamespace::Pallet, stake_clauses.clone())
 		.fee_clauses(AttributeNamespace::Pallet, fee_clauses.clone());
 	let contract_id = H256::random();
@@ -52,7 +54,7 @@ fn works_with_token_reward() {
 			assert_eq!(Balances::free_balance(&technical_account_id), reward_amount);
 
 			let accepted_at = ContractAccepted::<Test>::get(contract_id).unwrap();
-			run_to_block(accepted_at + stake_duration);
+			run_to_block(accepted_at + stake_duration + claim_duration + 1);
 
 			assert_ok!(NftStake::claim(RuntimeOrigin::signed(BOB), contract_id));
 
@@ -67,8 +69,10 @@ fn works_with_token_reward() {
 
 			let contract_collection_id = ContractCollectionId::<Test>::get().unwrap();
 			assert_eq!(Nft::owner(contract_collection_id, contract_id), None);
+			assert_eq!(Contracts::<Test>::get(contract_id), None);
 			assert_eq!(ContractAccepted::<Test>::get(contract_id), None);
 			assert_eq!(ContractStakedItems::<Test>::get(contract_id), None);
+			assert_eq!(ContractsMetadata::<Test>::get(contract_id), None);
 			assert_eq!(ContractIds::<Test>::get(BOB), None);
 
 			System::assert_last_event(RuntimeEvent::NftStake(crate::Event::Claimed {
@@ -243,40 +247,3 @@ fn rejects_when_contract_is_staking() {
 			assert_ok!(NftStake::claim(RuntimeOrigin::signed(BOB), contract_id));
 		});
 }
-
-#[test]
-fn rejects_when_contract_is_expired() {
-	let (stake_duration, claim_duration) = (3, 5);
-	let contract_id = H256::random();
-
-	ExtBuilder::default()
-		.set_creator(ALICE)
-		.create_contract_collection()
-		.create_contract_with_funds(
-			contract_id,
-			Contract::default()
-				.stake_duration(stake_duration)
-				.claim_duration(claim_duration),
-		)
-		.accept_contract(
-			vec![(BOB, Default::default())],
-			vec![(BOB, Default::default())],
-			contract_id,
-			BOB,
-		)
-		.build()
-		.execute_with(|| {
-			let accepted_at = ContractAccepted::<Test>::get(contract_id).unwrap();
-
-			// After expiry.
-			run_to_block(accepted_at + stake_duration + claim_duration + 1);
-			assert_noop!(
-				NftStake::claim(RuntimeOrigin::signed(BOB), contract_id),
-				Error::<Test>::Expired
-			);
-
-			// Regression for happy case, after staking finishes.
-			System::set_block_number(accepted_at + stake_duration);
-			assert_ok!(NftStake::claim(RuntimeOrigin::signed(BOB), contract_id));
-		});
-}

pallets/ajuna-nft-staking/src/tests/ext/remove.rs

@@ -37,7 +37,8 @@ fn works() {
 
 			assert_ok!(NftStake::remove(RuntimeOrigin::signed(ALICE), contract_id));
 
-			assert!(Contracts::<Test>::get(contract_id).is_none());
+			assert_eq!(Contracts::<Test>::get(contract_id), None);
+			assert_eq!(ContractsMetadata::<Test>::get(contract_id), None);
 
 			assert_eq!(Nft::owner(contract_collection_id, contract_id), None);
 

pallets/ajuna-nft-staking/src/tests/ext/snipe.rs

@@ -280,7 +280,7 @@ fn rejects_when_contract_is_staking() {
 				run_to_block(accepted_at + n);
 				assert_noop!(
 					NftStake::snipe(RuntimeOrigin::signed(CHARLIE), contract_id),
-					Error::<Test>::Staking
+					Error::<Test>::Claimable
 				);
 			}
 
@@ -291,41 +291,53 @@ fn rejects_when_contract_is_staking() {
 }
 
 #[test]
-fn rejects_when_contract_is_claimable() {
-	let (stake_duration, claim_duration) = (3, 5);
+fn snipe_with_maliciously_burned_contract_nft() {
+	let stake_clauses = vec![(0, Clause::HasAttribute(RESERVED_COLLECTION_0, bounded_vec![4]))];
+	let fee_clauses = vec![(0, Clause::HasAttribute(RESERVED_COLLECTION_2, bounded_vec![33]))];
+	let (stake_duration, claim_duration) = (4, 3);
+	let reward = 135;
+	let contract = Contract::default()
+		.reward(Reward::Tokens(reward))
+		.stake_duration(stake_duration)
+		.claim_duration(claim_duration)
+		.stake_clauses(AttributeNamespace::Pallet, stake_clauses.clone())
+		.fee_clauses(AttributeNamespace::Pallet, fee_clauses.clone());
 	let contract_id = H256::random();
 
+	let stakes = MockMints::from(MockClauses(stake_clauses.into_iter().map(|(_, c)| c).collect()));
+	let fees = MockMints::from(MockClauses(fee_clauses.into_iter().map(|(_, c)| c).collect()));
+
 	ExtBuilder::default()
 		.set_creator(ALICE)
 		.create_contract_collection()
-		.create_contract_with_funds(
-			contract_id,
-			Contract::default()
-				.stake_duration(stake_duration)
-				.claim_duration(claim_duration),
-		)
-		.accept_contract(
-			vec![(BOB, Default::default())],
-			vec![(BOB, Default::default())],
-			contract_id,
-			BOB,
-		)
-		.create_sniper(CHARLIE, Contract::default().stake_duration(10).claim_duration(10))
+		.create_contract_with_funds(contract_id, contract.clone())
+		.accept_contract(vec![(BOB, stakes)], vec![(BOB, fees)], contract_id, BOB)
+		.create_sniper(CHARLIE, Contract::default().stake_duration(10).claim_duration(20))
 		.build()
 		.execute_with(|| {
 			let accepted_at = ContractAccepted::<Test>::get(contract_id).unwrap();
+			run_to_block(accepted_at + stake_duration + claim_duration + 1);
 
-			// During claim phase.
-			for n in stake_duration..=(stake_duration + claim_duration) {
-				run_to_block(accepted_at + n);
-				assert_noop!(
-					NftStake::snipe(RuntimeOrigin::signed(CHARLIE), contract_id),
-					Error::<Test>::Claimable
-				);
-			}
+			let contract_collection =
+				ContractCollectionId::<Test>::get().expect("Should get collection");
+			assert_ok!(pallet_nfts::Pallet::<Test>::burn(
+				RuntimeOrigin::signed(BOB),
+				contract_collection,
+				contract_id
+			));
 
-			// Regression for happy case, after expiry.
-			System::set_block_number(accepted_at + stake_duration + claim_duration + 1);
 			assert_ok!(NftStake::snipe(RuntimeOrigin::signed(CHARLIE), contract_id));
+
+			let contract_collection_id = ContractCollectionId::<Test>::get().unwrap();
+			assert_eq!(Nft::owner(contract_collection_id, contract_id), None);
+			assert_eq!(ContractAccepted::<Test>::get(contract_id), None);
+			assert_eq!(ContractStakedItems::<Test>::get(contract_id), None);
+			assert_eq!(ContractIds::<Test>::get(BOB), Some(bounded_vec![contract_id]));
+
+			System::assert_last_event(RuntimeEvent::NftStake(crate::Event::Sniped {
+				by: CHARLIE,
+				contract_id,
+				reward: contract.reward,
+			}));
 		});
 }