proxy-keystore: Revoke keys with purpose

centrifuge · Jun 21, 2022 · 03863bf · 03863bf
1 parent fc322ae
commit 03863bf
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 78 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pallets/proxy-keystore/Cargo.toml b/pallets/proxy-keystore/Cargo.toml
@@ -21,8 +21,6 @@ frame-system = { git = "https://github.com/paritytech/substrate", default-featur
 sp-runtime = { git = "https://github.com/paritytech/substrate", default-features = false , branch = "polkadot-v0.9.20" }
 sp-std = { git = "https://github.com/paritytech/substrate",  default-features = false , branch = "polkadot-v0.9.20" }
 
-enum-iterator = { version = "1.1.1" }
-
 [dev-dependencies]
 sp-core = { git = "https://github.com/paritytech/substrate", default-features = false , branch = "polkadot-v0.9.20" }
 sp-io = { git = "https://github.com/paritytech/substrate", default-features = false , branch = "polkadot-v0.9.20" }

diff --git a/pallets/proxy-keystore/src/benchmarking.rs b/pallets/proxy-keystore/src/benchmarking.rs
@@ -47,7 +47,7 @@ benchmarks! {
 
 		let key_hashes: Vec<T::Hash> = test_keys.iter().map(|add_key| add_key.key).collect();
 		let origin = RawOrigin::Signed(caller.clone());
-	}: revoke_keys(origin, key_hashes)
+	}: revoke_keys(origin, key_hashes, KeyPurpose::P2PDiscovery)
 	verify {
 		assert_eq!(Keys::<T>::iter().collect::<Vec<_>>().len() as u32, n);
 		assert!(all_keys_are_revoked::<T>(caller.clone()));

diff --git a/pallets/proxy-keystore/src/lib.rs b/pallets/proxy-keystore/src/lib.rs
@@ -12,7 +12,6 @@
 // GNU General Public License for more details.
 #![cfg_attr(not(feature = "std"), no_std)]
 
-use enum_iterator::{all, Sequence};
 use frame_support::pallet_prelude::*;
 pub use pallet::*;
 use scale_info::TypeInfo;
@@ -28,9 +27,7 @@ mod tests;
 
 pub mod weights;
 
-// make sure representation is 1 byte
-// TODO(cdamian): Given the above, should we use #[pallet::compact]?
-#[derive(Encode, Decode, Clone, Eq, PartialEq, RuntimeDebug, TypeInfo, Sequence)]
+#[derive(Encode, Decode, Clone, Eq, PartialEq, RuntimeDebug, TypeInfo)]
 pub enum KeyPurpose {
 	P2PDiscovery,
 	P2PDocumentSigning,
@@ -184,9 +181,9 @@ pub mod pallet {
 			Ok(())
 		}
 
-		/// Remove keys from the storages.
+		/// Revoke keys with specified purpose.
 		#[pallet::weight(T::WeightInfo::revoke_keys(T::MaxKeys::get() as u32))]
-		pub fn revoke_keys(origin: OriginFor<T>, keys: Vec<T::Hash>) -> DispatchResult {
+		pub fn revoke_keys(origin: OriginFor<T>, keys: Vec<T::Hash>, key_purpose: KeyPurpose) -> DispatchResult {
 			let account_id = ensure_signed(origin)?;
 
 			ensure!(keys.len() > 0, Error::<T>::NoKeys);
@@ -196,7 +193,7 @@ pub mod pallet {
 			);
 
 			for key in keys {
-				Self::revoke_key(account_id.clone(), key.clone())?;
+				Self::revoke_key(account_id.clone(), key, key_purpose.clone())?;
 			}
 
 			Ok(())
@@ -265,47 +262,35 @@ pub mod pallet {
 
 		/// Revoke a key at the current `block_number` in the `Keys` storage
 		/// if the key is found and it's *not* already revoked.
-		fn revoke_key(account_id: T::AccountId, key: T::Hash) -> DispatchResult {
-			let mut key_found = false;
-
-			for key_purpose in all::<KeyPurpose>() {
-				let key_id: KeyId<T::Hash> = (key.clone(), key_purpose.clone());
+		fn revoke_key(account_id: T::AccountId, key: T::Hash, key_purpose: KeyPurpose) -> DispatchResult {
+			let key_id: KeyId<T::Hash> = (key, key_purpose);
 
-				<Keys<T>>::mutate(
-					account_id.clone(),
-					key_id,
-					|storage_key_opt| -> DispatchResult {
-						match storage_key_opt {
-							Some(storage_key) => {
-								if let Some(_) = storage_key.revoked_at {
-									return Err(Error::<T>::KeyAlreadyRevoked.into());
-								}
-
-								key_found = true;
+			<Keys<T>>::try_mutate(
+				account_id.clone(),
+				key_id,
+				|storage_key_opt| -> DispatchResult {
+					match storage_key_opt {
+						Some(storage_key) => {
+							if let Some(_) = storage_key.revoked_at {
+								return Err(Error::<T>::KeyAlreadyRevoked.into());
+							}
 
-								let block_number = <frame_system::Pallet<T>>::block_number();
+							let block_number = <frame_system::Pallet<T>>::block_number();
 
-								storage_key.revoked_at = Some(block_number.clone());
+							storage_key.revoked_at = Some(block_number.clone());
 
-								Self::deposit_event(Event::KeyRevoked(
-									account_id.clone(),
-									key.clone(),
-									block_number,
-								));
+							Self::deposit_event(Event::KeyRevoked(
+								account_id.clone(),
+								key.clone(),
+								block_number,
+							));
 
-								Ok(())
-							}
-							None => Ok(()),
+							Ok(())
 						}
-					},
-				)?;
-			}
-
-			if !key_found {
-				return Err(Error::<T>::KeyNotFound.into());
-			}
-
-			Ok(().into())
+						None => return Err(Error::<T>::KeyNotFound.into()),
+					}
+				},
+			)
 		}
 	}
 }
diff --git a/pallets/proxy-keystore/src/tests.rs b/pallets/proxy-keystore/src/tests.rs
@@ -147,19 +147,27 @@ fn revoke_keys() {
 
 		insert_test_keys_in_storage(origin, keys.clone());
 
-		let key_hashes: Vec<H256> = keys.iter().map(|add_key| add_key.key).collect();
+		for key in keys.clone() {
+			let vec: Vec<H256> = vec![key.key];
+
+			assert_ok!(
+				ProxyKeystore::revoke_keys(
+					Origin::signed(origin),
+					vec,
+					key.purpose,
+				),
+			);
+		}
 
-		assert_ok!(ProxyKeystore::revoke_keys(
-			Origin::signed(origin),
-			key_hashes.clone()
-		));
 		// Keys are still in storage but should be revoked.
 		assert_eq!(
 			Keys::<MockRuntime>::iter().collect::<Vec<_>>().len(),
 			2,
 			"keys should still be in storage"
 		);
 
+		let key_hashes: Vec<H256> = keys.iter().map(|add_key| add_key.key).collect();
+
 		keys_are_revoked(key_hashes);
 
 		assert_eq!(
@@ -187,13 +195,13 @@ fn revoke_keys() {
 fn revoke_keys_key_errors() {
 	// 0 Keys
 	new_test_ext().execute_with(|| {
-		let keys: Vec<AddKey<H256>> = Vec::new();
-		let key_hashes: Vec<H256> = keys.iter().map(|add_key| add_key.key).collect();
+		let keys: Vec<H256> = Vec::new();
 
 		assert_err!(
-			ProxyKeystore::revoke_keys(Origin::signed(1), key_hashes.clone()),
+			ProxyKeystore::revoke_keys(Origin::signed(1), keys, KeyPurpose::P2PDocumentSigning),
 			Error::<MockRuntime>::NoKeys
 		);
+
 		assert_eq!(
 			Keys::<MockRuntime>::iter().collect::<Vec<_>>().len(),
 			0,
@@ -209,7 +217,7 @@ fn revoke_keys_key_errors() {
 		let key_hashes: Vec<H256> = keys.iter().map(|add_key| add_key.key).collect();
 
 		assert_err!(
-			ProxyKeystore::revoke_keys(Origin::signed(1), key_hashes),
+			ProxyKeystore::revoke_keys(Origin::signed(1), key_hashes, KeyPurpose::P2PDocumentSigning),
 			Error::<MockRuntime>::TooManyKeys
 		);
 		assert_eq!(
@@ -228,7 +236,12 @@ fn revoke_keys_key_not_found() {
 		let key_hashes: Vec<H256> = keys.iter().map(|add_key| add_key.key).collect();
 
 		assert_err!(
-			ProxyKeystore::revoke_keys(Origin::signed(origin), key_hashes.clone()),
+			ProxyKeystore::revoke_keys(Origin::signed(origin), key_hashes.clone(), KeyPurpose::P2PDocumentSigning),
+			Error::<MockRuntime>::KeyNotFound
+		);
+
+		assert_err!(
+			ProxyKeystore::revoke_keys(Origin::signed(origin), key_hashes.clone(), KeyPurpose::P2PDiscovery),
 			Error::<MockRuntime>::KeyNotFound
 		);
 	});
@@ -238,8 +251,9 @@ fn revoke_keys_key_not_found() {
 fn revoke_keys_key_already_revoked() {
 	new_test_ext().execute_with(|| {
 		let origin: u64 = 1;
+		let key_purpose = KeyPurpose::P2PDocumentSigning;
 		let key = Key {
-			purpose: KeyPurpose::P2PDocumentSigning,
+			purpose: key_purpose.clone(),
 			key_type: KeyType::EDDSA,
 			revoked_at: Some(1),
 			deposit: 1,
@@ -252,7 +266,7 @@ fn revoke_keys_key_already_revoked() {
 		let key_hashes: Vec<H256> = vec![key_id.0];
 
 		assert_err!(
-			ProxyKeystore::revoke_keys(Origin::signed(origin), key_hashes.clone()),
+			ProxyKeystore::revoke_keys(Origin::signed(origin), key_hashes.clone(), key_purpose),
 			Error::<MockRuntime>::KeyAlreadyRevoked
 		);
 	});