Skip to content

Commit 139594e

Browse files
silva-fjsilva-fj
authored
Make client_id Dynamic (#3508)
1 parent a11b602 commit 139594e

File tree

21 files changed

+160
-135
lines changed

21 files changed

+160
-135
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ export default {
1111
task: "NativeTask",
1212
nonce: "Option<Nonce>",
1313
auth: "Option<OmniAuth>",
14+
client_id: "String",
1415
},
1516
MrEnclave: "H256",
1617
NativeTask: {

‎tee-worker/omni-executor/client-sdk/packages/client-sdk/src/lib/type-creators/raw-task.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { createOmniAuth, OmniAuthData } from './omni-auth';
1717
* @param data.nonce - Optional nonce value
1818
* @param data.authData - Optional authentication data
1919
* @param data.plain - Flag to return unencrypted raw task
20+
* @param data.clientId - Client identifier for the task
2021
* @param {Enclave} enclaveInstance - The enclave instance use to interact with Enclave.
2122
* @returns Promise resolving to RawTask
2223
*/
@@ -27,6 +28,7 @@ export async function createRawTaskType(
2728
nonce?: Index;
2829
authData?: OmniAuthData;
2930
plain: true;
31+
clientId: string;
3032
},
3133
enclaveInstance?: Enclave,
3234
): Promise<RawTask>;
@@ -50,6 +52,7 @@ export async function createRawTaskType(
5052
nonce?: Index;
5153
authData?: OmniAuthData;
5254
plain?: false;
55+
clientId: string;
5356
},
5457
enclaveInstance?: Enclave,
5558
): Promise<{ rawTask: RawTask, encryptionKey: CryptoKey }>;
@@ -90,16 +93,18 @@ export async function createRawTaskType(
9093
nonce?: Index;
9194
authData?: OmniAuthData;
9295
plain?: boolean;
96+
clientId: string;
9397
},
9498
enclaveInstance: Enclave = enclave,
9599
): Promise<RawTask | { rawTask: RawTask, encryptionKey: CryptoKey }> {
96-
const { authData, task, nonce, plain = false } = data;
100+
const { authData, task, nonce, plain = false, clientId } = data;
97101

98102
const auth = authData ? createOmniAuth(api.registry, authData) : undefined;
99103
const wrappedTask = api.createType<NativeTaskWrapper>('NativeTaskWrapper', {
100104
task,
101105
nonce: api.createType('Option<Nonce>', nonce),
102106
auth: api.createType('Option<OmniAuth>', auth),
107+
client_id: clientId,
103108
});
104109

105110
if (plain) {

‎tee-worker/omni-executor/executor-core/src/native_task.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,13 @@ pub struct NativeTaskWrapper<T: NativeTaskTrait> {
2424
pub task: T,
2525
pub nonce: Option<Nonce>,
2626
pub auth: Option<OmniAuth>,
27+
pub client_id: String,
2728
}
2829

2930
impl<T: NativeTaskTrait + Debug> NativeTaskWrapper<T> {
30-
pub fn new(task: T, nonce: Option<Nonce>, auth: Option<OmniAuth>) -> Self {
31+
pub fn new(task: T, nonce: Option<Nonce>, auth: Option<OmniAuth>, client_id: String) -> Self {
3132
let id: String = Uuid::new_v4().into();
32-
Self { id, task, nonce, auth }
33+
Self { id, task, nonce, auth, client_id }
3334
}
3435
}
3536

‎tee-worker/omni-executor/native-task-handler/src/lib.rs

Lines changed: 21 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,7 @@ use executor_storage::{
2020
};
2121
use heima_authentication::{
2222
auth_token::*,
23-
constants::{
24-
AUTH_TOKEN_ACCESS_TYPE, AUTH_TOKEN_EXPIRATION_DAYS, AUTH_TOKEN_ID_TYPE, CLIENT_ID_HEIMA,
25-
CLIENT_ID_PUMPX,
26-
},
23+
constants::{AUTH_TOKEN_ACCESS_TYPE, AUTH_TOKEN_EXPIRATION_DAYS, AUTH_TOKEN_ID_TYPE},
2724
};
2825
use heima_identity_verification::{get_verification_message, web2, web3};
2926
use parentchain_api_interface::runtime_types::{
@@ -206,6 +203,7 @@ async fn handle_native_task<
206203
};
207204

208205
let auth_type: Option<OmniAccountAuthType> = wrapper.auth.map(|t| t.into());
206+
let client_id = &wrapper.client_id;
209207

210208
let (response_sender, tx) = match wrapper.task {
211209
NativeTask::RequestAuthToken(sender) => {
@@ -236,14 +234,14 @@ async fn handle_native_task<
236234
AuthTokenClaims::new(
237235
email.to_string(),
238236
AUTH_TOKEN_ID_TYPE.to_string(),
239-
CLIENT_ID_HEIMA.to_string(),
237+
client_id.to_string(),
240238
auth_options,
241239
)
242240
},
243241
_ => AuthTokenClaims::new(
244242
sender.hash().to_string(),
245243
AUTH_TOKEN_ID_TYPE.to_string(),
246-
CLIENT_ID_HEIMA.to_string(),
244+
client_id.to_string(),
247245
auth_options,
248246
),
249247
};
@@ -275,8 +273,7 @@ async fn handle_native_task<
275273
return;
276274
},
277275
NativeTask::RequestIntent(sender, intent_id, intent) => {
278-
// TODO: fix this as part of P-1560
279-
let omni_account = sender.to_omni_account_with_client_id(CLIENT_ID_HEIMA);
276+
let omni_account = sender.to_omni_account_with_client_id(client_id);
280277

281278
debug!("Intent requested");
282279

@@ -622,12 +619,12 @@ async fn handle_native_task<
622619

623620
debug!("get_account_user_id ok, email: {}, user_id: {}", email, user_id);
624621
let omni_account = Identity::from_web2_account(&user_id, Web2IdentityType::Pumpx)
625-
.to_omni_account_with_client_id(CLIENT_ID_PUMPX);
622+
.to_omni_account_with_client_id(client_id);
626623

627624
let access_token_claims = AuthTokenClaims::new(
628625
omni_account.to_hex(),
629626
AUTH_TOKEN_ACCESS_TYPE.to_string(),
630-
CLIENT_ID_PUMPX.to_string(),
627+
client_id.to_string(),
631628
auth_options.clone(),
632629
);
633630
let Ok(access_token) = jwt::create(&access_token_claims, &ctx.jwt_rsa_private_key)
@@ -675,7 +672,7 @@ async fn handle_native_task<
675672
let id_token_claims = AuthTokenClaims::new(
676673
omni_account.to_hex(),
677674
AUTH_TOKEN_ID_TYPE.to_string(),
678-
CLIENT_ID_PUMPX.to_string(),
675+
client_id.to_string(),
679676
auth_options,
680677
);
681678
let Ok(id_token) = jwt::create(&id_token_claims, &ctx.jwt_rsa_private_key) else {
@@ -713,10 +710,9 @@ async fn handle_native_task<
713710
expected_wallet_address,
714711
) => {
715712
let storage = HeimaJwtStorage::new(ctx.storage_db.clone());
716-
let Ok(Some(access_token)) = storage.get(&(
717-
sender.to_omni_account_with_client_id(CLIENT_ID_PUMPX),
718-
AUTH_TOKEN_ACCESS_TYPE,
719-
)) else {
713+
let Ok(Some(access_token)) = storage
714+
.get(&(sender.to_omni_account_with_client_id(client_id), AUTH_TOKEN_ACCESS_TYPE))
715+
else {
720716
send_error(
721717
format!("Failed to get pumpx_{}_jwt_token", AUTH_TOKEN_ACCESS_TYPE),
722718
response_sender,
@@ -755,7 +751,7 @@ async fn handle_native_task<
755751
.export_wallet(
756752
chain,
757753
pumpx_wallet_index,
758-
sender.to_omni_account_with_client_id(CLIENT_ID_PUMPX).into(),
754+
sender.to_omni_account_with_client_id(client_id).into(),
759755
// TODO: theoretically we could pass the aes_key from initial RPC to signer, so that
760756
// we don't have to do double encryption/decryption
761757
ctx.aes256_key.to_vec(),
@@ -780,8 +776,8 @@ async fn handle_native_task<
780776
};
781777

782778
let omni_account_profile_storage = PumpxProfileStorage::new(ctx.storage_db.clone());
783-
if let Ok(maybe_profile) = omni_account_profile_storage
784-
.get(&sender.to_omni_account_with_client_id(CLIENT_ID_PUMPX))
779+
if let Ok(maybe_profile) =
780+
omni_account_profile_storage.get(&sender.to_omni_account_with_client_id(client_id))
785781
{
786782
let profile = maybe_profile
787783
.map(|mut p| {
@@ -790,7 +786,7 @@ async fn handle_native_task<
790786
})
791787
.unwrap_or_else(|| PumpxAccountProfile { wallet_exported: true });
792788
if let Err(e) = omni_account_profile_storage
793-
.insert(&sender.to_omni_account_with_client_id(CLIENT_ID_PUMPX), profile)
789+
.insert(&sender.to_omni_account_with_client_id(client_id), profile)
794790
{
795791
error!("Failed to update pumpx account profile: {:?}", e);
796792
send_error(
@@ -813,10 +809,9 @@ async fn handle_native_task<
813809
},
814810
NativeTask::PumpxAddWallet(sender) => {
815811
let storage = HeimaJwtStorage::new(ctx.storage_db.clone());
816-
let Ok(Some(access_token)) = storage.get(&(
817-
sender.to_omni_account_with_client_id(CLIENT_ID_PUMPX),
818-
AUTH_TOKEN_ACCESS_TYPE,
819-
)) else {
812+
let Ok(Some(access_token)) = storage
813+
.get(&(sender.to_omni_account_with_client_id(client_id), AUTH_TOKEN_ACCESS_TYPE))
814+
else {
820815
send_error(
821816
format!("Failed to get pumpx_{}_jwt_token", AUTH_TOKEN_ACCESS_TYPE),
822817
response_sender,
@@ -891,10 +886,9 @@ async fn handle_native_task<
891886
) => {
892887
// 1. Verify we have a valid Pumpx "access" token for the user
893888
let storage = HeimaJwtStorage::new(ctx.storage_db.clone());
894-
let Ok(Some(access_token)) = storage.get(&(
895-
sender.to_omni_account_with_client_id(CLIENT_ID_PUMPX),
896-
AUTH_TOKEN_ACCESS_TYPE,
897-
)) else {
889+
let Ok(Some(access_token)) = storage
890+
.get(&(sender.to_omni_account_with_client_id(client_id), AUTH_TOKEN_ACCESS_TYPE))
891+
else {
898892
send_error(
899893
"Failed to get access_token within NativeTask::PumpxTransferWidthdraw"
900894
.to_string(),

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,17 @@ pub struct RPCAddWalletResponse {
2323
pub backend_response: AddWalletResponse,
2424
}
2525

26-
impl From<AddWalletParams> for NativeTaskWrapper<NativeTask> {
27-
fn from(p: AddWalletParams) -> Self {
28-
let sender = Identity::from_web2_account(p.user_email.as_str(), Web2IdentityType::Email);
29-
NativeTaskWrapper::new(NativeTask::PumpxAddWallet(sender.clone()), None, None)
26+
impl AddWalletParams {
27+
pub fn into_native_task_wrapper(self, client_id: String) -> NativeTaskWrapper<NativeTask> {
28+
let sender = Identity::from_web2_account(self.user_email.as_str(), Web2IdentityType::Email);
29+
NativeTaskWrapper::new(NativeTask::PumpxAddWallet(sender.clone()), None, None, client_id)
3030
}
3131
}
3232

3333
pub fn register_add_wallet(module: &mut RpcModule<RpcContext>) {
3434
module
3535
.register_async_method("omni_addWallet", |params, ctx, ext| async move {
36-
check_auth(&ext).map_err(|e| {
36+
let user = check_auth(&ext).map_err(|e| {
3737
error!("Authentication check failed: {:?}", e);
3838
PumpxRpcError::from_error_code(ErrorCode::ServerError(
3939
AUTH_VERIFICATION_FAILED_CODE,
@@ -46,7 +46,7 @@ pub fn register_add_wallet(module: &mut RpcModule<RpcContext>) {
4646

4747
debug!("Received omni_addWallet, user_email: {}", params.user_email);
4848

49-
let wrapper: NativeTaskWrapper<NativeTask> = params.into();
49+
let wrapper = params.into_native_task_wrapper(user.client_id);
5050

5151
handle_omni_native_task(&ctx, wrapper, |task_ok| match task_ok {
5252
NativeTaskOk::PumpxAddWallet(response) => {

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ where
124124

125125
pub struct User {
126126
pub omni_account: String,
127+
pub client_id: String,
127128
}
128129

129130
/// This is used to verify that the request is authenticated.
@@ -132,7 +133,10 @@ pub struct User {
132133
/// Check rpc_middleware.rs
133134
pub fn check_auth(ext: &Extensions) -> Result<User, ()> {
134135
if let Some(rpc_extensions) = ext.get::<RpcExtensions>() {
135-
return Ok(User { omni_account: rpc_extensions.sender.clone() });
136+
return Ok(User {
137+
omni_account: rpc_extensions.sender.clone(),
138+
client_id: rpc_extensions.client_id.clone(),
139+
});
136140
}
137141
Err(())
138142
}

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

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ use super::common::handle_omni_native_task;
1717

1818
#[derive(Debug, Deserialize)]
1919
pub struct ExportWalletParams {
20-
pub user_email: String,
2120
pub client_id: String,
21+
pub user_email: String,
2222
pub key: Bytes, // RSA-encrypted AES key to encrypt the wallet private key, in 0x-hex-string
2323
pub google_code: String,
2424
pub chain_id: PumpxChainId,
@@ -27,25 +27,25 @@ pub struct ExportWalletParams {
2727
pub email_code: String,
2828
}
2929

30-
impl From<ExportWalletParams> for NativeTaskWrapper<NativeTask> {
31-
fn from(p: ExportWalletParams) -> Self {
30+
impl ExportWalletParams {
31+
pub fn into_native_task_wrapper(self) -> NativeTaskWrapper<NativeTask> {
3232
NativeTaskWrapper::new(
3333
NativeTask::PumpxExportWallet(
34-
Identity::from_web2_account(p.user_email.as_str(), Web2IdentityType::Pumpx),
35-
p.google_code,
36-
p.chain_id,
37-
p.wallet_index,
38-
p.wallet_address,
34+
Identity::from_web2_account(self.user_email.as_str(), Web2IdentityType::Pumpx),
35+
self.google_code,
36+
self.chain_id,
37+
self.wallet_index,
38+
self.wallet_address,
3939
),
4040
None,
41-
Some(OmniAuth::Email(p.client_id.clone(), p.user_email, p.email_code)),
41+
Some(OmniAuth::Email(self.client_id.clone(), self.user_email, self.email_code)),
42+
self.client_id,
4243
)
4344
}
4445
}
4546

4647
pub fn register_export_wallet(module: &mut RpcModule<RpcContext>) {
47-
module
48-
.register_async_method("omni_exportWallet", |params, ctx, _| async move {
48+
module .register_async_method("omni_exportWallet", |params, ctx, _ext| async move {
4949
let params = params.parse::<ExportWalletParams>().map_err(|e| {
5050
error!("Failed to parse params: {:?}", e);
5151
PumpxRpcError::from_error_code(ErrorCode::ParseError)
@@ -72,7 +72,7 @@ pub fn register_export_wallet(module: &mut RpcModule<RpcContext>) {
7272
)
7373
})?;
7474

75-
let wrapper: NativeTaskWrapper<NativeTask> = params.into();
75+
let wrapper = params.into_native_task_wrapper();
7676

7777
if wrapper.task.require_auth() {
7878
let Some(ref auth) = wrapper.auth else {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ pub fn register_notify_limit_order_result(module: &mut RpcModule<RpcContext>) {
4949
),
5050
None,
5151
None,
52+
user.client_id,
5253
);
5354

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

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

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@ use super::common::{check_omni_api_response, handle_omni_native_task};
1515

1616
#[derive(Debug, Deserialize)]
1717
pub struct RequestJwtParams {
18-
pub user_email: String,
1918
pub client_id: String,
19+
pub user_email: String,
2020
pub invite_code: Option<String>,
2121
pub google_code: String,
2222
pub language: Option<String>,
@@ -30,33 +30,37 @@ pub struct RequestJwtResponse {
3030
pub backend_response: UserConnectResponse,
3131
}
3232

33-
impl From<RequestJwtParams> for NativeTaskWrapper<NativeTask> {
34-
fn from(p: RequestJwtParams) -> Self {
33+
impl RequestJwtParams {
34+
pub fn into_native_task_wrapper(self) -> NativeTaskWrapper<NativeTask> {
3535
NativeTaskWrapper::new(
3636
NativeTask::PumpxRequestJwt(
37-
Identity::from_web2_account(p.user_email.as_str(), Web2IdentityType::Email), // actually unused
38-
p.user_email.clone(),
39-
p.invite_code,
40-
p.google_code,
41-
p.language,
37+
Identity::from_web2_account(self.user_email.as_str(), Web2IdentityType::Email), // actually unused
38+
self.user_email.clone(),
39+
self.invite_code,
40+
self.google_code,
41+
self.language,
4242
),
4343
None,
44-
Some(OmniAuth::Email(p.client_id.clone(), p.user_email, p.email_code)),
44+
Some(OmniAuth::Email(self.client_id.clone(), self.user_email, self.email_code)),
45+
self.client_id,
4546
)
4647
}
4748
}
4849

4950
pub fn register_request_jwt(module: &mut RpcModule<RpcContext>) {
5051
module
51-
.register_async_method("omni_requestJwt", |params, ctx, _| async move {
52+
.register_async_method("omni_requestJwt", |params, ctx, _ext| async move {
5253
let params = params.parse::<RequestJwtParams>().map_err(|e| {
5354
error!("Failed to parse params: {:?}", e);
5455
PumpxRpcError::from_error_code(ErrorCode::ParseError)
5556
})?;
5657

57-
debug!("Received omni_requestJwt, user_email: {}", params.user_email);
58+
debug!(
59+
"Received omni_requestJwt, user_email: {}, client_id: {}",
60+
params.user_email, params.client_id
61+
);
5862

59-
let wrapper: NativeTaskWrapper<NativeTask> = params.into();
63+
let wrapper = params.into_native_task_wrapper();
6064

6165
if wrapper.task.require_auth() {
6266
let Some(ref auth) = wrapper.auth else {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ pub fn register_sign_limit_order_params(module: &mut RpcModule<RpcContext>) {
8383
),
8484
None,
8585
None,
86+
user.client_id,
8687
);
8788

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

0 commit comments

Comments
 (0)