duplicate rpc call with new name (#3425)

Author: BillyWooo

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

@@ -0,0 +1,70 @@
+use crate::{
+	error_code::*, methods::omni::PumpxRpcError, server::RpcContext, verify_auth::verify_auth,
+	Deserialize, ErrorCode,
+};
+use executor_core::native_task::*;
+use executor_primitives::OmniAuth;
+use heima_primitives::{Identity, Web2IdentityType};
+use jsonrpsee::RpcModule;
+use native_task_handler::NativeTaskOk;
+use pumpx::types::AddWalletResponse;
+use serde::Serialize;
+
+use super::common::{check_omni_api_response, handle_omni_native_task};
+
+#[derive(Debug, Deserialize)]
+pub struct AddWalletParams {
+	pub user_email: String,
+	pub auth_token: String,
+}
+
+#[derive(Serialize, Clone)]
+pub struct RPCAddWalletResponse {
+	pub backend_response: AddWalletResponse,
+}
+
+impl From<AddWalletParams> for NativeTaskWrapper<NativeTask> {
+	fn from(p: AddWalletParams) -> Self {
+		Self {
+			task: NativeTask::PumpxAddWallet(Identity::from_web2_account(
+				p.user_email.as_str(),
+				Web2IdentityType::Email,
+			)),
+			nonce: None,
+			auth: Some(OmniAuth::AuthToken(p.auth_token)),
+		}
+	}
+}
+
+pub fn register_add_wallet(module: &mut RpcModule<RpcContext>) {
+	module
+		.register_async_method("omni_addWallet", |params, ctx, _| async move {
+			let params = params.parse::<AddWalletParams>().map_err(|e| {
+				log::error!("Failed to parse params: {:?}", e);
+				PumpxRpcError::from_error_code(ErrorCode::ParseError)
+			})?;
+
+			log::debug!("Received omni_addWallet, user_email: {}", params.user_email);
+
+			let wrapper: NativeTaskWrapper<NativeTask> = params.into();
+
+			if wrapper.task.require_auth() && verify_auth(ctx.clone(), &wrapper).await.is_err() {
+				return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
+					AUTH_VERIFICATION_FAILED_CODE,
+				)));
+			}
+
+			handle_omni_native_task(&ctx, wrapper, |task_ok| match task_ok {
+				NativeTaskOk::PumpxAddWallet(response) => {
+					check_omni_api_response(response.clone(), "Add wallet".into())?;
+					Ok(RPCAddWalletResponse { backend_response: response })
+				},
+				_ => {
+					log::error!("Unexpected response type");
+					Err(PumpxRpcError::from_error_code(ErrorCode::InternalError))
+				},
+			})
+			.await
+		})
+		.expect("Failed to register omni_addWallet method");
+}

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

@@ -0,0 +1,121 @@
+use jsonrpsee::types::{ErrorCode, ErrorObjectOwned};
+use parity_scale_codec::Codec;
+use pumpx::types::ApiResponse;
+use serde::Serialize;
+
+use crate::{error_code::*, oneshot, server::RpcContext, Decode};
+use executor_core::native_task::*;
+use native_task_handler::{NativeTaskError, NativeTaskOk, NativeTaskResponse};
+
+#[derive(Serialize, Debug)]
+pub struct PumpxRpcError {
+	pub code: i32,
+	pub message: String,
+	#[serde(skip_serializing_if = "Option::is_none")]
+	pub data: Option<PumpxRpcErrorData>,
+}
+
+#[derive(Serialize, Debug)]
+pub struct PumpxRpcErrorData {
+	#[serde(skip_serializing_if = "Option::is_none")]
+	pub backend_response: Option<PumpxRpcErrorBackendResponse>,
+}
+
+#[derive(Serialize, Debug)]
+pub struct PumpxRpcErrorBackendResponse {
+	pub code: i32,
+	pub message: String,
+}
+
+impl PumpxRpcError {
+	pub fn from_code_and_message(code: i32, message: String) -> Self {
+		Self { code, message, data: None }
+	}
+
+	pub fn from_error_code(error_code: ErrorCode) -> Self {
+		Self { code: error_code.code(), message: error_code.message().to_string(), data: None }
+	}
+
+	pub fn from_api_response<T>(api_response: ApiResponse<T>) -> Self
+	where
+		T: Codec,
+	{
+		Self {
+			code: ErrorCode::InternalError.code(),
+			message: ErrorCode::InternalError.message().to_string(),
+			data: Some(PumpxRpcErrorData {
+				backend_response: Some(PumpxRpcErrorBackendResponse {
+					code: api_response.code as i32,
+					message: api_response.message,
+				}),
+			}),
+		}
+	}
+}
+
+impl From<PumpxRpcError> for ErrorObjectOwned {
+	fn from(error: PumpxRpcError) -> Self {
+		ErrorObjectOwned::owned(error.code, error.message, error.data)
+	}
+}
+
+/// Process native task and handle response
+pub async fn handle_omni_native_task<F, R>(
+	ctx: &RpcContext,
+	wrapper: NativeTaskWrapper<NativeTask>,
+	task_ok_handler: F,
+) -> Result<R, PumpxRpcError>
+where
+	F: FnOnce(NativeTaskOk) -> Result<R, PumpxRpcError>,
+{
+	// Create channel for response
+	let (response_sender, response_receiver) = oneshot::channel();
+
+	// Send task to executor
+	ctx.native_task_sender.send((wrapper, response_sender)).await.map_err(|_| {
+		log::error!("Failed to send request to native call executor");
+		PumpxRpcError::from_error_code(ErrorCode::InternalError)
+	})?;
+
+	// Receive response
+	let response = response_receiver.await.map_err(|e| {
+		log::error!("Failed to receive response from native call handler: {:?}", e);
+		PumpxRpcError::from_error_code(ErrorCode::InternalError)
+	})?;
+
+	// Decode response
+	let native_task_response: NativeTaskResponse = Decode::decode(&mut response.as_slice())
+		.map_err(|_| {
+			log::error!("Failed to decode native task response");
+			PumpxRpcError::from_error_code(ErrorCode::InternalError)
+		})?;
+
+	// Process response
+	match native_task_response {
+		Ok(task_ok) => task_ok_handler(task_ok),
+		Err(NativeTaskError::InternalError) => {
+			log::error!("Internal error in native task");
+			Err(PumpxRpcError::from_error_code(ErrorCode::InternalError))
+		},
+		Err(native_task_error) => {
+			log::error!("Failed to execute native task: {:?}", native_task_error);
+			Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(get_native_task_error_code(
+				&native_task_error,
+			))))
+		},
+	}
+}
+
+pub fn check_omni_api_response<T>(
+	response: ApiResponse<T>,
+	name: String,
+) -> Result<(), PumpxRpcError>
+where
+	T: Codec,
+{
+	if response.code != 10000 {
+		log::error!("{} failed: code={}, message={}", name, response.code, response.message);
+		return Err(PumpxRpcError::from_api_response(response));
+	}
+	Ok(())
+}

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

@@ -0,0 +1,95 @@
+use crate::{
+	error_code::*, methods::omni::PumpxRpcError, server::RpcContext, verify_auth::verify_auth,
+	Deserialize, ErrorCode,
+};
+use ethers::types::Bytes;
+use executor_core::native_task::*;
+use executor_crypto::aes256::{aes_encrypt_default, Aes256Key, SerdeAesOutput};
+use executor_primitives::OmniAuth;
+use heima_primitives::{Identity, Web2IdentityType};
+use jsonrpsee::RpcModule;
+use native_task_handler::NativeTaskOk;
+use rsa::Oaep;
+use sha2::Sha256;
+
+use super::common::handle_omni_native_task;
+
+#[derive(Debug, Deserialize)]
+pub struct ExportWalletParams {
+	pub user_email: String,
+	pub key: Bytes, // RSA-encrypted AES key to encrypt the wallet private key, in 0x-hex-string
+	pub google_code: String,
+	pub chain_id: PumpxChainId,
+	pub wallet_index: PumxWalletIndex,
+	pub wallet_address: String,
+	pub email_code: String,
+}
+
+impl From<ExportWalletParams> for NativeTaskWrapper<NativeTask> {
+	fn from(p: ExportWalletParams) -> Self {
+		Self {
+			task: NativeTask::PumpxExportWallet(
+				Identity::from_web2_account(p.user_email.as_str(), Web2IdentityType::Pumpx),
+				p.google_code,
+				p.chain_id,
+				p.wallet_index,
+				p.wallet_address,
+			),
+			nonce: None,
+			auth: Some(OmniAuth::Email(p.user_email, p.email_code)),
+		}
+	}
+}
+
+pub fn register_export_wallet(module: &mut RpcModule<RpcContext>) {
+	module
+		.register_async_method("omni_exportWallet", |params, ctx, _| async move {
+			let params = params.parse::<ExportWalletParams>().map_err(|e| {
+				log::error!("Failed to parse params: {:?}", e);
+				PumpxRpcError::from_error_code(ErrorCode::ParseError)
+			})?;
+
+			log::debug!("Received omni_exportWallet, user_email: {}, chain_id: {}, wallet_index: {}, expected_wallet_address: {}", params.user_email, params.chain_id, params.wallet_index, params.wallet_address);
+
+			let aes_key = ctx
+				.shielding_key
+				.private_key()
+				.decrypt(Oaep::new::<Sha256>(), &params.key)
+				.map_err(|e| {
+					log::error!("Failed to decrypt shielded value: {:?}", e);
+					PumpxRpcError::from_code_and_message(
+						DECRYPT_REQUEST_FAILED_CODE,
+						"Shielded value decryption failed".into(),
+					)
+				})?;
+			let aes_key: Aes256Key = aes_key.try_into().map_err(|_| {
+				log::error!("Failed to convert AesKey");
+				PumpxRpcError::from_code_and_message(
+					AES_KEY_CONVERT_FAILED_CODE,
+					"AesKey convert failed".into(),
+				)
+			})?;
+
+			let wrapper: NativeTaskWrapper<NativeTask> = params.into();
+
+			if wrapper.task.require_auth() && verify_auth(ctx.clone(), &wrapper).await.is_err() {
+				return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
+					AUTH_VERIFICATION_FAILED_CODE,
+				)));
+			}
+
+			handle_omni_native_task(&ctx, wrapper, |task_ok| match task_ok {
+				NativeTaskOk::PumpxExportWallet(wallet) => {
+					let encrypted_wallet: SerdeAesOutput =
+						aes_encrypt_default(&aes_key, &wallet).into();
+					Ok(encrypted_wallet)
+				},
+				_ => {
+					log::error!("Unexpected response type");
+					Err(PumpxRpcError::from_error_code(ErrorCode::InternalError))
+				},
+			})
+			.await
+		})
+		.expect("Failed to register omni_exportWallet method");
+}

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

@@ -1,6 +1,9 @@
 use crate::server::RpcContext;
 use jsonrpsee::RpcModule;
 
+mod common;
+use common::*;
+
 mod get_health;
 use get_health::*;
 
@@ -19,11 +22,40 @@ use request_email_verification_code::*;
 mod submit_native_task;
 use submit_native_task::*;
 
+mod add_wallet;
+use add_wallet::*;
+
+mod export_wallet;
+use export_wallet::*;
+
+mod notify_limit_order_result;
+use notify_limit_order_result::*;
+
+mod request_jwt;
+use request_jwt::*;
+
+mod sign_limit_order;
+use sign_limit_order::*;
+
+mod submit_swap_order;
+use submit_swap_order::*;
+
+mod transfer_widthdraw;
+use transfer_widthdraw::*;
+
 pub fn register_omni(module: &mut RpcModule<RpcContext>) {
 	register_get_health(module);
 	register_get_next_intent_id(module);
 	register_get_shielding_key(module);
 	register_submit_native_task(module);
 	register_request_email_verification_code(module);
 	register_get_oauth2_google_authorization_url(module);
+
+	register_request_jwt(module);
+	register_export_wallet(module);
+	register_add_wallet(module);
+	register_transfer_withdraw(module);
+	register_submit_swap_order(module);
+	register_sign_limit_order_params(module);
+	register_notify_limit_order_result(module);
 }