Skip to content

Commit d543ed9

Browse files
silva-fjKailai-Wang
silva-fj
and
Kailai-Wang
authored
Remove the need to include sender for jwt authentication (#3476)
* feat: remove sender requirement from JWT authentication - Simplified OmniAuth::AuthToken enum variant to only accept JWT token - Removed OmniAccount parameter from AuthToken constructor calls - Updated all RPC method implementations to use simplified AuthToken - Removed unused OmniAccount type definition - Cleaned up imports removing ToHexPrefixed where no longer needed This change streamlines JWT authentication by extracting sender information directly from the JWT token rather than requiring it as a separate parameter, reducing potential inconsistencies and simplifying the API surface. * feat: remove subject validation from JWT auth tokens This change streamlines JWT validation by removing the redundant subject check, as the subject information is already embedded within the token claims and doesn't need separate validation. * feat: extract account ID from JWT claims instead of parameters Remove the need to pass omni_account as a separate parameter to JWT authentication. The account ID is now extracted directly from the JWT token's subject claim, improving security and reducing parameter coupling. * refactor: centralize JWT auth and extract account from token Replace duplicated JWT validation code with centralized verify_auth_token_authentication function. Account ID is now extracted directly from JWT token claims instead of being passed as separate parameter, improving security and reducing code duplication. * updating ts definitions --------- Co-authored-by: Kai <7630809+Kailai-Wang@users.noreply.github.com>
1 parent 4010d57 commit d543ed9

File tree

12 files changed

+134
-211
lines changed

12 files changed

+134
-211
lines changed

‎tee-worker/client-api/src/omni/interfaces/omniExecutor/definitions.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ export default {
4444
_enum: {
4545
Web3: "(Identity, HeimaMultiSignature)",
4646
Email: "(Text, Text)",
47-
AuthToken: "(Text, Text)",
47+
AuthToken: "(Text)",
4848
OAuth2: "(Identity, OAuth2Data)",
4949
},
5050
},

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ use serde::{Deserialize, Serialize};
66

77
pub type VerificationCode = String;
88
type Email = String;
9-
type OmniAccount = String;
109
type JwtToken = String;
1110

1211
/// A serializable representation of Identity for JSON interchange.
@@ -89,15 +88,15 @@ impl TryFrom<IdentitySerde> for Identity {
8988
pub enum OmniAuth {
9089
Web3(Identity, HeimaMultiSignature), // (Signer, Signature)
9190
Email(Email, VerificationCode),
92-
AuthToken(OmniAccount, JwtToken),
91+
AuthToken(JwtToken),
9392
OAuth2(Identity, OAuth2Data), // (Sender, OAuth2Data)
9493
}
9594

9695
#[derive(Deserialize)]
9796
pub enum OmniAuthSerde {
9897
Web3(IdentitySerde, HeimaMultiSignature),
9998
Email(Email, VerificationCode),
100-
AuthToken(OmniAccount, JwtToken),
99+
AuthToken(JwtToken),
101100
OAuth2(IdentitySerde, OAuth2Data),
102101
}
103102

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

Lines changed: 7 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -37,20 +37,16 @@ impl AuthTokenClaims {
3737
}
3838

3939
pub struct Validation {
40-
pub sub: String,
4140
pub typ: String,
4241
pub skip_exp_check: bool,
4342
}
4443

4544
impl Validation {
46-
pub fn new(sub: String, typ: String, skip_exp_check: bool) -> Self {
47-
Self { sub, typ, skip_exp_check }
45+
pub fn new(typ: String, skip_exp_check: bool) -> Self {
46+
Self { typ, skip_exp_check }
4847
}
4948

5049
pub fn validate(&self, claims: &AuthTokenClaims) -> Result<(), Error> {
51-
if self.sub != claims.sub {
52-
return Err(Error::JwtError(jwt::ErrorKind::InvalidSubject));
53-
}
5450
if self.typ != claims.typ {
5551
return Err(Error::JwtError(jwt::ErrorKind::InvalidToken));
5652
}
@@ -144,8 +140,7 @@ mod tests {
144140
);
145141
let token = jwt::create(&claims, private_key.as_bytes()).unwrap();
146142

147-
let validation =
148-
Validation::new(omni_account.to_hex(), AUTH_TOKEN_ID_TYPE.to_string(), false);
143+
let validation = Validation::new(AUTH_TOKEN_ID_TYPE.to_string(), false);
149144
let result = token.validate(private_key.as_bytes(), validation);
150145

151146
assert_eq!(result, Ok(claims));
@@ -168,8 +163,7 @@ mod tests {
168163
);
169164
let token = jwt::create(&claims, private_key.as_bytes()).unwrap();
170165

171-
let validation =
172-
Validation::new(omni_account.to_hex(), AUTH_TOKEN_ID_TYPE.to_string(), false);
166+
let validation = Validation::new(AUTH_TOKEN_ID_TYPE.to_string(), false);
173167
let result = token.validate(private_key.as_bytes(), validation);
174168

175169
assert_eq!(result, Err(Error::JwtError(jwt::ErrorKind::ExpiredSignature)));
@@ -191,41 +185,14 @@ mod tests {
191185
TestAuthTokenNoSubClaims { typ: AUTH_TOKEN_ID_TYPE.to_string(), exp: expires_at };
192186
let token = jwt::create(&claims, private_key.as_bytes()).unwrap();
193187

194-
let validation =
195-
Validation::new("test-subject".to_string(), AUTH_TOKEN_ID_TYPE.to_string(), false);
188+
let validation = Validation::new(AUTH_TOKEN_ID_TYPE.to_string(), false);
196189
let result = token.validate(private_key.as_bytes(), validation);
197190
let Err(Error::JwtError(jwt::ErrorKind::Json(e))) = result else {
198191
panic!("Expected JsonError, got {:?}", result);
199192
};
200193
assert!(e.to_string().contains("missing field `sub`"));
201194
}
202195

203-
#[test]
204-
fn test_auth_token_invalid_subject() {
205-
let mut rng = rand::thread_rng();
206-
let rsa_private_key =
207-
RsaPrivateKey::new(&mut rng, 2048).expect("Failed to generate private key");
208-
let private_key = rsa_private_key.to_pkcs1_der().unwrap();
209-
210-
let expires_at = Utc::now()
211-
.checked_add_days(Days::new(1))
212-
.expect("Failed to calculate expiration")
213-
.timestamp();
214-
215-
let claims = AuthTokenClaims {
216-
sub: "test-sub".to_string(),
217-
typ: AUTH_TOKEN_ID_TYPE.to_string(),
218-
exp: expires_at,
219-
};
220-
let token = jwt::create(&claims, private_key.as_bytes()).unwrap();
221-
222-
let validation =
223-
Validation::new("invalid-sub".to_string(), AUTH_TOKEN_ID_TYPE.to_string(), false);
224-
let result = token.validate(private_key.as_bytes(), validation);
225-
226-
assert_eq!(result, Err(Error::JwtError(jwt::ErrorKind::InvalidSubject)));
227-
}
228-
229196
#[test]
230197
fn test_auth_token_missing_typ() {
231198
let mut rng = rand::thread_rng();
@@ -241,8 +208,7 @@ mod tests {
241208
let claims = TestAuthTokenNoTypClaims { sub: "test-sub".to_string(), exp: expires_at };
242209
let token = jwt::create(&claims, private_key.as_bytes()).unwrap();
243210

244-
let validation =
245-
Validation::new("test-sub".to_string(), AUTH_TOKEN_ID_TYPE.to_string(), false);
211+
let validation = Validation::new(AUTH_TOKEN_ID_TYPE.to_string(), false);
246212
let result = token.validate(private_key.as_bytes(), validation);
247213
let Err(Error::JwtError(jwt::ErrorKind::Json(e))) = result else {
248214
panic!("Expected JsonError, got {:?}", result);
@@ -272,8 +238,7 @@ mod tests {
272238
);
273239
let token = jwt::create(&claims, private_key.as_bytes()).unwrap();
274240

275-
let validation =
276-
Validation::new(omni_account.to_hex(), AUTH_TOKEN_ID_TYPE.to_string(), false);
241+
let validation = Validation::new(AUTH_TOKEN_ID_TYPE.to_string(), false);
277242
let result = token.validate(private_key.as_bytes(), validation);
278243

279244
assert_eq!(result, Err(Error::JwtError(jwt::ErrorKind::InvalidToken)));

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::{
33
Deserialize, ErrorCode,
44
};
55
use executor_core::native_task::*;
6-
use executor_primitives::{utils::hex::ToHexPrefixed, OmniAuth};
6+
use executor_primitives::OmniAuth;
77
use heima_primitives::{Identity, Web2IdentityType};
88
use jsonrpsee::RpcModule;
99
use native_task_handler::NativeTaskOk;
@@ -30,7 +30,7 @@ impl From<AddWalletParams> for NativeTaskWrapper<NativeTask> {
3030
NativeTaskWrapper::new(
3131
NativeTask::PumpxAddWallet(sender.clone()),
3232
None,
33-
Some(OmniAuth::AuthToken(sender.to_omni_account().to_hex(), p.auth_token)),
33+
Some(OmniAuth::AuthToken(p.auth_token)),
3434
)
3535
}
3636
}

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

Lines changed: 16 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
11
use crate::methods::omni::PumpxRpcError;
2+
use crate::verify_auth::verify_auth_token_authentication;
23
use crate::{error_code::*, server::RpcContext, Deserialize, ErrorCode};
34
use executor_core::native_task::*;
4-
use executor_crypto::jwt;
55
use executor_primitives::{utils::hex::FromHexPrefixed, OmniAuth};
6-
use heima_authentication::auth_token::AuthTokenClaims;
6+
use heima_authentication::auth_token::AUTH_TOKEN_ACCESS_TYPE;
77
use heima_primitives::{Address32, Identity};
88
use jsonrpsee::RpcModule;
99
use native_task_handler::NativeTaskOk;
10-
use rsa::pkcs1::DecodeRsaPrivateKey;
11-
use rsa::pkcs1::EncodeRsaPublicKey;
12-
use rsa::RsaPrivateKey;
1310
use tracing::{debug, error};
1411

1512
use super::common::handle_omni_native_task;
@@ -35,32 +32,21 @@ pub fn register_notify_limit_order_result(module: &mut RpcModule<RpcContext>) {
3532
params.intent_id, params.result, params.message
3633
);
3734

38-
let private_key =
39-
RsaPrivateKey::from_pkcs1_der(&ctx.jwt_rsa_private_key).map_err(|e| {
40-
error!("Failed to parse private key: {:?}", e);
41-
PumpxRpcError::from_error_code(ErrorCode::InternalError)
42-
})?;
43-
44-
let public_key = private_key.to_public_key().to_pkcs1_der().map_err(|e| {
45-
error!("Failed to generate public key: {:?}", e);
46-
PumpxRpcError::from_error_code(ErrorCode::InternalError)
47-
})?;
48-
49-
// this validates jwt - we skip exp check for this call
50-
let Ok(token) =
51-
jwt::decode::<AuthTokenClaims>(&params.auth_token, public_key.as_bytes(), true)
52-
else {
53-
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
54-
AUTH_VERIFICATION_FAILED_CODE,
55-
)));
35+
let omni_account = match verify_auth_token_authentication(
36+
ctx.clone(),
37+
&params.auth_token,
38+
AUTH_TOKEN_ACCESS_TYPE,
39+
true,
40+
) {
41+
Ok(claims) => claims.sub,
42+
Err(_) => {
43+
error!("Failed to verify auth token");
44+
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
45+
AUTH_VERIFICATION_FAILED_CODE,
46+
)));
47+
},
5648
};
57-
if token.typ != "access" {
58-
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
59-
AUTH_VERIFICATION_FAILED_CODE,
60-
)));
61-
}
6249

63-
let omni_account = token.sub;
6450
let Ok(address) = Address32::from_hex(&omni_account) else {
6551
error!("Failed to parse from omni account token");
6652
return Err(PumpxRpcError::from_error_code(ErrorCode::InternalError));
@@ -74,7 +60,7 @@ pub fn register_notify_limit_order_result(module: &mut RpcModule<RpcContext>) {
7460
params.message,
7561
),
7662
None,
77-
Some(OmniAuth::AuthToken(omni_account, params.auth_token)),
63+
Some(OmniAuth::AuthToken(params.auth_token)),
7864
);
7965

8066
handle_omni_native_task(&ctx, wrapper, |task_ok| match task_ok {

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

Lines changed: 19 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -14,28 +14,24 @@
1414
// You should have received a copy of the GNU General Public License
1515
// along with Litentry. If not, see <https://www.gnu.org/licenses/>.
1616

17-
use executor_primitives::utils::hex::FromHexPrefixed;
18-
use heima_primitives::Address32;
19-
use heima_primitives::IntentId;
20-
use rsa::pkcs1::DecodeRsaPrivateKey;
21-
use rsa::pkcs1::EncodeRsaPublicKey;
22-
2317
use crate::error_code::AUTH_VERIFICATION_FAILED_CODE;
2418
use crate::methods::omni::PumpxRpcError;
2519
use crate::server::RpcContext;
20+
use crate::verify_auth::verify_auth_token_authentication;
2621
use crate::ErrorCode;
2722
use ethers::types::Bytes;
2823
use executor_core::native_task::NativeTask;
2924
use executor_core::native_task::NativeTaskWrapper;
3025
use executor_core::native_task::PumpxChainId;
3126
use executor_core::native_task::PumxWalletIndex;
32-
use executor_crypto::jwt;
27+
use executor_primitives::utils::hex::FromHexPrefixed;
3328
use executor_primitives::OmniAuth;
34-
use heima_authentication::auth_token::AuthTokenClaims;
29+
use heima_authentication::auth_token::AUTH_TOKEN_ACCESS_TYPE;
30+
use heima_primitives::Address32;
3531
use heima_primitives::Identity;
32+
use heima_primitives::IntentId;
3633
use jsonrpsee::RpcModule;
3734
use native_task_handler::NativeTaskOk;
38-
use rsa::RsaPrivateKey;
3935
use serde::Deserialize;
4036
use serde::Serialize;
4137
use tracing::{debug, error};
@@ -70,32 +66,21 @@ pub fn register_sign_limit_order_params(module: &mut RpcModule<RpcContext>) {
7066

7167
debug!("Received omni_signLimitOrder, intent_id: {}, order_id: {}, chain_id: {}, wallet_index: {}", params.intent_id, params.order_id, params.chain_id, params.wallet_index);
7268

73-
let private_key =
74-
RsaPrivateKey::from_pkcs1_der(&ctx.jwt_rsa_private_key).map_err(|e| {
75-
error!("Failed to parse private key: {:?}", e);
76-
PumpxRpcError::from_error_code(ErrorCode::InternalError)
77-
})?;
78-
79-
let public_key = private_key.to_public_key().to_pkcs1_der().map_err(|e| {
80-
error!("Failed to generate public key: {:?}", e);
81-
PumpxRpcError::from_error_code(ErrorCode::InternalError)
82-
})?;
83-
84-
// this validates jwt - we skip exp check for this call
85-
let Ok(token) =
86-
jwt::decode::<AuthTokenClaims>(&params.auth_token, public_key.as_bytes(), true)
87-
else {
88-
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
89-
AUTH_VERIFICATION_FAILED_CODE,
90-
)));
69+
let omni_account = match verify_auth_token_authentication(
70+
ctx.clone(),
71+
&params.auth_token,
72+
AUTH_TOKEN_ACCESS_TYPE,
73+
true,
74+
) {
75+
Ok(claims) => claims.sub,
76+
Err(_) => {
77+
error!("Failed to verify auth token");
78+
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
79+
AUTH_VERIFICATION_FAILED_CODE,
80+
)));
81+
},
9182
};
92-
if token.typ != "access" {
93-
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
94-
AUTH_VERIFICATION_FAILED_CODE,
95-
)));
96-
}
9783

98-
let omni_account = token.sub;
9984
let Ok(address) = Address32::from_hex(&omni_account) else {
10085
error!("Failed to parse from omni account token");
10186
return Err(PumpxRpcError::from_error_code(ErrorCode::InternalError));
@@ -109,7 +94,7 @@ pub fn register_sign_limit_order_params(module: &mut RpcModule<RpcContext>) {
10994
params.unsigned_tx.iter().map(|tx| tx.to_vec()).collect(),
11095
),
11196
None,
112-
Some(OmniAuth::AuthToken(omni_account, params.auth_token)),
97+
Some(OmniAuth::AuthToken(params.auth_token)),
11398
);
11499

115100
handle_omni_native_task(&ctx, wrapper, |task_ok| match task_ok {

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

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@ use crate::{
44
Deserialize, ErrorCode,
55
};
66
use executor_core::native_task::*;
7-
use executor_primitives::{utils::hex::ToHexPrefixed, OmniAuth};
7+
use executor_primitives::OmniAuth;
88
use executor_storage::{PumpxJwtStorage, Storage};
99
use heima_authentication::auth_token::{AUTH_TOKEN_ACCESS_TYPE, AUTH_TOKEN_ID_TYPE};
1010
use heima_hex_utils::decode_hex;
1111
use heima_primitives::{
12-
Address20, Address32, BinanceConfig, BoundedVec, ChainAsset, CrossChainSwapProvider,
12+
AccountId, Address20, Address32, BinanceConfig, BoundedVec, ChainAsset, CrossChainSwapProvider,
1313
EthereumToken, Identity, Intent, PumpxConfig, PumpxOrderType, SingleChainSwapProvider,
1414
SolanaToken, SwapOrder, Web2IdentityType,
1515
};
@@ -19,6 +19,7 @@ use pumpx::constants::*;
1919
use pumpx::methods::common::{OrderInfoResponse, SwapType};
2020
use pumpx::methods::send_order_tx::SendOrderTxResponse;
2121
use serde::Serialize;
22+
use std::str::FromStr;
2223
use tracing::{debug, error};
2324

2425
use super::common::{check_omni_api_response, handle_omni_native_task};
@@ -116,16 +117,20 @@ pub fn register_submit_swap_order(module: &mut RpcModule<RpcContext>) {
116117
debug!("Received omni_submitSwapOrder, user_email: {}, intent_id: {}, order_type: {:?}, swap_type: {:?}, from_chain_id: {}, from_token_ca: {:?}, from_amount: {}, to_chain_id: {}, to_token_ca: {:?}, wallet_index: {}",
117118
params.user_email, params.intent_id, params.order_type, params.swap_type, params.from_chain_id, params.from_token_ca, params.from_amount, params.to_chain_id, params.to_token_ca, params.wallet_index);
118119

119-
let user_identity =
120-
Identity::from_web2_account(&params.user_email, Web2IdentityType::Email);
121-
if verify_auth_token_authentication(ctx.clone(), user_identity.to_omni_account().to_hex(), &params.auth_token, AUTH_TOKEN_ID_TYPE, false)
122-
.is_err()
123-
{
124-
error!("Failed to verify auth token");
125-
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
126-
AUTH_VERIFICATION_FAILED_CODE,
127-
)));
128-
}
120+
let omni_account = match verify_auth_token_authentication(ctx.clone(), &params.auth_token, AUTH_TOKEN_ID_TYPE, false) {
121+
Ok(claims) => {
122+
AccountId::from_str(&claims.sub).map_err(|_| {
123+
error!("Failed to parse account id from auth token");
124+
PumpxRpcError::from_error_code(ErrorCode::InternalError)
125+
})?
126+
},
127+
Err(_) => {
128+
error!("Failed to verify auth token");
129+
return Err(PumpxRpcError::from_error_code(ErrorCode::ServerError(
130+
AUTH_VERIFICATION_FAILED_CODE,
131+
)));
132+
}
133+
};
129134

130135
let from_chain_asset = params.try_get_from_chain_asset().map_err(|_| {
131136
error!("Failed to get from chain asset");
@@ -151,7 +156,7 @@ pub fn register_submit_swap_order(module: &mut RpcModule<RpcContext>) {
151156

152157
let storage = PumpxJwtStorage::new(ctx.storage_db.clone());
153158
let Ok(Some(access_token)) =
154-
storage.get(&(user_identity.to_omni_account(), AUTH_TOKEN_ACCESS_TYPE))
159+
storage.get(&(omni_account, AUTH_TOKEN_ACCESS_TYPE))
155160
else {
156161
error!("Failed to get access token from storage");
157162
return Err(PumpxRpcError::from_error_code(ErrorCode::InternalError));
@@ -225,10 +230,12 @@ pub fn register_submit_swap_order(module: &mut RpcModule<RpcContext>) {
225230
error!("Failed to convert single chain swap provider");
226231
PumpxRpcError::from_error_code(ErrorCode::InternalError)
227232
})?);
233+
let user_identity =
234+
Identity::from_web2_account(&params.user_email, Web2IdentityType::Pumpx);
228235
let wrapper = NativeTaskWrapper::new(
229-
NativeTask::RequestIntent(user_identity.clone(), params.intent_id, intent),
236+
NativeTask::RequestIntent(user_identity, params.intent_id, intent),
230237
None,
231-
Some(OmniAuth::AuthToken(user_identity.to_omni_account().to_hex(), params.auth_token)),
238+
Some(OmniAuth::AuthToken(params.auth_token)),
232239
);
233240

234241
handle_omni_native_task(&ctx, wrapper, |task_ok| match task_ok {

0 commit comments

Comments
 (0)