Tracking OmniAccount authentication on chain (#3218)

* adding auth_token_requested call * including auth_type to omni_account calls dispatchers * adjusting worker to send auth_type on dispatch_as_omni_account and dispatch_as_signed * sending auth_token_requested call on token requested * including extrinsic params and adding error logs
litentry · Dec 24, 2024 · d21a22c · d21a22c
1 parent 6b179a2
commit d21a22c
Show file tree

Hide file tree

Showing 8 changed files with 781 additions and 1,157 deletions.
diff --git a/common/primitives/core/Cargo.lock b/common/primitives/core/Cargo.lock
diff --git a/common/primitives/core/src/omni_account.rs b/common/primitives/core/src/omni_account.rs
@@ -19,6 +19,14 @@ use parity_scale_codec::{Decode, Encode};
 use scale_info::TypeInfo;
 use sp_runtime::RuntimeDebug;
 
+#[derive(Encode, Decode, TypeInfo, Clone, PartialEq, Eq, RuntimeDebug)]
+pub enum OmniAccountAuthType {
+    Web3,
+    Email,
+    OAuth2,
+    AuthToken,
+}
+
 #[derive(Encode, Decode, TypeInfo, Clone, PartialEq, Eq, RuntimeDebug)]
 pub enum MemberAccount {
     Public(Identity),

diff --git a/parachain/pallets/omni-account/src/lib.rs b/parachain/pallets/omni-account/src/lib.rs
@@ -21,7 +21,9 @@ mod mock;
 #[cfg(test)]
 mod tests;
 
-pub use core_primitives::{Identity, Intent, MemberAccount, OmniAccountConverter};
+pub use core_primitives::{
+	Identity, Intent, MemberAccount, OmniAccountAuthType, OmniAccountConverter,
+};
 pub use frame_system::{self as system, pallet_prelude::BlockNumberFor};
 pub use pallet::*;
 
@@ -136,13 +138,23 @@ pub mod pallet {
 		/// An account store is updated
 		AccountStoreUpdated { who: T::AccountId, account_store: MemberAccounts<T> },
 		/// Some call is dispatched as omni-account origin
-		DispatchedAsOmniAccount { who: T::AccountId, result: DispatchResult },
+		DispatchedAsOmniAccount {
+			who: T::AccountId,
+			auth_type: OmniAccountAuthType,
+			result: DispatchResult,
+		},
 		/// Some call is dispatched as signed origin
-		DispatchedAsSigned { who: T::AccountId, result: DispatchResult },
+		DispatchedAsSigned {
+			who: T::AccountId,
+			auth_type: OmniAccountAuthType,
+			result: DispatchResult,
+		},
 		/// Intent is requested
 		IntentRequested { who: T::AccountId, intent: Intent },
 		/// Intent is executed
 		IntentExecuted { who: T::AccountId, intent: Intent, result: IntentExecutionResult },
+		/// An auth token is requested
+		AuthTokenRequested { who: T::AccountId, expires_at: BlockNumberFor<T> },
 	}
 
 	#[pallet::error]
@@ -164,6 +176,7 @@ pub mod pallet {
 			origin: OriginFor<T>,
 			member_account_hash: H256,
 			call: Box<<T as Config>::RuntimeCall>,
+			auth_type: OmniAccountAuthType,
 		) -> DispatchResultWithPostInfo {
 			let _ = T::TEECallOrigin::ensure_origin(origin)?;
 			let omni_account = MemberAccountHash::<T>::get(member_account_hash)
@@ -172,6 +185,7 @@ pub mod pallet {
 			system::Pallet::<T>::inc_account_nonce(&omni_account);
 			Self::deposit_event(Event::DispatchedAsOmniAccount {
 				who: omni_account,
+				auth_type,
 				result: result.map(|_| ()).map_err(|e| e.error),
 			});
 			Ok(Pays::No.into())
@@ -185,6 +199,7 @@ pub mod pallet {
 			origin: OriginFor<T>,
 			member_account_hash: H256,
 			call: Box<<T as Config>::RuntimeCall>,
+			auth_type: OmniAccountAuthType,
 		) -> DispatchResultWithPostInfo {
 			let _ = T::TEECallOrigin::ensure_origin(origin)?;
 			let omni_account = MemberAccountHash::<T>::get(member_account_hash)
@@ -196,6 +211,7 @@ pub mod pallet {
 			system::Pallet::<T>::inc_account_nonce(&omni_account);
 			Self::deposit_event(Event::DispatchedAsSigned {
 				who: omni_account,
+				auth_type,
 				result: result.map(|_| ()).map_err(|e| e.error),
 			});
 			Ok(Pays::No.into())
@@ -359,6 +375,18 @@ pub mod pallet {
 			Self::deposit_event(Event::IntentExecuted { who, intent, result });
 			Ok(())
 		}
+
+		#[pallet::call_index(9)]
+		#[pallet::weight((195_000_000, DispatchClass::Normal))]
+		pub fn auth_token_requested(
+			origin: OriginFor<T>,
+			who: T::AccountId,
+			expires_at: BlockNumberFor<T>,
+		) -> DispatchResult {
+			let _ = T::TEECallOrigin::ensure_origin(origin)?;
+			Self::deposit_event(Event::AuthTokenRequested { who, expires_at });
+			Ok(())
+		}
 	}
 
 	impl<T: Config> Pallet<T> {

diff --git a/parachain/pallets/omni-account/src/tests.rs b/parachain/pallets/omni-account/src/tests.rs
@@ -88,6 +88,7 @@ fn add_account_without_creating_store_fails() {
 				RuntimeOrigin::signed(tee_signer.clone()),
 				alice().identity.hash(),
 				call,
+				OmniAccountAuthType::Web3
 			),
 			Error::<Test>::AccountNotFound
 		);
@@ -115,6 +116,7 @@ fn add_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
@@ -142,6 +144,7 @@ fn add_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 		let expected_member_accounts: MemberAccounts<Test> = BoundedVec::truncate_from(vec![
 			public_member_account(alice()),
@@ -200,17 +203,20 @@ fn add_account_with_already_linked_account_fails() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call.clone(),
+			OmniAccountAuthType::Web3
 		));
 
 		assert_ok!(OmniAccount::dispatch_as_omni_account(
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: Err(DispatchError::Module(ModuleError {
 					index: 5,
 					error: [0, 0, 0, 0],
@@ -231,11 +237,13 @@ fn add_account_with_already_linked_account_fails() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			charlie().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: Err(DispatchError::Module(ModuleError {
 					index: 5,
 					error: [0, 0, 0, 0],
@@ -276,11 +284,13 @@ fn add_account_store_len_limit_reached_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: Err(DispatchError::Module(ModuleError {
 					index: 5,
 					error: [1, 0, 0, 0],
@@ -308,6 +318,7 @@ fn remove_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		// normal signed origin should give `BadOrigin`, no matter
@@ -333,11 +344,13 @@ fn remove_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: DispatchResult::Ok(()),
 			}
 			.into(),
@@ -372,6 +385,7 @@ fn remove_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		assert!(!AccountStore::<Test>::contains_key(alice().omni_account));
@@ -394,6 +408,7 @@ fn remove_account_empty_account_check_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		let call = remove_accounts_call(vec![]);
@@ -402,10 +417,12 @@ fn remove_account_empty_account_check_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: Err(DispatchError::Module(ModuleError {
 					index: 5,
 					error: [5, 0, 0, 0],
@@ -434,6 +451,7 @@ fn publicize_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		let expected_member_accounts: MemberAccounts<Test> =
@@ -448,11 +466,13 @@ fn publicize_account_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: DispatchResult::Ok(()),
 			}
 			.into(),
@@ -495,6 +515,7 @@ fn publicize_account_identity_not_found_works() {
 				RuntimeOrigin::signed(tee_signer.clone()),
 				alice().identity.hash(),
 				call,
+				OmniAccountAuthType::Web3
 			),
 			Error::<Test>::AccountNotFound
 		);
@@ -509,17 +530,20 @@ fn publicize_account_identity_not_found_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		let call = publicize_account_call(charlie().identity);
 		assert_ok!(OmniAccount::dispatch_as_omni_account(
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: Err(DispatchError::Module(ModuleError {
 					index: 5,
 					error: [2, 0, 0, 0],
@@ -547,6 +571,7 @@ fn request_intent_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		let intent =
@@ -557,11 +582,13 @@ fn request_intent_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		System::assert_has_event(
 			Event::DispatchedAsOmniAccount {
 				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
 				result: DispatchResult::Ok(()),
 			}
 			.into(),
@@ -594,17 +621,23 @@ fn dispatch_as_signed_works() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 
 		let call = make_balance_transfer_call(bob().native_account, 5);
 		assert_ok!(OmniAccount::dispatch_as_signed(
 			RuntimeOrigin::signed(tee_signer),
 			alice().identity.hash(),
-			call
+			call,
+			OmniAccountAuthType::Web3
 		));
 		System::assert_has_event(
-			Event::DispatchedAsSigned { who: alice().omni_account, result: DispatchResult::Ok(()) }
-				.into(),
+			Event::DispatchedAsSigned {
+				who: alice().omni_account,
+				auth_type: OmniAccountAuthType::Web3,
+				result: DispatchResult::Ok(()),
+			}
+			.into(),
 		);
 
 		assert_eq!(Balances::free_balance(bob().native_account), 5);
@@ -630,6 +663,7 @@ fn dispatch_as_omni_account_increments_omni_account_nonce() {
 			RuntimeOrigin::signed(tee_signer.clone()),
 			alice().identity.hash(),
 			call,
+			OmniAccountAuthType::Web3
 		));
 		assert_eq!(System::account_nonce(alice().omni_account), 1);
 	});
@@ -658,8 +692,27 @@ fn dispatch_as_signed_account_increments_omni_account_nonce() {
 		assert_ok!(OmniAccount::dispatch_as_signed(
 			RuntimeOrigin::signed(tee_signer),
 			alice().identity.hash(),
-			call
+			call,
+			OmniAccountAuthType::Web3
 		));
 		assert_eq!(System::account_nonce(alice().omni_account), 1);
 	});
 }
+
+#[test]
+fn auth_token_requested_works() {
+	new_test_ext().execute_with(|| {
+		let tee_signer = get_tee_signer();
+
+		assert_ok!(OmniAccount::auth_token_requested(
+			RuntimeOrigin::signed(tee_signer.clone()),
+			alice().identity.to_omni_account(),
+			10
+		));
+
+		System::assert_last_event(
+			Event::AuthTokenRequested { who: alice().identity.to_omni_account(), expires_at: 10 }
+				.into(),
+		);
+	});
+}
diff --git a/tee-worker/identity/litentry/core/authentication/src/jwt.rs b/tee-worker/identity/litentry/core/authentication/src/jwt.rs
@@ -48,7 +48,7 @@ pub struct Header {
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
 pub struct Payload {
 	sub: String,
-	exp: BlockNumber,
+	pub exp: BlockNumber,
 }
 
 impl Payload {