atomic IMT pallet operations

tee-worker/app-libs/stf/src/trusted_call_litentry.rs

@@ -68,23 +68,20 @@ impl TrustedCallSigned {
 		let key = IdentityManagement::user_shielding_keys(&who)
 			.ok_or(StfError::CreateIdentityFailed(ErrorDetail::UserShieldingKeyNotFound))?;
 
+		// generate challenge code
+		let code = generate_challenge_code();
+
 		IMTCall::create_identity {
-			who: who.clone(),
-			identity: identity.clone(),
+			who,
+			identity,
 			metadata,
 			creation_request_block: bn,
 			parent_ss58_prefix,
+			code,
 		}
 		.dispatch_bypass_filter(RuntimeOrigin::root())
 		.map_err(|e| StfError::CreateIdentityFailed(e.into()))?;
 
-		// generate challenge code
-		let code = generate_challenge_code();
-
-		IMTCall::set_challenge_code { who, identity, code }
-			.dispatch_bypass_filter(RuntimeOrigin::root())
-			.map_err(|e| StfError::CreateIdentityFailed(e.into()))?;
-
 		Ok((key, code))
 	}
 
@@ -192,16 +189,7 @@ impl TrustedCallSigned {
 		let key = IdentityManagement::user_shielding_keys(&who)
 			.ok_or(StfError::VerifyIdentityFailed(ErrorDetail::UserShieldingKeyNotFound))?;
 
-		IMTCall::verify_identity {
-			who: who.clone(),
-			identity: identity.clone(),
-			verification_request_block: bn,
-		}
-		.dispatch_bypass_filter(RuntimeOrigin::root())
-		.map_err(|e| StfError::VerifyIdentityFailed(e.into()))?;
-
-		// remove challenge code
-		IMTCall::remove_challenge_code { who, identity }
+		IMTCall::verify_identity { who, identity, verification_request_block: bn }
 			.dispatch_bypass_filter(RuntimeOrigin::root())
 			.map_err(|e| StfError::VerifyIdentityFailed(e.into()))?;
 

tee-worker/litentry/pallets/identity-management/src/lib.rs

@@ -219,8 +219,7 @@ pub mod pallet {
 				ChallengeCodes::<T>::contains_key(&who, &identity),
 				Error::<T>::ChallengeCodeNotExist
 			);
-			ChallengeCodes::<T>::remove(&who, &identity);
-			Self::deposit_event(Event::ChallengeCodeRemoved { who, identity });
+			Self::do_remove_challenge_code(who, identity);
 			Ok(())
 		}
 
@@ -233,6 +232,7 @@ pub mod pallet {
 			metadata: Option<MetadataOf<T>>,
 			creation_request_block: ParentchainBlockNumber,
 			parent_ss58_prefix: u16,
+			code: ChallengeCode,
 		) -> DispatchResult {
 			T::ManageOrigin::ensure_origin(origin)?;
 
@@ -259,7 +259,12 @@ pub mod pallet {
 				..Default::default()
 			};
 			Self::insert_identity_with_limit(&who, &identity, context)?;
-			Self::deposit_event(Event::IdentityCreated { who, identity });
+			Self::deposit_event(Event::IdentityCreated {
+				who: who.clone(),
+				identity: identity.clone(),
+			});
+			ChallengeCodes::<T>::insert(&who, &identity, code);
+			Self::deposit_event(Event::ChallengeCodeSet { who, identity, code });
 			Ok(())
 		}
 
@@ -302,27 +307,33 @@ pub mod pallet {
 			verification_request_block: ParentchainBlockNumber,
 		) -> DispatchResult {
 			T::ManageOrigin::ensure_origin(origin)?;
-			IDGraphs::<T>::try_mutate(&who, &identity, |context| -> DispatchResult {
-				let mut c = context.take().ok_or(Error::<T>::IdentityNotExist)?;
-				ensure!(!c.is_verified, Error::<T>::IdentityAlreadyVerified);
-
-				if let Some(b) = c.creation_request_block {
-					ensure!(
-						b <= verification_request_block,
-						Error::<T>::VerificationRequestTooEarly
-					);
-					ensure!(
-						verification_request_block - b <= T::MaxVerificationDelay::get(),
-						Error::<T>::VerificationRequestTooLate
-					);
-					c.is_verified = true;
-					c.verification_request_block = Some(verification_request_block);
-					*context = Some(c);
-					Ok(())
-				} else {
-					Err(Error::<T>::IdentityNotCreated.into())
-				}
-			})
+			let verify_result =
+				IDGraphs::<T>::try_mutate(&who, &identity, |context| -> DispatchResult {
+					let mut c = context.take().ok_or(Error::<T>::IdentityNotExist)?;
+					ensure!(!c.is_verified, Error::<T>::IdentityAlreadyVerified);
+
+					if let Some(b) = c.creation_request_block {
+						ensure!(
+							b <= verification_request_block,
+							Error::<T>::VerificationRequestTooEarly
+						);
+						ensure!(
+							verification_request_block - b <= T::MaxVerificationDelay::get(),
+							Error::<T>::VerificationRequestTooLate
+						);
+						c.is_verified = true;
+						c.verification_request_block = Some(verification_request_block);
+						*context = Some(c);
+						Ok(())
+					} else {
+						Err(Error::<T>::IdentityNotCreated.into())
+					}
+				});
+			// remove challenge code associated with this identity
+			if ChallengeCodes::<T>::contains_key(&who, &identity) {
+				Self::do_remove_challenge_code(who, identity);
+			}
+			verify_result
 		}
 	}
 
@@ -367,6 +378,11 @@ pub mod pallet {
 			IDGraphs::<T>::remove(owner, identity);
 		}
 
+		fn do_remove_challenge_code(who: T::AccountId, identity: Identity) {
+			ChallengeCodes::<T>::remove(&who, &identity);
+			Self::deposit_event(Event::ChallengeCodeRemoved { who, identity });
+		}
+
 		pub fn get_id_graph(who: &T::AccountId) -> IDGraph<T> {
 			IDGraphs::iter_prefix(who).collect::<Vec<_>>()
 		}

tee-worker/litentry/pallets/identity-management/src/tests.rs

@@ -19,12 +19,16 @@ use crate::{
 	UserShieldingKeyType,
 };
 use frame_support::{assert_err, assert_noop, assert_ok, traits::Get};
-use litentry_primitives::{Identity, IdentityString, Web2Network, USER_SHIELDING_KEY_LEN};
+use litentry_primitives::{
+	Identity, IdentityString, Web2Network, CHALLENGE_CODE_SIZE, USER_SHIELDING_KEY_LEN,
+};
 use sp_runtime::AccountId32;
 
 pub const ALICE: AccountId32 = AccountId32::new([1u8; 32]);
 pub const BOB: AccountId32 = AccountId32::new([2u8; 32]);
 
+pub const SAMPLE_CHALLENGE_CODE: [u8; CHALLENGE_CODE_SIZE] = [0u8; CHALLENGE_CODE_SIZE];
+
 #[test]
 fn set_user_shielding_key_works() {
 	new_test_ext(false).execute_with(|| {
@@ -59,6 +63,7 @@ fn create_identity_works() {
 			Some(metadata.clone()),
 			1,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
 		assert_eq!(
 			IMT::id_graphs(BOB, alice_web3_identity()).unwrap(),
@@ -70,6 +75,7 @@ fn create_identity_works() {
 			}
 		);
 		assert_eq!(crate::IDGraphLens::<Test>::get(&BOB), 2);
+		assert!(crate::ChallengeCodes::<Test>::get(&BOB, alice_web3_identity()).is_some())
 	});
 }
 
@@ -85,6 +91,7 @@ fn cannot_create_more_identities_for_account_than_limit() {
 				None,
 				i,
 				131_u16,
+				SAMPLE_CHALLENGE_CODE
 			));
 		}
 		assert_err!(
@@ -95,6 +102,7 @@ fn cannot_create_more_identities_for_account_than_limit() {
 				None,
 				max_id_graph_len + 1,
 				131_u16,
+				SAMPLE_CHALLENGE_CODE
 			),
 			Error::<Test>::IDGraphLenLimitReached
 		);
@@ -106,6 +114,7 @@ fn remove_identity_works() {
 	new_test_ext(false).execute_with(|| {
 		let ss58_prefix = 31_u16;
 		let shielding_key: UserShieldingKeyType = [0u8; USER_SHIELDING_KEY_LEN];
+
 		assert_ok!(IMT::set_user_shielding_key(
 			RuntimeOrigin::signed(ALICE),
 			BOB,
@@ -125,6 +134,7 @@ fn remove_identity_works() {
 			Some(metadata.clone()),
 			1,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
 		assert_eq!(
 			IMT::id_graphs(BOB, alice_web3_identity()).unwrap(),
@@ -167,7 +177,16 @@ fn verify_identity_works() {
 			Some(metadata.clone()),
 			1,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
+		assert_ok!(IMT::set_challenge_code(
+			RuntimeOrigin::signed(ALICE),
+			BOB,
+			alice_web3_identity(),
+			SAMPLE_CHALLENGE_CODE
+		));
+		assert!(IMT::challenge_codes(BOB, alice_web3_identity()).is_some());
+
 		assert_ok!(IMT::verify_identity(
 			RuntimeOrigin::signed(ALICE),
 			BOB,
@@ -183,6 +202,7 @@ fn verify_identity_works() {
 				is_verified: true,
 			}
 		);
+		assert!(IMT::challenge_codes(BOB, alice_web3_identity()).is_none());
 	});
 }
 
@@ -198,6 +218,7 @@ fn get_id_graph_works() {
 			Some(metadata3.clone()),
 			3,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
 		assert_ok!(IMT::verify_identity(
 			RuntimeOrigin::signed(ALICE),
@@ -218,6 +239,7 @@ fn get_id_graph_works() {
 			Some(metadata2.clone()),
 			2,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
 		assert_ok!(IMT::verify_identity(
 			RuntimeOrigin::signed(ALICE),
@@ -240,13 +262,15 @@ fn verify_identity_fails_when_too_early() {
 
 		let metadata: MetadataOf<Test> = vec![0u8; 16].try_into().unwrap();
 		let ss58_prefix = 131_u16;
+
 		assert_ok!(IMT::create_identity(
 			RuntimeOrigin::signed(ALICE),
 			BOB,
 			alice_web3_identity(),
 			Some(metadata.clone()),
 			CREATION_REQUEST_BLOCK,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
 		assert_noop!(
 			IMT::verify_identity(
@@ -277,13 +301,15 @@ fn verify_identity_fails_when_too_late() {
 
 		let metadata: MetadataOf<Test> = vec![0u8; 16].try_into().unwrap();
 		let ss58_prefix = 131_u16;
+
 		assert_ok!(IMT::create_identity(
 			RuntimeOrigin::signed(ALICE),
 			BOB,
 			alice_web3_identity(),
 			Some(metadata.clone()),
 			CREATION_REQUEST_BLOCK,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE
 		));
 		assert_noop!(
 			IMT::verify_identity(
@@ -318,6 +344,7 @@ fn get_id_graph_with_max_len_works() {
 				None,
 				i,
 				131_u16,
+				SAMPLE_CHALLENGE_CODE
 			));
 		}
 		// the full id_graph should have 22 elements, including the prime_id
@@ -344,13 +371,15 @@ fn id_graph_stats_works() {
 	new_test_ext(true).execute_with(|| {
 		let metadata: MetadataOf<Test> = vec![0u8; 16].try_into().unwrap();
 		let ss58_prefix = 131_u16;
+
 		IMT::create_identity(
 			RuntimeOrigin::signed(ALICE),
 			ALICE,
 			alice_web3_identity(),
 			Some(metadata.clone()),
 			1,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE,
 		)
 		.unwrap();
 		IMT::create_identity(
@@ -360,6 +389,7 @@ fn id_graph_stats_works() {
 			Some(metadata.clone()),
 			1,
 			ss58_prefix,
+			SAMPLE_CHALLENGE_CODE,
 		)
 		.unwrap();
 

tee-worker/ts-tests/interfaces/augment-api-tx.ts

@@ -213,9 +213,10 @@ declare module '@polkadot/api-base/types/submittable' {
                         | Uint8Array,
                     metadata: Option<Bytes> | null | Uint8Array | Bytes | string,
                     creationRequestBlock: u32 | AnyNumber | Uint8Array,
-                    parentSs58Prefix: u16 | AnyNumber | Uint8Array
+                    parentSs58Prefix: u16 | AnyNumber | Uint8Array,
+                    code: U8aFixed | string | Uint8Array
                 ) => SubmittableExtrinsic<ApiType>,
-                [AccountId32, LitentryPrimitivesIdentity, Option<Bytes>, u32, u16]
+                [AccountId32, LitentryPrimitivesIdentity, Option<Bytes>, u32, u16, U8aFixed]
             >;
             removeChallengeCode: AugmentedSubmittable<
                 (

tee-worker/ts-tests/interfaces/lookup.ts

@@ -601,6 +601,7 @@ export default {
                 metadata: 'Option<Bytes>',
                 creationRequestBlock: 'u32',
                 parentSs58Prefix: 'u16',
+                code: '[u8;16]',
             },
             remove_identity: {
                 who: 'AccountId32',

tee-worker/ts-tests/interfaces/types-lookup.ts

@@ -709,6 +709,7 @@ declare module '@polkadot/types/lookup' {
             readonly metadata: Option<Bytes>;
             readonly creationRequestBlock: u32;
             readonly parentSs58Prefix: u16;
+            readonly code: U8aFixed;
         } & Struct;
         readonly isRemoveIdentity: boolean;
         readonly asRemoveIdentity: {