Clean up to_hex() usage (#3785)

* init

* update

* update oa type check

* remove unneeded

* fix tests

* fix tests

Author: Kailai-Wang

tee-worker/omni-executor/executor-primitives/src/auth.rs

@@ -3,7 +3,7 @@ use crate::{
 		BitcoinSignature, EthereumSignature, HeimaMultiSignature, SolanaSignature,
 		SubstrateSignature,
 	},
-	utils::hex::{decode_hex, ToHexPrefixed},
+	utils::hex::{decode_hex, hex_encode},
 	OmniAccountAuthType,
 };
 use base58::{FromBase58, ToBase58};
@@ -123,9 +123,9 @@ impl TryFrom<Identity> for UserId {
 			Identity::Github(handle) => {
 				Ok(UserId::Github(String::from_utf8(handle.inner.to_vec()).map_err(|_| ())?))
 			},
-			Identity::Substrate(address) => Ok(UserId::Substrate(address.to_hex())),
-			Identity::Evm(address) => Ok(UserId::Evm(address.to_hex())),
-			Identity::Bitcoin(address) => Ok(UserId::Bitcoin(address.to_hex())),
+			Identity::Substrate(address) => Ok(UserId::Substrate(hex_encode(&address.encode()))),
+			Identity::Evm(address) => Ok(UserId::Evm(hex_encode(&address.encode()))),
+			Identity::Bitcoin(address) => Ok(UserId::Bitcoin(hex_encode(&address.encode()))),
 			Identity::Solana(address) => Ok(UserId::Solana(address.as_ref().to_base58())),
 			Identity::Email(handle) => {
 				Ok(UserId::Email(String::from_utf8(handle.inner.to_vec()).map_err(|_| ())?))

tee-worker/omni-executor/executor-primitives/src/utils/hex.rs

@@ -15,51 +15,8 @@
 // along with Litentry.  If not, see <https://www.gnu.org/licenses/>.
 
 use hex::FromHexError;
-use parity_scale_codec::{Decode, Encode, Error as CodecError};
 use std::{string::String, vec::Vec};
 
-#[derive(Debug)]
-pub enum Error {
-	Hex(FromHexError),
-	Codec(CodecError),
-}
-
-impl std::fmt::Display for Error {
-	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-		match self {
-			Error::Hex(e) => write!(f, "Hex error: {}", e),
-			Error::Codec(e) => write!(f, "Codec error: {}", e),
-		}
-	}
-}
-
-/// Trait to encode a given value to a hex string, prefixed with "0x".
-pub trait ToHexPrefixed {
-	fn to_hex(&self) -> String;
-}
-
-impl<T: Encode> ToHexPrefixed for T {
-	fn to_hex(&self) -> String {
-		hex_encode(&self.encode())
-	}
-}
-
-/// Trait to decode a hex string to a given output.
-pub trait FromHexPrefixed {
-	type Output;
-
-	fn from_hex(msg: &str) -> Result<Self::Output, Error>;
-}
-
-impl<T: Decode> FromHexPrefixed for T {
-	type Output = T;
-
-	fn from_hex(msg: &str) -> Result<Self::Output, Error> {
-		let byte_array = decode_hex(msg).map_err(Error::Hex)?;
-		Decode::decode(&mut byte_array.as_slice()).map_err(Error::Codec)
-	}
-}
-
 /// Hex encodes given data and preappends a "0x".
 pub fn hex_encode(data: &[u8]) -> String {
 	let mut hex_str = hex::encode(data);
@@ -82,38 +39,24 @@ pub fn decode_hex<T: AsRef<[u8]>>(message: T) -> Result<Vec<u8>, FromHexError> {
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use parity_scale_codec::{Decode, Encode};
-	use std::string::ToString;
 
 	#[test]
 	fn hex_encode_decode_works() {
-		let data = "Hello World!".to_string();
+		let data = "Hello World!".as_bytes();
 
-		let hex_encoded_data = hex_encode(&data.encode());
-		let decoded_data =
-			String::decode(&mut decode_hex(hex_encoded_data).unwrap().as_slice()).unwrap();
+		let hex_encoded_data = hex_encode(data);
+		let decoded_data = decode_hex(&hex_encoded_data).unwrap();
 
-		assert_eq!(data, decoded_data);
+		assert_eq!(data, decoded_data.as_slice());
 	}
 
 	#[test]
 	fn hex_encode_decode_works_empty_input() {
-		let data = String::new();
-
-		let hex_encoded_data = hex_encode(&data.encode());
-		let decoded_data =
-			String::decode(&mut decode_hex(hex_encoded_data).unwrap().as_slice()).unwrap();
-
-		assert_eq!(data, decoded_data);
-	}
-
-	#[test]
-	fn to_hex_from_hex_works() {
-		let data = "Hello World!".to_string();
+		let data = "".as_bytes();
 
-		let hex_encoded_data = data.to_hex();
-		let decoded_data = String::from_hex(&hex_encoded_data).unwrap();
+		let hex_encoded_data = hex_encode(data);
+		let decoded_data = decode_hex(&hex_encoded_data).unwrap();
 
-		assert_eq!(data, decoded_data);
+		assert_eq!(data, decoded_data.as_slice());
 	}
 }

tee-worker/omni-executor/heima/authentication/src/auth_token.rs

@@ -111,7 +111,8 @@ mod tests {
 
 	use super::*;
 	use chrono::{Days, Utc};
-	use executor_primitives::{utils::hex::ToHexPrefixed, Identity, Web2IdentityType};
+	use executor_primitives::{utils::hex::hex_encode, Identity, Web2IdentityType};
+	use parity_scale_codec::Encode;
 	use rsa::{pkcs1::EncodeRsaPrivateKey, RsaPrivateKey};
 
 	#[derive(PartialEq, Debug, Serialize, Deserialize)]
@@ -142,7 +143,7 @@ mod tests {
 		let omni_account = uid.to_omni_account(CLIENT_ID_HEIMA);
 
 		let claims = AuthTokenClaims::new(
-			omni_account.to_hex(),
+			hex_encode(omni_account.as_ref()),
 			AUTH_TOKEN_ID_TYPE.to_string(),
 			CLIENT_ID_HEIMA.to_string(),
 			AuthOptions { expires_at },
@@ -166,7 +167,7 @@ mod tests {
 		let omni_account = uid.to_omni_account(CLIENT_ID_HEIMA);
 
 		let claims = AuthTokenClaims::new(
-			omni_account.to_hex(),
+			hex_encode(omni_account.as_ref()),
 			AUTH_TOKEN_ID_TYPE.to_string(),
 			CLIENT_ID_HEIMA.to_string(),
 			AuthOptions { expires_at: 100 },
@@ -242,7 +243,7 @@ mod tests {
 		let omni_account = email_identity.to_omni_account(CLIENT_ID_HEIMA);
 
 		let claims = AuthTokenClaims::new(
-			omni_account.to_hex(),
+			hex_encode(omni_account.as_ref()),
 			AUTH_TOKEN_ACCESS_TYPE.to_string(),
 			CLIENT_ID_HEIMA.to_string(),
 			AuthOptions { expires_at },

tee-worker/omni-executor/rpc-server/src/auth_utils.rs

@@ -6,7 +6,7 @@ use executor_core::types::SerializablePackedUserOperation;
 use executor_crypto::ecdsa;
 use executor_primitives::{
 	signature::{EthereumSignature, HeimaMultiSignature},
-	utils::hex::{decode_hex, FromHexPrefixed},
+	utils::hex::decode_hex,
 	ChainId,
 };
 use executor_storage::{Storage, WildmetaTimestampStorage};
@@ -34,8 +34,12 @@ pub fn verify_wildmeta_signature(
 		})?;
 	let heima_sig = HeimaMultiSignature::Ethereum(ethereum_signature);
 
-	let agent_address = Address20::from_hex(agent_address).map_err(|e| {
-		error!("Failed to parse agent address: {:?}", e);
+	let agent_address_bytes = decode_hex(agent_address).map_err(|e| {
+		error!("Failed to decode signature: {:?}", e);
+		ErrorObject::from(ErrorCode::ParseError)
+	})?;
+	let agent_address = Address20::try_from(agent_address_bytes.as_slice()).map_err(|_| {
+		error!("Failed to parse agent address");
 		ErrorObject::from(ErrorCode::ParseError)
 	})?;
 

tee-worker/omni-executor/rpc-server/src/error_code.rs

@@ -57,7 +57,6 @@ pub const INVALID_EMAIL_FORMAT_CODE: i32 = -32107;
 
 // Account/Identity Error Codes (-32120 to -32139)
 pub const ACCOUNT_PARSE_ERROR_CODE: i32 = -32121;
-pub const INVALID_ACCOUNT_LENGTH_CODE: i32 = -32124;
 
 // External Service Error Codes (-32160 to -32179)
 pub const SIGNER_SERVICE_ERROR_CODE: i32 = -32160;

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

@@ -4,12 +4,11 @@ use crate::{
 	error_code::*,
 	methods::omni::{common::check_auth, PumpxRpcError},
 	server::RpcContext,
+	utils::omni::to_omni_account,
 };
 use executor_core::intent_executor::IntentExecutor;
-use executor_primitives::{utils::hex::FromHexPrefixed, AccountId};
 use executor_storage::{HeimaJwtStorage, Storage};
 use heima_authentication::constants::AUTH_TOKEN_ACCESS_TYPE;
-use heima_primitives::Address32;
 use jsonrpsee::RpcModule;
 use pumpx::methods::add_wallet::AddWalletResponse;
 use serde::Serialize;
@@ -25,7 +24,7 @@ pub fn register_add_wallet<CrossChainIntentExecutor: IntentExecutor + Send + Syn
 ) {
 	module
 		.register_async_method("omni_addWallet", |_params, ctx, ext| async move {
-			let omni_account = check_auth(&ext).map_err(|e| {
+			let oa_str = check_auth(&ext).map_err(|e| {
 				error!("Authentication check failed: {:?}", e);
 				PumpxRpcError::from(
 					DetailedError::new(
@@ -38,18 +37,16 @@ pub fn register_add_wallet<CrossChainIntentExecutor: IntentExecutor + Send + Syn
 
 			debug!("Received omni_addWallet");
 
-			let Ok(address) = Address32::from_hex(&omni_account) else {
-				error!("Failed to parse from omni account token");
-				return Err(PumpxRpcError::from(
-					DetailedError::new(INTERNAL_ERROR_CODE, "Internal error")
-						.with_reason("Failed to parse omni account from authentication token"),
-				));
-			};
-			let omni_account_id = AccountId::from(address);
+			let omni_account = to_omni_account(&oa_str).map_err(|_| {
+				PumpxRpcError::from(DetailedError::new(
+					PARSE_ERROR_CODE,
+					"Failed to parse omni account",
+				))
+			})?;
 
 			// Inlined handler logic from handle_pumpx_add_wallet
 			let storage = HeimaJwtStorage::new(ctx.storage_db.clone());
-			let Ok(Some(access_token)) = storage.get(&(omni_account_id, AUTH_TOKEN_ACCESS_TYPE))
+			let Ok(Some(access_token)) = storage.get(&(omni_account, AUTH_TOKEN_ACCESS_TYPE))
 			else {
 				error!("Failed to get pumpx_{}_jwt_token", AUTH_TOKEN_ACCESS_TYPE);
 				return Err(PumpxRpcError::from(

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

@@ -16,18 +16,17 @@
 
 use super::common::PumpxRpcError;
 use crate::detailed_error::DetailedError;
+use crate::error_code::PARSE_ERROR_CODE;
 use crate::server::RpcContext;
 use crate::utils::gas_estimation::estimate_user_op_gas;
+use crate::utils::omni::to_omni_account;
 use crate::utils::user_op::convert_to_packed_user_op;
-use crate::validation_helpers::{
-	validate_ethereum_address, validate_omni_account_hex, validate_omni_account_length,
-};
+use crate::validation_helpers::validate_ethereum_address;
 use alloy::primitives::utils::format_units;
 use executor_core::intent_executor::IntentExecutor;
 use executor_core::types::SerializablePackedUserOperation;
-use executor_primitives::{AccountId, ChainId};
+use executor_primitives::ChainId;
 use jsonrpsee::RpcModule;
-use parity_scale_codec::Decode;
 use serde::{Deserialize, Serialize};
 use tracing::{debug, error, info};
 
@@ -108,14 +107,10 @@ pub fn register_estimate_user_op_gas<
 
 			debug!("Received omni_estimateUserOpGas, params: {:?}", params);
 
-			let account_bytes = validate_omni_account_hex(&params.omni_account, "omni_account")
-				.map_err(PumpxRpcError::from)?;
-			validate_omni_account_length(&account_bytes, "omni_account")
-				.map_err(PumpxRpcError::from)?;
-			let account_id = AccountId::decode(&mut &account_bytes[..]).map_err(|e| {
-				PumpxRpcError::from(DetailedError::account_parse_error(
-					&params.omni_account,
-					&e.to_string(),
+			let omni_account = to_omni_account(&params.omni_account).map_err(|_| {
+				PumpxRpcError::from(DetailedError::new(
+					PARSE_ERROR_CODE,
+					"Failed to parse omni account",
 				))
 			})?;
 
@@ -125,7 +120,7 @@ pub fn register_estimate_user_op_gas<
 			// Inlined handler logic from handle_estimate_user_op_gas
 			info!(
 				"Processing EstimateUserOpGas for account {:?}, wallet_index: {}, chain_id: {}",
-				account_id, params.wallet_index, params.chain_id
+				omni_account, params.wallet_index, params.chain_id
 			);
 
 			// Get EntryPoint client for this chain

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

@@ -3,17 +3,17 @@ use crate::{
 	error_code::{INTERNAL_ERROR_CODE, PARSE_ERROR_CODE, *},
 	methods::omni::{common::check_auth, PumpxRpcError},
 	server::RpcContext,
+	utils::omni::to_omni_account,
 	Deserialize,
 };
 use ::pumpx::signer_client::PumpxChainId as _;
 use ethers::types::Bytes;
 use executor_core::intent_executor::IntentExecutor;
 use executor_core::native_task::*;
 use executor_crypto::aes256::{aes_decrypt, aes_encrypt_default, Aes256Key, SerdeAesOutput};
-use executor_primitives::{utils::hex::FromHexPrefixed, AccountId, PumpxAccountProfile};
+use executor_primitives::PumpxAccountProfile;
 use executor_storage::{HeimaJwtStorage, PumpxProfileStorage, Storage};
 use heima_authentication::constants::AUTH_TOKEN_ACCESS_TYPE;
-use heima_primitives::Address32;
 use jsonrpsee::RpcModule;
 use rsa::Oaep;
 use sha2::Sha256;
@@ -34,7 +34,7 @@ pub fn register_export_wallet<CrossChainIntentExecutor: IntentExecutor + Send +
 ) {
 	module
 		.register_async_method("omni_exportWallet", |params, ctx, ext| async move {
-			let omni_account = check_auth(&ext).map_err(|e| {
+			let oa_str = check_auth(&ext).map_err(|e| {
 				error!("Authentication check failed: {:?}", e);
 				PumpxRpcError::from(DetailedError::new(
 					AUTH_VERIFICATION_FAILED_CODE,
@@ -52,14 +52,12 @@ pub fn register_export_wallet<CrossChainIntentExecutor: IntentExecutor + Send +
 
 			debug!("Received omni_exportWallet, chain_id: {}, wallet_index: {}, expected_wallet_address: {}", params.chain_id, params.wallet_index, params.wallet_address);
 
-			let Ok(address) = Address32::from_hex(&omni_account) else {
-				error!("Failed to parse from omni account token");
-				return Err(PumpxRpcError::from(DetailedError::new(
-					INTERNAL_ERROR_CODE,
-					"Internal error"
-				).with_reason("Failed to parse omni account from authentication token")));
-			};
-			let omni_account_id = AccountId::from(address);
+			let omni_account = to_omni_account(&oa_str).map_err(|_| {
+				PumpxRpcError::from(DetailedError::new(
+					PARSE_ERROR_CODE,
+					"Failed to parse omni account",
+				))
+			})?;
 
 			let aes_key = ctx
 				.shielding_key
@@ -82,7 +80,7 @@ pub fn register_export_wallet<CrossChainIntentExecutor: IntentExecutor + Send +
 
 			// Inlined handler logic from handle_pumpx_export_wallet
 			let storage = HeimaJwtStorage::new(ctx.storage_db.clone());
-			let Ok(Some(access_token)) = storage.get(&(omni_account_id.clone(), AUTH_TOKEN_ACCESS_TYPE))
+			let Ok(Some(access_token)) = storage.get(&(omni_account.clone(), AUTH_TOKEN_ACCESS_TYPE))
 			else {
 				error!("Failed to get pumpx_{}_jwt_token", AUTH_TOKEN_ACCESS_TYPE);
 				return Err(PumpxRpcError::from(
@@ -133,7 +131,7 @@ pub fn register_export_wallet<CrossChainIntentExecutor: IntentExecutor + Send +
 				.export_wallet(
 					chain,
 					params.wallet_index,
-					omni_account_id.clone().into(),
+					omni_account.clone().into(),
 					ctx.aes256_key.to_vec(),
 					params.wallet_address,
 				)
@@ -157,14 +155,14 @@ pub fn register_export_wallet<CrossChainIntentExecutor: IntentExecutor + Send +
 			};
 
 			let omni_account_profile_storage = PumpxProfileStorage::new(ctx.storage_db.clone());
-			if let Ok(maybe_profile) = omni_account_profile_storage.get(&omni_account_id) {
+			if let Ok(maybe_profile) = omni_account_profile_storage.get(&omni_account) {
 				let profile = maybe_profile
 					.map(|mut p| {
 						p.wallet_exported = true;
 						p
 					})
 					.unwrap_or_else(|| PumpxAccountProfile { wallet_exported: true });
-				if let Err(e) = omni_account_profile_storage.insert(&omni_account_id, profile) {
+				if let Err(e) = omni_account_profile_storage.insert(&omni_account, profile) {
 					error!("Failed to update pumpx account profile: {:?}", e);
 					return Err(PumpxRpcError::from(
 						DetailedError::new(INTERNAL_ERROR_CODE, "Internal error")

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

@@ -18,7 +18,7 @@ use crate::detailed_error::DetailedError;
 use crate::error_code::PARSE_ERROR_CODE;
 use crate::{methods::omni::common::PumpxRpcError, server::RpcContext};
 use executor_core::intent_executor::IntentExecutor;
-use executor_primitives::{utils::hex::ToHexPrefixed, Web2IdentityType};
+use executor_primitives::{utils::hex::hex_encode, Web2IdentityType};
 use heima_primitives::Identity;
 use jsonrpsee::{types::ErrorObject, RpcModule};
 use serde::{Deserialize, Serialize};
@@ -49,7 +49,7 @@ pub fn register_get_omni_account<
 			let account =
 				Identity::from_web2_account(params.user_email.as_str(), Web2IdentityType::Email)
 					.to_omni_account(&params.client_id);
-			Ok::<String, ErrorObject>(account.to_hex())
+			Ok::<String, ErrorObject>(hex_encode(account.as_ref()))
 		})
 		.expect("Failed to register omni_getOmniAccount method");
 }