From 05759a5149911dfaa38282c8c066660a019a8433 Mon Sep 17 00:00:00 2001
From: Noor <noormohammedb@protonmail.com>
Date: Wed, 28 May 2025 17:17:26 +0530
Subject: [PATCH] improve query handling in new_vm

enhanced transaction error handling
bind all string input into NewVmReq submission query
---
 src/db/vm.rs | 87 ++++++++++++++++++++++++++++++----------------------
 1 file changed, 51 insertions(+), 36 deletions(-)

diff --git a/src/db/vm.rs b/src/db/vm.rs
index eeb449c..f7413a1 100644
--- a/src/db/vm.rs
+++ b/src/db/vm.rs
@@ -202,52 +202,67 @@ impl NewVmReq {
 
     pub async fn submit(self, db: &Surreal<Client>) -> Result<(), Error> {
         let locked_nano = self.locked_nano;
-        let account = self.admin.key().to_string();
-        let vm_id = self.id.key().to_string();
-        let vm_node = self.vm_node.key().to_string();
-        // TODO: check for possible injection and maybe use .bind()
-        let query = format!(
-            "
+        let tx_query = format!("
             BEGIN TRANSACTION;
-            UPDATE account:{account} SET balance -= {locked_nano};
-            IF account:{account}.balance < 0 {{
-                THROW 'Insufficient funds.'
-            }};
-            UPDATE account:{account} SET tmp_locked += {locked_nano};
-            RELATE
-                account:{account}
-                ->new_vm_req:{vm_id}
-                ->vm_node:{vm_node}
-            CONTENT {{
-                created_at: time::now(), hostname: '{}', vcpus: {}, memory_mb: {}, disk_size_gb: {},
-                extra_ports: {:?}, public_ipv4: {:?}, public_ipv6: {:?},
-                dtrfs_url: '{}', dtrfs_sha: '{}', kernel_url: '{}', kernel_sha: '{}',
-                price_per_unit: {}, locked_nano: {locked_nano}, error: ''
-            }};
-            COMMIT TRANSACTION;
-        ",
-            self.hostname,
+
+                LET $account = $account_input;
+                LET $new_vm_req = $new_vm_req_input;
+                LET $vm_node = $vm_node_input;
+                LET $hostname = $hostname_input;
+                LET $dtrfs_url = $dtrfs_url_input;
+                LET $dtrfs_sha = $dtrfs_sha_input;
+                LET $kernel_url = $kernel_url_input;
+                LET $kernel_sha = $kernel_sha_input;
+
+                UPDATE $account SET balance -= {locked_nano};
+                IF $account.balance < 0 {{
+                    THROW 'Insufficient funds.'
+                }};
+                UPDATE $account SET tmp_locked += {locked_nano};
+                RELATE
+                    $account
+                    ->$new_vm_req
+                    ->$vm_node
+                CONTENT {{
+                    created_at: time::now(), hostname: $hostname, vcpus: {}, memory_mb: {}, disk_size_gb: {},
+                    extra_ports: {:?}, public_ipv4: {}, public_ipv6: {},
+                    dtrfs_url: $dtrfs_url, dtrfs_sha: $dtrfs_sha, kernel_url: $kernel_url, kernel_sha: $kernel_sha,
+                    price_per_unit: {}, locked_nano: {locked_nano}, error: ''
+                }};
+            
+            COMMIT TRANSACTION;",
             self.vcpus,
             self.memory_mb,
             self.disk_size_gb,
             self.extra_ports,
             self.public_ipv4,
             self.public_ipv6,
-            self.dtrfs_url,
-            self.dtrfs_sha,
-            self.kernel_url,
-            self.kernel_sha,
             self.price_per_unit
         );
-        //let _: Vec<Self> = db.insert(NEW_VM_REQ).relation(self).await?;
-        let mut query_resp = db.query(query).await?;
-        let resp_err = query_resp.take_errors();
+        let mut query_resp = db
+            .query(tx_query)
+            .bind(("account_input", self.admin))
+            .bind(("new_vm_req_input", self.id))
+            .bind(("vm_node_input", self.vm_node))
+            .bind(("hostname_input", self.hostname))
+            .bind(("dtrfs_url_input", self.dtrfs_url))
+            .bind(("dtrfs_sha_input", self.dtrfs_sha))
+            .bind(("kernel_url_input", self.kernel_url))
+            .bind(("kernel_sha_input", self.kernel_sha))
+            .await?;
 
-        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);
+        let query_err = query_resp.take_errors();
+        if !query_err.is_empty() {
+            let tx_fail_err_str =
+                String::from("The query was not executed due to a failed transaction");
+
+            if query_err.contains_key(&9) && query_err[&9].to_string() != tx_fail_err_str {
+                log::error!("Transaction error: {}", query_err[&9]);
+                return Err(Error::InsufficientFunds);
+            } else {
+                log::error!("Unknown error in submit_new_vm_req: {query_err:?}");
+                return Err(Error::Unknown("submit_new_vm_req".to_string()));
+            }
         }
 
         Ok(())