feat: better handling of DID deletions

pallets/did/src/lib.rs

@@ -88,6 +88,7 @@ pub mod errors;
 pub mod migrations;
 pub mod origin;
 pub mod service_endpoints;
+pub mod traits;
 
 #[cfg(test)]
 mod mock;
@@ -156,6 +157,7 @@ pub mod pallet {
 			tokens::{Fortitude, Precision, Preservation},
 			StorageVersion,
 		},
+		weights::Weight,
 	};
 	use frame_system::pallet_prelude::*;
 	use kilt_support::{
@@ -171,6 +173,7 @@ pub mod pallet {
 			DidEncryptionKey, DidSignature, DidVerifiableIdentifier, DidVerificationKey, RelationshipDeriveError,
 		},
 		service_endpoints::{utils as service_endpoints_utils, ServiceEndpointId},
+		traits::{DeletionHelper, SteppedDeletion},
 	};
 
 	/// The current storage version.
@@ -328,6 +331,10 @@ pub mod pallet {
 
 		/// Migration manager to handle new created entries
 		type BalanceMigrationManager: BalanceMigrationManager<AccountIdOf<Self>, BalanceOf<Self>>;
+
+		/// Helper trait to aid in cleaning up DID-related resources across all
+		/// runtime pallets.
+		type DeletionHelper: DeletionHelper<Self>;
 	}
 
 	#[pallet::pallet]
@@ -388,6 +395,8 @@ pub mod pallet {
 			/// The new deposit owner.
 			to: AccountIdOf<T>,
 		},
+		CleanupIncomplete,
+		CleanupComplete,
 	}
 
 	#[pallet::error]
@@ -455,6 +464,10 @@ pub mod pallet {
 		/// The number of service endpoints stored under the DID is larger than
 		/// the number of endpoints to delete.
 		MaxStoredEndpointsCountExceeded,
+		/// The DID is linked to other runtime elements that would be left
+		/// dangling if the DID were to be deleted.
+		NonZeroReferences,
+		TooExpensive,
 		/// An error that is not supposed to take place, yet it happened.
 		Internal,
 	}
@@ -531,6 +544,7 @@ pub mod pallet {
 	impl<T: Config> Pallet<T>
 	where
 		T::AccountId: AsRef<[u8; 32]> + From<[u8; 32]>,
+		<<T::DeletionHelper as DeletionHelper<T>>::DeletionIter as Iterator>::Item: SteppedDeletion,
 	{
 		/// Store a new DID on chain, after verifying that the creation
 		/// operation has been signed by the KILT account associated with the
@@ -1254,11 +1268,62 @@ pub mod pallet {
 
 			Ok(())
 		}
+
+		#[pallet::call_index(17)]
+		// TODO: Benchmark base call dispatch + the weight specified in the call
+		#[pallet::weight(*max_weight)]
+		pub fn cleanup_linked_resources(
+			origin: OriginFor<T>,
+			did: DidIdentifierOf<T>,
+			max_weight: Weight,
+		) -> DispatchResult {
+			// If it's a regular transaction, verify the submitter is the deposit owner.
+			if let Ok(sender) = ensure_signed(origin.clone()) {
+				let did_entry = Did::<T>::get(&did).ok_or(Error::<T>::NotFound)?;
+				ensure!(did_entry.deposit.owner == sender, Error::<T>::BadDidOrigin);
+			// If it's a DID-authenticated operation, verify the specified `did`
+			// matches the origin.
+			} else if let Ok(did_origin) = T::EnsureOrigin::ensure_origin(origin) {
+				ensure!(did_origin.subject() == did, Error::<T>::BadDidOrigin);
+			// We don't accept any other origin -> fail.
+			} else {
+				return Err(DispatchError::BadOrigin);
+			};
+
+			// The cleanup logic does not need to get benchmarked as it consumes the
+			// provided weight.
+			#[cfg(not(feature = "runtime-benchmarks"))]
+			{
+				let deletion_iter = T::DeletionHelper::deletion_iter(&did);
+				let mut remaining_weight = max_weight;
+				for next_deletion_step in deletion_iter {
+					let Some(ticket) = next_deletion_step.pre_check(remaining_weight) else {
+						// If we run out of gas while there's still stuff to cleanup, apply the cleanups
+						// so far and generate an event.
+						Self::deposit_event(Event::<T>::CleanupIncomplete);
+						return Ok(());
+					};
+					let consumed_weight = next_deletion_step.execute(ticket);
+					// The `pre_check` should tell us if this step can underflow. In case this still
+					// happens, handle it accordingly.
+					let Some(new_weight) = max_weight.checked_sub(&consumed_weight) else {
+						Self::deposit_event(Event::<T>::CleanupIncomplete);
+						return Ok(());
+					};
+					remaining_weight = new_weight;
+				}
+				// If the whole loop completes, the cleanup is successful.
+				Self::deposit_event(Event::<T>::CleanupComplete);
+			}
+
+			Ok(())
+		}
 	}
 
 	impl<T: Config> Pallet<T>
 	where
 		T::AccountId: AsRef<[u8; 32]> + From<[u8; 32]>,
+		<<T::DeletionHelper as DeletionHelper<T>>::DeletionIter as Iterator>::Item: SteppedDeletion,
 	{
 		/// Try creating a DID.
 		///
@@ -1496,6 +1561,8 @@ pub mod pallet {
 				current_endpoints_count <= endpoints_to_remove,
 				Error::<T>::MaxStoredEndpointsCountExceeded
 			);
+			let is_deletion_possible = T::DeletionHelper::deletion_iter(&did_subject).next().is_none();
+			ensure!(is_deletion_possible, Error::<T>::NonZeroReferences);
 
 			// This one can fail, albeit this should **never** be the case as we check for
 			// the preconditions above.

pallets/did/src/mock.rs

@@ -175,6 +175,7 @@ impl Config for Test {
 	type MaxNumberOfTypesPerService = MaxNumberOfTypesPerService;
 	type MaxNumberOfUrlsPerService = MaxNumberOfUrlsPerService;
 	type BalanceMigrationManager = ();
+	type DeletionHelper = ();
 }
 
 parameter_types! {

pallets/did/src/signature.rs

@@ -24,13 +24,15 @@ use sp_weights::Weight;
 use crate::{
 	did_details::{DidSignature, DidVerificationKeyRelationship},
 	errors::DidError,
+	traits::{DeletionHelper, SteppedDeletion},
 	Config, Did, Pallet, WeightInfo,
 };
 
 pub struct DidSignatureVerify<T>(PhantomData<T>);
 impl<T: Config> VerifySignature for DidSignatureVerify<T>
 where
 	T::AccountId: AsRef<[u8; 32]> + From<[u8; 32]>,
+	<<T::DeletionHelper as DeletionHelper<T>>::DeletionIter as Iterator>::Item: SteppedDeletion,
 {
 	type SignerId = <T as Config>::DidIdentifier;
 	type Payload = Vec<u8>;

pallets/did/src/traits.rs

@@ -0,0 +1,75 @@
+// KILT Blockchain – https://botlabs.org
+// Copyright (C) 2019-2024 BOTLabs GmbH
+
+// The KILT Blockchain is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// The KILT Blockchain is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+// If you feel like getting in touch with us, you can do so at [email protected]
+
+use sp_weights::Weight;
+
+use crate::{Config, DidIdentifierOf};
+
+/// Runtime-injected logic to support the DID pallet in making sure no dangling
+/// references is left before deleting a given DID.
+pub trait DeletionHelper<T>
+where
+	T: Config,
+{
+	type DeletionIter: Iterator;
+
+	fn deletion_iter(did: &DidIdentifierOf<T>) -> (Self::DeletionIter, Weight)
+	where
+		<Self::DeletionIter as Iterator>::Item: SteppedDeletion;
+}
+
+impl<T> DeletionHelper<T> for ()
+where
+	T: Config,
+{
+	type DeletionIter = EmptyIterator;
+
+	fn deletion_iter(_did: &DidIdentifierOf<T>) -> (Self::DeletionIter, Weight) {
+		(EmptyIterator, Weight::zero())
+	}
+}
+
+pub struct EmptyIterator;
+
+impl Iterator for EmptyIterator {
+	type Item = ();
+
+	fn next(&mut self) -> Option<Self::Item> {
+		Some(())
+	}
+}
+
+pub trait SteppedDeletion {
+	type VerifiedInfo;
+
+	fn pre_check(&self, remaining_weight: Weight) -> Option<Self::VerifiedInfo>;
+
+	fn execute(self, info: Self::VerifiedInfo) -> Weight;
+}
+
+impl SteppedDeletion for () {
+	type VerifiedInfo = ();
+
+	fn pre_check(&self, _remaining_weight: Weight) -> Option<Self::VerifiedInfo> {
+		None
+	}
+
+	fn execute(self, _info: Self::VerifiedInfo) -> Weight {
+		Weight::zero()
+	}
+}

runtimes/peregrine/Cargo.toml

@@ -14,12 +14,12 @@ version       = { workspace = true }
 substrate-wasm-builder = { workspace = true }
 
 [dev-dependencies]
-enum-iterator = { workspace = true }
-sp-io         = { workspace = true }
+sp-io = { workspace = true }
 
 [dependencies]
 # External dependencies
 cfg-if             = { workspace = true }
+enum-iterator      = { workspace = true }
 log                = { workspace = true }
 parity-scale-codec = { workspace = true, features = ["derive"] }
 scale-info         = { workspace = true, features = ["derive"] }

runtimes/peregrine/src/lib.rs

@@ -33,12 +33,20 @@ use core::str;
 
 // Polkadot-sdk crates
 use cumulus_pallet_parachain_system::register_validate_block;
-use frame_support::construct_runtime;
+use did::traits::{DeletionHelper, SteppedDeletion};
+use enum_iterator::Sequence;
+use frame_support::{
+	construct_runtime,
+	pallet_prelude::{OptionQuery, ValueQuery},
+	storage_alias,
+};
 use frame_system::{
 	ChainContext, CheckEra, CheckGenesis, CheckNonZeroSender, CheckNonce, CheckSpecVersion, CheckTxVersion, CheckWeight,
 };
 use pallet_transaction_payment::ChargeTransactionPayment;
-use runtime_common::opaque::Header;
+use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
+use runtime_common::{opaque::Header, DidIdentifier};
+use scale_info::TypeInfo;
 use sp_runtime::{create_runtime_str, generic};
 use sp_std::{prelude::*, vec::Vec};
 
@@ -58,9 +66,10 @@ use runtime_apis::_InternalImplRuntimeApis;
 pub use runtime_apis::{api, RuntimeApi};
 mod system;
 use sp_version::RuntimeVersion;
+use sp_weights::Weight;
 pub use system::{SessionKeys, SS_58_PREFIX};
 
-use crate::runtime_apis::RUNTIME_API_VERSION;
+use crate::{runtime_apis::RUNTIME_API_VERSION, weights::rocksdb_weights::constants::RocksDbWeight};
 mod weights;
 pub mod xcm;
 
@@ -215,3 +224,65 @@ pub type SignedExtra = (
 	CheckWeight<Runtime>,
 	ChargeTransactionPayment<Runtime>,
 );
+
+// Stores the new map inside the DID pallet. If value is not present, it returns
+// its default.
+#[storage_alias]
+type DidsUnderDeletion = StorageMap<Did, Blake2_128Concat, DidIdentifier, DidDeletionState, ValueQuery>;
+
+#[derive(Encode, Decode, MaxEncodedLen, TypeInfo, Clone, Debug, Default, Sequence)]
+enum DidDeletionState {
+	#[default]
+	Web3Name,
+	DotName,
+	DidLookup,
+	UniqueLinking,
+}
+
+pub struct TestDidHelper;
+
+impl DeletionHelper for TestDidHelper {
+	type DeletionIter = TestDidDeletionIter;
+
+	fn deletion_iter(did: &DidIdentifier) -> (Self::DeletionIter, Weight)
+	where
+		<Self::DeletionIter as Iterator>::Item: SteppedDeletion,
+	{
+		// Returns the iterator + the cost for reading it from storage
+		(
+			TestDidDeletionIter(Some(DidsUnderDeletion::get(did))),
+			RocksDbWeight::get().read,
+		)
+	}
+}
+
+pub struct TestDidDeletionIter(Option<DidDeletionState>);
+
+impl Iterator for TestDidDeletionIter {
+	type Item = DidDeletionState;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		let current = self.0;
+		self.0 = self.0.next();
+		current
+	}
+}
+
+impl SteppedDeletion for DidDeletionState {
+	type VerifiedInfo = ();
+
+	fn pre_check(&self, remaining_weight: Weight) -> Option<Self::VerifiedInfo> {
+		// All pallets have double maps, so we need to read and write twice per pallet.
+		if remaining_weight > RocksDbWeight::get().reads_writes(2, 2) {
+			Some(())
+		} else {
+			None
+		}
+	}
+
+	fn execute(self, info: Self::VerifiedInfo) -> Weight {
+		match self {
+			DidDeletionState::DidLookup => pallet_did_lookup::ConnectedAccounts::<Runtime>::take(key)...
+		}
+	}
+}