refactoring submit_aes_request method

tee-worker/omni-executor/rpc-server/src/methods/mod.rs

@@ -1,12 +1,12 @@
 mod get_shielding_key;
-mod submit_native_aes_request;
+mod submit_aes_request;
 
 use crate::server::RpcContext;
 use get_shielding_key::register_get_shielding_key;
 use jsonrpsee::RpcModule;
-use submit_native_aes_request::register_native_submit_aes_request;
+use submit_aes_request::register_submit_aes_request;
 
 pub fn register_methods(module: &mut RpcModule<RpcContext>) {
 	register_get_shielding_key(module);
-	register_native_submit_aes_request(module);
+	register_submit_aes_request(module);
 }

tee-worker/omni-executor/rpc-server/src/methods/submit_native_aes_request.rs → tee-worker/omni-executor/rpc-server/src/methods/submit_aes_request.rs

@@ -1,60 +1,31 @@
 use crate::{
+	authentication::{
+		verify_auth_token_authentication, verify_email_authentication, verify_web3_authentication,
+		Authentication,
+	},
 	error_code::*,
+	request::{AesRequest, DecryptableRequest},
 	server::RpcContext,
 	utils::hex::{FromHexPrefixed, ToHexPrefixed},
 };
-use crypto::{
-	aes256::{aes_decrypt, Aes256Key, AesOutput},
-	traits::Decrypt,
-};
 use jsonrpsee::{
 	types::{ErrorCode, ErrorObject},
 	RpcModule,
 };
 use native_call_executor::NativeCall;
+use parentchain_primitives::Nonce;
 use parity_scale_codec::{Decode, Encode};
-use std::{fmt::Debug, sync::Arc, vec::Vec};
+use std::{fmt::Debug, sync::Arc};
 use tokio::sync::oneshot;
 
-type MrEnclave = [u8; 32];
-
-#[derive(Encode, Decode, Default, Clone, PartialEq, Eq, Debug)]
-pub struct AesRequest {
-	pub mrenclave: MrEnclave,
-	pub key: Vec<u8>,
-	pub payload: AesOutput,
-}
-
-impl AesRequest {
-	fn decrypt<T: Debug>(
-		&mut self,
-		shielding_key: Arc<impl Decrypt<Error = T>>,
-	) -> Result<Vec<u8>, ()> {
-		let aes_key: Aes256Key =
-			shielding_key.decrypt(&self.key).map_err(|_| ())?.try_into().map_err(|_| ())?;
-
-		aes_decrypt(&aes_key, &mut self.payload).ok_or(())
-	}
-}
-
-pub type VerificationCode = String;
-
-#[derive(Encode, Decode, Clone, Debug, PartialEq, Eq)]
-pub enum Authentication {
-	Email(VerificationCode),
-	AuthToken(String),
-	// OAuth2(OAuth2Data),
-	// Web3(LitentryMultiSignature),
-}
-
 #[derive(Encode, Decode, Clone, Debug, PartialEq, Eq)]
 pub struct AuthenticatedCall {
 	pub call: NativeCall,
-	pub nonce: u32,
+	pub nonce: Nonce,
 	pub authentication: Authentication,
 }
 
-pub fn register_native_submit_aes_request(module: &mut RpcModule<RpcContext>) {
+pub fn register_submit_aes_request(module: &mut RpcModule<RpcContext>) {
 	module
 		.register_async_method("native_submitAesRequest", |params, ctx, _| async move {
 			let Ok(hex_request) = params.one::<String>() else {
@@ -63,7 +34,7 @@ pub fn register_native_submit_aes_request(module: &mut RpcModule<RpcContext>) {
 			let Ok(mut request) = AesRequest::from_hex(&hex_request) else {
 				return Err(ErrorCode::ServerError(INVALID_AES_REQUEST_CODE).into());
 			};
-			let native_call = get_native_call_from_aes_request(&mut request, ctx.clone())?;
+			let native_call = get_native_call_from_aes_request(&mut request, ctx.clone()).await?;
 
 			let (response_sender, response_receiver) = oneshot::channel();
 
@@ -83,32 +54,42 @@ pub fn register_native_submit_aes_request(module: &mut RpcModule<RpcContext>) {
 		.expect("Failed to register submitNativeRequest method");
 }
 
-fn get_native_call_from_aes_request<'a>(
+async fn get_native_call_from_aes_request<'a>(
 	request: &mut AesRequest,
 	ctx: Arc<RpcContext>,
 ) -> Result<NativeCall, ErrorObject<'a>> {
-	if request.mrenclave != ctx.mrenclave {
-		log::error!("Invalid mrenclave");
-		return Err(ErrorCode::ServerError(INVALID_MRENCLAVE_CODE).into());
+	if request.shard().encode() != ctx.mrenclave.encode() {
+		return Err(ErrorCode::ServerError(INVALID_SHARD_CODE).into());
 	}
 
-	let Ok(encoded_auth_call) = request.decrypt(ctx.shielding_key.clone()) else {
+	let Ok(encoded_auth_call) = request.decrypt(Box::new(ctx.shielding_key.clone())) else {
 		return Err(ErrorCode::ServerError(REQUEST_DECRYPTION_FAILED_CODE).into());
 	};
 
-	let Ok(auth_call) = AuthenticatedCall::decode(&mut encoded_auth_call.as_slice()) else {
-		return Err(ErrorCode::ServerError(INVALID_AUTHENTICATED_CALL_CODE).into());
-	};
-
-	let authentication_result: Result<(), &str> = match auth_call.authentication {
-		// TODO:
-		Authentication::Email(ref _verification_code) => {
-			// Verify code
-			Ok(())
+	let auth_call: AuthenticatedCall =
+		match AuthenticatedCall::decode(&mut encoded_auth_call.as_slice()) {
+			Ok(auth_call) => auth_call,
+			Err(e) => {
+				log::error!("Failed to decode authenticated call: {:?}", e);
+				return Err(ErrorCode::ServerError(INVALID_AUTHENTICATED_CALL_CODE).into());
+			},
+		};
+
+	let authentication_result = match auth_call.authentication {
+		Authentication::Web3(ref signature) => verify_web3_authentication(
+			signature,
+			&auth_call.call,
+			auth_call.nonce,
+			&ctx.mrenclave,
+			&request.shard,
+		),
+		Authentication::Email(ref verification_code) => {
+			verify_email_authentication(ctx, auth_call.call.sender_identity(), verification_code)
+				.await
 		},
-		Authentication::AuthToken(ref _token) => {
-			// Verify token
-			Ok(())
+		Authentication::AuthToken(ref auth_token) => {
+			verify_auth_token_authentication(ctx, auth_call.call.sender_identity(), auth_token)
+				.await
 		},
 	};