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

@@ -128,7 +128,7 @@ operators:
       - 7Xw3RxbP5pvfjZ8U6yA3HHVSS9YXjKH5Vkas3JRbQYd9
     app_nodes: []
   7V3rEuh6j8VuwMVB5PyGqWKLmjJ4fYSv6WtrTL51NZTB:
-    escrow: 0
+    escrow: 888888888899999
     email: ""
     banned_users: []
     vm_nodes: []

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] =
+pub const DB_SCHEMA_FILES: [&str; 2] = ["surql/tables.sql", "surql/functions.sql"];
-    ["surql/tables.sql", "surql/timer.sql", "surql/functions.sql"];
 
 pub static ADMIN_ACCOUNTS: LazyLock<Vec<String>> = LazyLock::new(|| {
     let default_admin_keys = vec![

src/db/app.rs

@@ -351,8 +351,12 @@ impl From<ActiveAppWithNode> for ActiveApp {
 
 impl ActiveAppWithNode {
     pub async fn get_by_uuid(db: &Surreal<Client>, uuid: &str) -> Result<Option<Self>, Error> {
-        let contract: Option<Self> =
+        let contract: Option<Self> = db
-            db.query(format!("select * from {ACTIVE_APP}:{uuid} fetch out;")).await?.take(0)?;
+            .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<ActiveApp> {
                 .clone()
                 .unwrap_or_default()
                 .iter()
-                .map(|byte| format!("{:02X}", byte))
+                .map(|byte| format!("{byte:02X}"))
                 .collect();
 
             contracts.push(ActiveApp {

src/db/general.rs

@@ -276,114 +276,60 @@ impl Operator {
     }
 }
 
-pub enum WrapperContract {
+pub async fn kick_contract(
-    Vm(ActiveVmWithNode),
+    db: &Surreal<Client>,
-    App(ActiveAppWithNode),
+    operator_wallet: &str,
-}
+    contract_uuid: &str,
-
+    reason: &str,
-impl WrapperContract {
+) -> Result<u64, Error> {
-    pub async fn kick_contract(
+    let (contract_id, operator_id, admin_id, app_or_vm) =
-        db: &Surreal<Client>,
+        if let Some(active_vm) = ActiveVmWithNode::get_by_uuid(db, contract_uuid).await? {
-        operator_wallet: &str,
+            (active_vm.id, active_vm.vm_node.operator, active_vm.admin, "vm")
-        contract_uuid: &str,
-        reason: &str,
-    ) -> Result<u64, Error> {
-        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,
-                ",
-            )
         } else if let Some(active_app) = ActiveAppWithNode::get_by_uuid(db, contract_uuid).await? {
-            let price_per_minute = Into::<ActiveApp>::into(active_app.clone()).price_per_minute();
+            (active_app.id, active_app.app_node.operator, active_app.admin, "app")
-
-            (
-                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,
-                ",
-            )
         } else {
             return Err(Error::ContractNotFound);
         };
 
-        let operator = operator_id.key().to_string();
+    if operator_id.key().to_string() != operator_wallet {
-        let admin = admin_id.key().to_string();
+        return Err(Error::AccessDenied);
+    }
 
-        if operator != operator_wallet {
+    log::debug!("Kicking contract {contract_id} by operator {operator_id} for reason: '{reason}'",);
-            return Err(Error::AccessDenied);
-        }
-        let mut minutes_to_refund =
-            chrono::Utc::now().signed_duration_since(*collected_at).num_minutes().unsigned_abs();
 
-        let one_week_minute = 10080;
+    let transaction_query = format!(
-
+        "
-        if minutes_to_refund > one_week_minute {
+        BEGIN TRANSACTION;
-            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 $contract = {contract_id};
             LET $operator_account = {operator_id};
-            LET $reason = '{reason}';
+            LET $reason = $reason_str_input;
-            LET $refund_amount = {refund_amount};
+            LET $contract_id = record::id($contract.id);
-            LET $deleted_contract = {deleted_table}:{contract_uuid};
-            LET $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
+            -- calculating refund minutes
-            SET id = $id,
+            LET $one_week_minutes = duration::mins(1w);
-                {platform_specific_query}
+            LET $uncollected_minutes =  (time::now() - $active_contract.collected_at).mins();
-                mapped_ports = $contract.mapped_ports,
+            
-                disk_size_gb = $contract.disk_size_gb,
+            LET $minutes_to_refund = if $uncollected_minutes > $one_week_minutes {{
-                vcpus = $contract.vcpus,
+                $one_week_minutes;
-                memory_mb = $contract.memory_mb,
+            }} ELSE {{
-                created_at = $contract.created_at,
+                $uncollected_minutes;
-                deleted_at = time::now(),
+            }};
-                price_per_unit = $contract.price_per_unit
-            ;
-
-            DELETE $contract;
 
             -- 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,
+                contract = $deleted_contract.id,
-                node = $node,
+                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);
+    log::trace!("kick_contract transaction_query: {}", &transaction_query);
-        let refunded: Option<u64> = db.query(transaction_query).await?.take(14)?;
-        let refunded_amount = refunded.ok_or(Error::FailedToCreateDBEntry("Refund".to_string()))?;
 
-        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<u64> = 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)
 }

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<ActiveVm> = old_data.into();
     let active_app: Vec<ActiveApp> = old_data.into();
 
-    for schema in DB_SCHEMA_FILES.map(std::fs::read_to_string) {
+    for schema_data in DB_SCHEMA_FILES.map(|path| (std::fs::read_to_string(path), path)) {
-        db.query(schema?).await?;
+        let schema_file = schema_data.1;
+        println!("Loading schema from {schema_file}");
+        let schema = schema_data.0?;
+        db.query(schema).await?;
     }
 
     println!("Inserting accounts...");

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 Some(surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error))) =
-            if let surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error)) =
+            resp_err.get(&1)
-                insufficient_funds_error
+        {
-            {
+            log::error!("Transaction error: {tx_query_error}");
-                log::error!("Transaction error: {}", tx_query_error);
+            return Err(Error::InsufficientFunds);
-                return Err(Error::InsufficientFunds);
-            }
         }
 
         Ok(())
@@ -753,8 +751,11 @@ impl From<ActiveVmWithNode> for ActiveVm {
 
 impl ActiveVmWithNode {
     pub async fn get_by_uuid(db: &Surreal<Client>, uuid: &str) -> Result<Option<Self>, Error> {
-        let contract: Option<Self> =
+        let contract: Option<Self> = db
-            db.query(format!("select * from {ACTIVE_VM}:{uuid} fetch out;")).await?.take(0)?;
+            .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)
     }
 

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(

src/grpc/general.rs

@@ -127,13 +127,9 @@ impl BrainGeneralCli for GeneralCliServer {
 
     async fn kick_contract(&self, req: Request<KickReq>) -> Result<Response<KickResp>, Status> {
         let req = check_sig_from_req(req)?;
-        match db::WrapperContract::kick_contract(
+        log::info!("Kicking contract: {}, by: {}", req.contract_uuid, req.operator_wallet);
-            &self.db,
+        match db::kick_contract(&self.db, &req.operator_wallet, &req.contract_uuid, &req.reason)
-            &req.operator_wallet,
+            .await
-            &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.",
                 ))

src/grpc/mod.rs

@@ -58,7 +58,7 @@ impl_pubkey_getter!(RegisterAppNodeReq);
 impl_pubkey_getter!(AppNodeFilters);
 
 pub fn check_sig_from_req<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) -> Result<T, Status> {
-    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<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) ->
     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 {}",
+            "Date is not within 4 sec of the time of the server: CLI {parsed_time} vs Server {now}",
-            parsed_time, 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 {}",
+            "Date is not within 4 sec of the time of the server: CLI {parsed_time} vs Server {now}",
-            parsed_time, now
         )));
     }
 

src/grpc/types.rs

@@ -263,15 +263,15 @@ impl From<db::ActiveAppWithNode> 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,
+                memory_mb: value.memory_mb,
-                disk_mb: value.disk_size_gb as u32,
+                disk_mb: value.disk_size_gb,
-                vcpu: value.vcpus as u32,
+                vcpu: value.vcpus,
-                ports: value.mapped_ports.iter().map(|(_, g)| *g as u32).collect(),
+                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<NewAppReq> 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<db::ActiveApp> 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(),

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;
+};

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<deleted_vm|deleted_app>;
-DEFINE FIELD node ON TABLE kick TYPE record<vm_node|app_name>;
+DEFINE FIELD node ON TABLE kick TYPE record<vm_node|app_node>;
 
 DEFINE TABLE report TYPE RELATION FROM account TO vm_node|app_node;
 DEFINE FIELD created_at ON TABLE report TYPE datetime;

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 = "e3d01f252b2a410b80e312f44e474334";
-    let contract_uuid = "26577f1c98674a1780a86cf0490f1270";
+    let operator_wallet = "7V3rEuh6j8VuwMVB5PyGqWKLmjJ4fYSv6WtrTL51NZTB";
-    let reason = "test reason";
+    let reason = "'; THROW 'Injected error'; --"; // sql injection query
 
-    let kick_response = surreal_brain::db::general::WrapperContract::kick_contract(
+    let kick_response =
-        &db_conn,
+        surreal_brain::db::general::kick_contract(&db_conn, operator_wallet, contract_uuid, reason)
-        operator_wallet,
+            .await;
-        contract_uuid,
-        reason,
-    )
-    .await;
     match kick_response {
         Ok(refund_amount) => {
             println!("Refund amount: {}", refund_amount);