From db60a3f8077c4574b58ae62f1345d83470b430d4 Mon Sep 17 00:00:00 2001 From: Noor Date: Fri, 23 May 2025 00:42:54 +0530 Subject: [PATCH] Enhanced kick contract calculating refund inside transaction query db function to calculate price per minute for app binded all kick contract query to prevent sql injection proper logging for kick contract fix table schema app_node typo and default time on deleted_app removed wrapper contract struct removed migrating timer.sql fixed some clippy warnings --- saved_data.yaml | 2 +- src/constants.rs | 3 +- src/db/app.rs | 10 ++- src/db/general.rs | 168 ++++++++++++++----------------------- src/db/mod.rs | 9 +- src/db/vm.rs | 23 ++--- src/grpc/app.rs | 2 +- src/grpc/general.rs | 13 ++- src/grpc/mod.rs | 8 +- src/grpc/types.rs | 14 ++-- surql/functions.sql | 10 +++ surql/tables.sql | 4 +- tests/grpc_general_test.rs | 18 ++-- 13 files changed, 126 insertions(+), 158 deletions(-) diff --git a/saved_data.yaml b/saved_data.yaml index 9b9dc16..26195ad 100644 --- a/saved_data.yaml +++ b/saved_data.yaml @@ -128,7 +128,7 @@ operators: - 7Xw3RxbP5pvfjZ8U6yA3HHVSS9YXjKH5Vkas3JRbQYd9 app_nodes: [] 7V3rEuh6j8VuwMVB5PyGqWKLmjJ4fYSv6WtrTL51NZTB: - escrow: 0 + escrow: 888888888899999 email: "" banned_users: [] vm_nodes: [] diff --git a/src/constants.rs b/src/constants.rs index 87f6210..33cc61f 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -5,8 +5,7 @@ pub const CERT_PATH: &str = "/etc/detee/brain/brain-crt.pem"; pub const CERT_KEY_PATH: &str = "/etc/detee/brain/brain-key.pem"; pub const CONFIG_PATH: &str = "/etc/detee/brain/config.ini"; -pub const DB_SCHEMA_FILES: [&str; 3] = - ["surql/tables.sql", "surql/timer.sql", "surql/functions.sql"]; +pub const DB_SCHEMA_FILES: [&str; 2] = ["surql/tables.sql", "surql/functions.sql"]; pub static ADMIN_ACCOUNTS: LazyLock> = LazyLock::new(|| { let default_admin_keys = vec![ diff --git a/src/db/app.rs b/src/db/app.rs index 2ce99ae..2245313 100644 --- a/src/db/app.rs +++ b/src/db/app.rs @@ -351,8 +351,12 @@ impl From for ActiveApp { impl ActiveAppWithNode { pub async fn get_by_uuid(db: &Surreal, uuid: &str) -> Result, Error> { - let contract: Option = - db.query(format!("select * from {ACTIVE_APP}:{uuid} fetch out;")).await?.take(0)?; + let contract: Option = db + .query(format!("select * from {ACTIVE_APP} where id = $uuid_input fetch out;")) + .bind(("uuid_input", RecordId::from((ACTIVE_APP, uuid)))) + .await? + .take(0)?; + Ok(contract) } @@ -460,7 +464,7 @@ impl From<&old_brain::BrainData> for Vec { .clone() .unwrap_or_default() .iter() - .map(|byte| format!("{:02X}", byte)) + .map(|byte| format!("{byte:02X}")) .collect(); contracts.push(ActiveApp { diff --git a/src/db/general.rs b/src/db/general.rs index ce7fe0a..7597e7a 100644 --- a/src/db/general.rs +++ b/src/db/general.rs @@ -276,114 +276,60 @@ impl Operator { } } -pub enum WrapperContract { - Vm(ActiveVmWithNode), - App(ActiveAppWithNode), -} - -impl WrapperContract { - pub async fn kick_contract( - db: &Surreal, - operator_wallet: &str, - contract_uuid: &str, - reason: &str, - ) -> Result { - let ( - operator_id, - admin_id, - contract_id, - collected_at, - price_per_mint, - deleted_table, - platform_specific_query, - ) = if let Some(active_vm) = ActiveVmWithNode::get_by_uuid(db, contract_uuid).await? { - let price_per_minute = active_vm.price_per_minute(); - - ( - active_vm.vm_node.operator, - active_vm.admin, - active_vm.id, - active_vm.collected_at, - price_per_minute, - "deleted_vm", - " - hostname = $contract.hostname, - public_ipv4 = $contract.public_ipv4, - public_ipv6 = $contract.public_ipv6, - dtrfs_sha = $contract.dtrfs_sha, - kernel_sha = $contract.kernel_sha, - ", - ) +pub async fn kick_contract( + db: &Surreal, + operator_wallet: &str, + contract_uuid: &str, + reason: &str, +) -> Result { + let (contract_id, operator_id, admin_id, app_or_vm) = + if let Some(active_vm) = ActiveVmWithNode::get_by_uuid(db, contract_uuid).await? { + (active_vm.id, active_vm.vm_node.operator, active_vm.admin, "vm") } else if let Some(active_app) = ActiveAppWithNode::get_by_uuid(db, contract_uuid).await? { - let price_per_minute = Into::::into(active_app.clone()).price_per_minute(); - - ( - active_app.app_node.operator, - active_app.admin, - active_app.id, - active_app.collected_at, - price_per_minute, - "deleted_app", - " - app_name = $contract.app_name, - host_ipv4 = $contract.host_ipv4, - mr_enclave = $contract.mr_enclave, - package_url= $contract.package_url, - hratls_pubkey = $contract.hratls_pubkey, - ", - ) + (active_app.id, active_app.app_node.operator, active_app.admin, "app") } else { return Err(Error::ContractNotFound); }; - let operator = operator_id.key().to_string(); - let admin = admin_id.key().to_string(); + if operator_id.key().to_string() != operator_wallet { + return Err(Error::AccessDenied); + } - if operator != operator_wallet { - return Err(Error::AccessDenied); - } - let mut minutes_to_refund = - chrono::Utc::now().signed_duration_since(*collected_at).num_minutes().unsigned_abs(); + log::debug!("Kicking contract {contract_id} by operator {operator_id} for reason: '{reason}'",); - let one_week_minute = 10080; - - if minutes_to_refund > one_week_minute { - minutes_to_refund = one_week_minute; - } - - let refund_amount = minutes_to_refund * price_per_mint; - - log::debug!("Removing {refund_amount} escrow from {} and giving it to {}", operator, admin); - - let transaction_query = format!( - " - BEGIN TRANSACTION; + let transaction_query = format!( + " + BEGIN TRANSACTION; LET $contract = {contract_id}; LET $operator_account = {operator_id}; - LET $reason = '{reason}'; - LET $refund_amount = {refund_amount}; - LET $deleted_contract = {deleted_table}:{contract_uuid}; - LET $id = record::id($contract.id); + LET $reason = $reason_str_input; + LET $contract_id = record::id($contract.id); LET $admin = $contract.in; LET $node = $contract.out; - -- move contract into deleted state + LET $active_contract = (select * from $contract)[0]; + LET $deleted_contract = $active_contract.patch([{{ + 'op': 'replace', + 'path': 'id', + 'value': type::record('deleted_{app_or_vm}:' + $contract_id) + }}]); + LET $deleted_contract = (INSERT RELATION INTO deleted_{app_or_vm} ( $deleted_contract ) RETURN AFTER)[0]; - RELATE $admin->{deleted_table}->$node - SET id = $id, - {platform_specific_query} - mapped_ports = $contract.mapped_ports, - disk_size_gb = $contract.disk_size_gb, - vcpus = $contract.vcpus, - memory_mb = $contract.memory_mb, - created_at = $contract.created_at, - deleted_at = time::now(), - price_per_unit = $contract.price_per_unit - ; - - DELETE $contract; + -- calculating refund minutes + LET $one_week_minutes = duration::mins(1w); + LET $uncollected_minutes = (time::now() - $active_contract.collected_at).mins(); + + LET $minutes_to_refund = if $uncollected_minutes > $one_week_minutes {{ + $one_week_minutes; + }} ELSE {{ + $uncollected_minutes; + }}; -- calculating refund amount + LET $prince_per_minute = fn::{app_or_vm}_price_per_minute($active_contract.id); + + LET $refund_amount = $prince_per_minute * $minutes_to_refund; + LET $refund = IF SELECT * FROM {KICK} WHERE out = $admin.id AND created_at > time::now() - 24h {{ 0 }} ELSE IF $operator_account.escrow <= $refund_amount {{ @@ -392,14 +338,14 @@ impl WrapperContract { $refund_amount; }}; - RELATE $operator_account->{KICK}->$admin - SET id = $id, + SET id = $contract_id, reason = $reason, - contract = $deleted_contract, - node = $node, + contract = $deleted_contract.id, + node = $node.id, created_at = time::now() ; + DELETE $active_contract.id; -- update balances UPDATE $operator_account SET escrow -= $refund; @@ -408,16 +354,28 @@ impl WrapperContract { }}; UPDATE $admin SET balance += $refund; - SELECT * FROM $refund; + $refund; - COMMIT TRANSACTION; - ", - ); + COMMIT TRANSACTION; + ", + ); - log::trace!("kick_contract transaction_query: {}", &transaction_query); - let refunded: Option = db.query(transaction_query).await?.take(14)?; - let refunded_amount = refunded.ok_or(Error::FailedToCreateDBEntry("Refund".to_string()))?; + log::trace!("kick_contract transaction_query: {}", &transaction_query); - Ok(refunded_amount) + let mut query_res = + db.query(transaction_query).bind(("reason_str_input", reason.to_string())).await?; + + log::trace!("transaction_query response: {:?}", &query_res); + + let query_error = query_res.take_errors(); + if !query_error.is_empty() { + log::error!("kick_contract query error: {query_error:?}"); + return Err(Error::FailedKickContract(contract_id.to_string())); } + + let refunded: Option = query_res.take(20)?; + let refunded_amount = refunded.ok_or(Error::FailedToCreateDBEntry("Refund".to_string()))?; + log::info!("Refunded: {refunded_amount} to {admin_id}"); + + Ok(refunded_amount) } diff --git a/src/db/mod.rs b/src/db/mod.rs index 818f6c3..cd06be1 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -43,6 +43,8 @@ pub enum Error { AccessDenied, #[error("Failed to delete contract {0}")] FailedToDeleteContract(String), + #[error("Failed to kick contract {0}")] + FailedKickContract(String), } pub mod prelude { @@ -77,8 +79,11 @@ pub async fn migration0( let active_vm: Vec = old_data.into(); let active_app: Vec = old_data.into(); - for schema in DB_SCHEMA_FILES.map(std::fs::read_to_string) { - db.query(schema?).await?; + for schema_data in DB_SCHEMA_FILES.map(|path| (std::fs::read_to_string(path), path)) { + let schema_file = schema_data.1; + println!("Loading schema from {schema_file}"); + let schema = schema_data.0?; + db.query(schema).await?; } println!("Inserting accounts..."); diff --git a/src/db/vm.rs b/src/db/vm.rs index ff5e3fe..c9ddcee 100644 --- a/src/db/vm.rs +++ b/src/db/vm.rs @@ -219,7 +219,7 @@ impl NewVmReq { ->vm_node:{vm_node} CONTENT {{ created_at: time::now(), hostname: '{}', vcpus: {}, memory_mb: {}, disk_size_gb: {}, - extra_ports: {}, public_ipv4: {:?}, public_ipv6: {:?}, + extra_ports: {:?}, public_ipv4: {:?}, public_ipv6: {:?}, dtrfs_url: '{}', dtrfs_sha: '{}', kernel_url: '{}', kernel_sha: '{}', price_per_unit: {}, locked_nano: {locked_nano}, error: '' }}; @@ -229,7 +229,7 @@ impl NewVmReq { self.vcpus, self.memory_mb, self.disk_size_gb, - format!("{:?}", self.extra_ports,), + self.extra_ports, self.public_ipv4, self.public_ipv6, self.dtrfs_url, @@ -242,13 +242,11 @@ impl NewVmReq { let mut query_resp = db.query(query).await?; let resp_err = query_resp.take_errors(); - if let Some(insufficient_funds_error) = resp_err.get(&1) { - if let surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error)) = - insufficient_funds_error - { - log::error!("Transaction error: {}", tx_query_error); - return Err(Error::InsufficientFunds); - } + if let Some(surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error))) = + resp_err.get(&1) + { + log::error!("Transaction error: {tx_query_error}"); + return Err(Error::InsufficientFunds); } Ok(()) @@ -753,8 +751,11 @@ impl From for ActiveVm { impl ActiveVmWithNode { pub async fn get_by_uuid(db: &Surreal, uuid: &str) -> Result, Error> { - let contract: Option = - db.query(format!("select * from {ACTIVE_VM}:{uuid} fetch out;")).await?.take(0)?; + let contract: Option = db + .query(format!("select * from {ACTIVE_VM} where id = $uuid_input fetch out;")) + .bind(("uuid_input", RecordId::from((ACTIVE_VM, uuid)))) + .await? + .take(0)?; Ok(contract) } diff --git a/src/grpc/app.rs b/src/grpc/app.rs index 1cf31ee..78be5af 100644 --- a/src/grpc/app.rs +++ b/src/grpc/app.rs @@ -115,7 +115,7 @@ impl BrainAppDaemon for AppDaemonServer { let mut req_stream = req.into_inner(); let pubkey: String; if let Some(Ok(msg)) = req_stream.next().await { - log::debug!("App daemon_messages received auth message: {:?}", msg); + log::debug!("App daemon_messages received auth message: {msg:?}"); if let Some(daemon_message_app::Msg::Auth(auth)) = msg.msg { pubkey = auth.pubkey.clone(); check_sig_from_parts( diff --git a/src/grpc/general.rs b/src/grpc/general.rs index 35e15b9..cf3bf2b 100644 --- a/src/grpc/general.rs +++ b/src/grpc/general.rs @@ -127,13 +127,9 @@ impl BrainGeneralCli for GeneralCliServer { async fn kick_contract(&self, req: Request) -> Result, Status> { let req = check_sig_from_req(req)?; - match db::WrapperContract::kick_contract( - &self.db, - &req.operator_wallet, - &req.contract_uuid, - &req.reason, - ) - .await + log::info!("Kicking contract: {}, by: {}", req.contract_uuid, req.operator_wallet); + match db::kick_contract(&self.db, &req.operator_wallet, &req.contract_uuid, &req.reason) + .await { Ok(nano_lp) => Ok(Response::new(KickResp { nano_lp })), Err(e) @@ -144,10 +140,11 @@ impl BrainGeneralCli for GeneralCliServer { | db::Error::FailedToDeleteContract(_) ) => { + log::warn!("Failed to kick contract: {e:?}"); Err(Status::failed_precondition(e.to_string())) } Err(e) => { - log::info!("Failed to kick contract: {e:?}"); + log::error!("Failed to kick contract: {e:?}"); Err(Status::unknown( "Unknown error. Please try again or contact the DeTEE devs team.", )) diff --git a/src/grpc/mod.rs b/src/grpc/mod.rs index 4fbf0d4..c78c346 100644 --- a/src/grpc/mod.rs +++ b/src/grpc/mod.rs @@ -58,7 +58,7 @@ impl_pubkey_getter!(RegisterAppNodeReq); impl_pubkey_getter!(AppNodeFilters); pub fn check_sig_from_req(req: Request) -> Result { - log::trace!("Checking signature from request: {:?}", req); + log::trace!("Checking signature from request: {req:?}"); let time = match req.metadata().get("timestamp") { Some(t) => t.clone(), None => return Err(Status::unauthenticated("Timestamp not found in metadata.")), @@ -73,8 +73,7 @@ pub fn check_sig_from_req(req: Request) -> let seconds_elapsed = now.signed_duration_since(parsed_time).num_seconds(); if !(-4..=4).contains(&seconds_elapsed) { return Err(Status::unauthenticated(format!( - "Date is not within 4 sec of the time of the server: CLI {} vs Server {}", - parsed_time, now + "Date is not within 4 sec of the time of the server: CLI {parsed_time} vs Server {now}", ))); } @@ -131,8 +130,7 @@ pub fn check_sig_from_parts(pubkey: &str, time: &str, msg: &str, sig: &str) -> R let seconds_elapsed = now.signed_duration_since(parsed_time).num_seconds(); if !(-4..=4).contains(&seconds_elapsed) { return Err(Status::unauthenticated(format!( - "Date is not within 4 sec of the time of the server: CLI {} vs Server {}", - parsed_time, now + "Date is not within 4 sec of the time of the server: CLI {parsed_time} vs Server {now}", ))); } diff --git a/src/grpc/types.rs b/src/grpc/types.rs index dfde396..ee7e716 100644 --- a/src/grpc/types.rs +++ b/src/grpc/types.rs @@ -263,15 +263,15 @@ impl From for AppContract { node_pubkey: value.app_node.id.key().to_string(), public_ipv4: value.host_ipv4, resource: Some(AppResource { - memory_mb: value.memory_mb as u32, - disk_mb: value.disk_size_gb as u32, - vcpu: value.vcpus as u32, - ports: value.mapped_ports.iter().map(|(_, g)| *g as u32).collect(), + memory_mb: value.memory_mb, + disk_mb: value.disk_size_gb, + vcpu: value.vcpus, + ports: value.mapped_ports.iter().map(|(_, g)| *g).collect(), }), mapped_ports: value .mapped_ports .iter() - .map(|(h, g)| MappedPort { host_port: *h as u32, guest_port: *g as u32 }) + .map(|(h, g)| MappedPort { host_port: *h, guest_port: *g }) .collect(), created_at: value.created_at.to_rfc3339(), @@ -294,7 +294,7 @@ impl From for db::NewAppReq { .public_package_mr_enclave .unwrap_or_default() .iter() - .fold(String::new(), |acc, x| acc + &format!("{:02x?}", x)); + .fold(String::new(), |acc, x| acc + &format!("{x:02x?}")); Self { id: RecordId::from((NEW_APP_REQ, nanoid!(40, &ID_ALPHABET))), @@ -377,7 +377,7 @@ impl From for NewAppRes { let mapped_ports = val .mapped_ports .iter() - .map(|(h, g)| MappedPort { host_port: *h as u32, guest_port: *g as u32 }) + .map(|(h, g)| MappedPort { host_port: *h, guest_port: *g }) .collect(); Self { uuid: val.id.key().to_string(), diff --git a/surql/functions.sql b/surql/functions.sql index 689635c..062b62d 100644 --- a/surql/functions.sql +++ b/surql/functions.sql @@ -25,3 +25,13 @@ DEFINE FUNCTION OVERWRITE fn::delete_vm( DELETE $vm.id; }; +DEFINE FUNCTION OVERWRITE fn::app_price_per_minute( + $app_id: record +) { + LET $app = (select * from $app_id)[0]; + RETURN + (($app.vcpus * 5) + + ($app.memory_mb / 200) + + ($app.disk_size_gb / 10)) + * $app.price_per_unit; +}; \ No newline at end of file diff --git a/surql/tables.sql b/surql/tables.sql index 35ecc1c..0071da4 100644 --- a/surql/tables.sql +++ b/surql/tables.sql @@ -129,7 +129,7 @@ DEFINE FIELD vcpus ON TABLE deleted_app TYPE int; DEFINE FIELD memory_mb ON TABLE deleted_app TYPE int; DEFINE FIELD disk_size_gb ON TABLE deleted_app TYPE int; DEFINE FIELD created_at ON TABLE deleted_app TYPE datetime; -DEFINE FIELD deleted_at ON TABLE deleted_app TYPE datetime; +DEFINE FIELD deleted_at ON TABLE deleted_app TYPE datetime DEFAULT time::now();; DEFINE FIELD price_per_unit ON TABLE deleted_app TYPE int; DEFINE FIELD mr_enclave ON TABLE deleted_app TYPE string; DEFINE FIELD package_url ON TABLE deleted_app TYPE string; @@ -142,7 +142,7 @@ DEFINE TABLE kick TYPE RELATION FROM account TO account; DEFINE FIELD created_at ON TABLE kick TYPE datetime; DEFINE FIELD reason ON TABLE kick TYPE string; DEFINE FIELD contract ON TABLE kick TYPE record; -DEFINE FIELD node ON TABLE kick TYPE record; +DEFINE FIELD node ON TABLE kick TYPE record; DEFINE TABLE report TYPE RELATION FROM account TO vm_node|app_node; DEFINE FIELD created_at ON TABLE report TYPE datetime; diff --git a/tests/grpc_general_test.rs b/tests/grpc_general_test.rs index 2c601a0..e57c35b 100644 --- a/tests/grpc_general_test.rs +++ b/tests/grpc_general_test.rs @@ -222,23 +222,19 @@ async fn test_kick_contract() { // 7. refund of multiple contract kick in a day for same user env_logger::builder() - .filter_level(log::LevelFilter::Trace) + .filter_level(log::LevelFilter::Info) .filter_module("tungstenite", log::LevelFilter::Debug) .filter_module("tokio_tungstenite", log::LevelFilter::Debug) .init(); let db_conn = prepare_test_db().await.unwrap(); - let operator_wallet = "BFopWmwcZAMF1h2PFECZNdEucdZfnZZ32p6R9ZaBiVsS"; - let contract_uuid = "26577f1c98674a1780a86cf0490f1270"; - let reason = "test reason"; + let contract_uuid = "e3d01f252b2a410b80e312f44e474334"; + let operator_wallet = "7V3rEuh6j8VuwMVB5PyGqWKLmjJ4fYSv6WtrTL51NZTB"; + let reason = "'; THROW 'Injected error'; --"; // sql injection query - let kick_response = surreal_brain::db::general::WrapperContract::kick_contract( - &db_conn, - operator_wallet, - contract_uuid, - reason, - ) - .await; + let kick_response = + surreal_brain::db::general::kick_contract(&db_conn, operator_wallet, contract_uuid, reason) + .await; match kick_response { Ok(refund_amount) => { println!("Refund amount: {}", refund_amount);