feat extend vm locked amount

extensive test for extend_vm feature
error handling for db transaction limits
mock data for test extend_vm
todo marking for refund tmp_locked

src/db/mod.rs

@@ -49,6 +49,10 @@ pub enum Error {
     AlreadyBanned(String),
     #[error("Failed to slash operator {0}")]
     FailedToSlashOperator(String),
+    #[error("Transation Too Big {0}")]
+    TooBigTransaction(String),
+    #[error("Unknown: {0}")]
+    Unknown(String),
 }
 
 pub mod prelude {

src/db/vm.rs

@@ -196,6 +196,7 @@ impl NewVmReq {
             error: String,
         }
         let _: Option<Self> = db.update((NEW_VM_REQ, id)).merge(NewVmError { error }).await?;
+        // TODO: IMP refund tmp_locked
         Ok(())
     }
 
@@ -492,6 +493,78 @@ impl ActiveVm {
             Ok(false)
         }
     }
+
+    pub async fn extend_time(
+        db: &Surreal<Client>,
+        id: &str,
+        admin: &str,
+        nano_lp: u64,
+    ) -> Result<(), Error> {
+        if nano_lp > 100_000_000_000_000 {
+            return Err(Error::TooBigTransaction(nano_lp.to_string()));
+        }
+
+        let tx_query = format!(
+            "
+            BEGIN TRANSACTION;
+
+                LET $contract = $contract_input;
+                LET $admin = $admin_input;
+                LET $lock_amt = {nano_lp};
+
+                if !record::exists($contract) {{
+                    THROW 'contract not exist ' + <string>$contract
+                }};
+                if $contract.in != $admin {{
+                    THROW 'Unauthorized'
+                }};
+                if $admin.balance + $contract.locked_nano < $lock_amt {{
+                    THROW 'InsufficientFunds'
+                }};
+
+                UPDATE $admin SET balance = $admin.balance + $contract.locked_nano - $lock_amt;
+                UPDATE $contract SET locked_nano = $lock_amt;
+
+            COMMIT TRANSACTION;
+            "
+        );
+
+        log::trace!("extend_time query: {tx_query}");
+
+        let mut query_res = db
+            .query(tx_query)
+            .bind(("contract_input", RecordId::from((ACTIVE_VM, id))))
+            .bind(("admin_input", RecordId::from((ACCOUNT, admin))))
+            .await?;
+
+        log::trace!("tx_query response: {query_res:?}");
+
+        let query_err = query_res.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(&3) && query_err[&3].to_string() != tx_fail_err_str {
+                log::error!("contract not exist: {}", query_err[&3]);
+                return Err(Error::ContractNotFound);
+            }
+
+            if query_err.contains_key(&4) && query_err[&4].to_string() != tx_fail_err_str {
+                log::error!("Unauthorized: {}", query_err[&4]);
+                return Err(Error::AccessDenied);
+            }
+
+            if query_err.contains_key(&5) && query_err[&5].to_string() != tx_fail_err_str {
+                log::error!("InsufficientFunds: {}", query_err[&5]);
+                return Err(Error::InsufficientFunds);
+            }
+
+            log::error!("Unknown error in extend_time: {query_err:?}");
+
+            return Err(Error::Unknown("extend_time".to_string()));
+        }
+        Ok(())
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]

src/grpc/vm.rs

@@ -313,15 +313,26 @@ impl BrainVmCli for VmCliServer {
     }
 
     async fn extend_vm(&self, req: Request<ExtendVmReq>) -> Result<Response<Empty>, Status> {
-        let _req = check_sig_from_req(req)?;
+        let req = check_sig_from_req(req)?;
-        todo!();
+        match db::ActiveVm::extend_time(&self.db, &req.uuid, &req.admin_pubkey, req.locked_nano)
-        // match self
+            .await
-        //     .data
+        {
-        //     .extend_vm_contract_time(&req.uuid, &req.admin_pubkey, req.locked_nano)
+            Ok(()) => Ok(Response::new(Empty {})),
-        // {
+            Err(e)
-        //     Ok(()) => Ok(Response::new(Empty {})),
+                if matches!(
-        //     Err(e) => Err(Status::unknown(format!("Could not extend contract: {e}"))),
+                    e,
-        // }
+                    db::Error::ContractNotFound
+                        | db::Error::AccessDenied
+                        | db::Error::InsufficientFunds
+                ) =>
+            {
+                Err(Status::failed_precondition(e.to_string()))
+            }
+            Err(e) => {
+                log::error!("Error extending VM contract {}: {e}", &req.uuid);
+                Err(Status::unknown(format!("Could not extend contract: {e}")))
+            }
+        }
     }
 
     async fn delete_vm(&self, req: Request<DeleteVmReq>) -> Result<Response<Empty>, Status> {

tests/common/vm_cli_utils.rs

@@ -161,3 +161,38 @@ pub async fn list_all_app_contracts(
 
     Ok(app_contracts)
 }
+
+pub async fn user_list_vm_contracts(
+    brain_channel: &Channel,
+    key: &Key,
+    as_operator: bool,
+    uuid: &str,
+) -> Result<Vec<vm_proto::VmContract>> {
+    let mut cli_client = BrainVmCliClient::new(brain_channel.clone());
+    let mut stream = cli_client
+        .list_vm_contracts(
+            key.sign_request(vm_proto::ListVmContractsReq {
+                wallet: key.pubkey.clone(),
+                as_operator,
+                uuid: uuid.to_string(),
+            })
+            .unwrap(),
+        )
+        .await?
+        .into_inner();
+
+    let mut vm_contracts = Vec::new();
+
+    while let Some(stream_data) = stream.next().await {
+        match stream_data {
+            Ok(vm_contract) => {
+                vm_contracts.push(vm_contract);
+            }
+            Err(e) => {
+                panic!("Error while listing vm_contracts: {e:?}");
+            }
+        }
+    }
+
+    Ok(vm_contracts)
+}

tests/grpc_vm_cli_test.rs

@@ -1,13 +1,14 @@
 use common::prepare_test_env::{prepare_test_db, run_service_for_stream};
 use common::test_utils::Key;
-use common::vm_cli_utils::{airdrop, create_new_vm};
+use common::vm_cli_utils::{airdrop, create_new_vm, user_list_vm_contracts};
 use common::vm_daemon_utils::{mock_vm_daemon, register_vm_node};
 use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient;
 use detee_shared::vm_proto::brain_vm_daemon_client::BrainVmDaemonClient;
-use detee_shared::vm_proto::{ListVmContractsReq, NewVmReq};
+use detee_shared::vm_proto::{ExtendVmReq, ListVmContractsReq, NewVmReq};
 use futures::StreamExt;
 use std::vec;
-use surreal_brain::constants::ACTIVE_VM;
+use surreal_brain::constants::{ACCOUNT, ACTIVE_VM};
+use surreal_brain::db::prelude as db;
 use surreal_brain::db::vm::ActiveVm;
 
 mod common;
@@ -102,3 +103,85 @@ async fn test_list_vm_contracts() {
 
     // verify report in db
 }
+
+/*
+Private Key: 69kPgwAzftzNf5oBFQwqocFAgsn17ZwKxCHFvpoZ1LRn
+Public Key : HYo88ncAPekykFDTFGdJV9ehJ81LL4onQ8xzYRAu2Ph4
+
+Public Key : Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG // operator
+Private Key: 2oQvpUMSzaZmaVLfVWsriRbhRK1JbnPEBed5UmpZSDNc */
+
+#[tokio::test]
+async fn test_extend_vm() {
+    let db_conn = prepare_test_db().await.unwrap();
+
+    let channel = run_service_for_stream().await.unwrap();
+    let mut cli_client = BrainVmCliClient::new(channel.clone());
+
+    /*
+    env_logger::builder()
+        .filter_level(log::LevelFilter::Debug)
+        .filter_module("tungstenite", log::LevelFilter::Debug)
+        .filter_module("tokio_tungstenite", log::LevelFilter::Debug)
+        .init();
+     */
+
+    let key = Key::from("YnaZccN852KYvnhV5pbeBrHV4dSwgzpQXghdyMiHC6i"); // GmE4JH3bL4NpmzwKCBJemJzRTumJAnbcXLGqce5mREgS admin
+    let vm_contract_id = "1d749a816c27a22efa0e574b023a6afef040fe5";
+
+    let locking_01 = 200;
+    let locking_02 = 86000000000u64;
+
+    let mut req = ExtendVmReq {
+        admin_pubkey: key.pubkey.clone(),
+        uuid: "foooooooo".to_string(),
+        locked_nano: u64::MAX,
+    };
+
+    let err_precondition =
+        cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.err().unwrap();
+    assert!(err_precondition.to_string().contains("Transation Too Big "));
+
+    req.locked_nano = 99999999999999;
+    let err_precondition =
+        cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.err().unwrap();
+    assert!(err_precondition.to_string().contains("Contract not found"));
+
+    req.uuid = vm_contract_id.to_string();
+    let err_precondition =
+        cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.err().unwrap();
+    assert!(err_precondition.to_string().contains("Insufficient funds"));
+
+    req.locked_nano = locking_01;
+    let err_precondition =
+        cli_client.extend_vm(Key::new().sign_request(req.clone()).unwrap()).await.err().unwrap();
+    assert!(err_precondition.to_string().contains("signature does not match"));
+
+    let acc: db::Account = db_conn.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
+    let contract: db::ActiveVm =
+        db_conn.select((ACTIVE_VM, vm_contract_id)).await.unwrap().unwrap();
+
+    let expected_bal_01 =
+        (acc.balance as i128 + (contract.locked_nano as i128 - locking_01 as i128)) as u64;
+
+    cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.unwrap();
+
+    let contract = &user_list_vm_contracts(&channel, &key, false, vm_contract_id).await.unwrap()[0];
+    assert_eq!(contract.locked_nano, locking_01);
+
+    let acc: db::Account = db_conn.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
+    assert_eq!(acc.balance, expected_bal_01);
+
+    let expected_bal_02 =
+        (acc.balance as i128 + (contract.locked_nano as i128 - locking_02 as i128)) as u64;
+
+    req.locked_nano = locking_02;
+
+    cli_client.extend_vm(key.sign_request(req.clone()).unwrap()).await.unwrap();
+
+    let contract = &user_list_vm_contracts(&channel, &key, false, vm_contract_id).await.unwrap()[0];
+    assert_eq!(contract.locked_nano, locking_02);
+
+    let acc: db::Account = db_conn.select((ACCOUNT, key.pubkey.clone())).await.unwrap().unwrap();
+    assert_eq!(acc.balance, expected_bal_02);
+}

tests/mock_data.yaml

@@ -114,6 +114,18 @@ accounts:
     kicked_for: []
     last_kick: 1970-01-01T00:00:00Z
     banned_by: []
+  Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG:
+    balance: 25949200000
+    tmp_locked: 0
+    kicked_for: []
+    last_kick: 1970-01-01T00:00:00Z
+    banned_by: []
+  GmE4JH3bL4NpmzwKCBJemJzRTumJAnbcXLGqce5mREgS:
+    balance: 500000000000
+    tmp_locked: 0
+    kicked_for: []
+    last_kick: 1970-01-01T00:00:00Z
+    banned_by: []
 
 operators:
   BFopWmwcZAMF1h2PFECZNdEucdZfnZZ32p6R9ZaBiVsS:
@@ -151,6 +163,14 @@ operators:
       - 7fujZQeTme52RdXTLmQST5jBgAbvzic5iERtH5EWoYjk
     app_nodes: []
 
+  Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG:
+    escrow: 5499700480000
+    email: "test_mock_extend@operator"
+    banned_users: []
+    vm_nodes:
+      - 8ue3VHMnJg2i8pwTQ6mJtvYuS2kd9n1XLLco8GUPfT95
+    app_nodes: []
+
 vm_nodes:
   - public_key: 7Xw3RxbP5pvfjZ8U6yA3HHVSS9YXjKH5Vkas3JRbQYd9
     operator_wallet: x52w7jARC5erhWWK65VZmjdGXzBK6ZDgfv1A283d8XK
@@ -265,6 +285,22 @@ vm_nodes:
     price: 24000
     reports: {}
     offline_minutes: 0
+  - public_key: 8ue3VHMnJg2i8pwTQ6mJtvYuS2kd9n1XLLco8GUPfT95
+    operator_wallet: Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG
+    country: GB
+    region: England
+    city: London
+    ip: 193.234.17.2
+    avail_mem_mb: 28000
+    avail_vcpus: 24
+    avail_storage_gbs: 1680
+    avail_ipv4: 1
+    avail_ipv6: 0
+    avail_ports: 19999
+    max_ports_per_vm: 10
+    price: 24000
+    reports: {}
+    offline_minutes: 0
 
 vm_contracts:
   - uuid: 958165e3-dea8-407d-8c42-dd17002ef79c
@@ -406,6 +442,24 @@ vm_contracts:
     locked_nano: 12730960000
     collected_at: 2025-04-20T00:34:15.461240342Z
 
+  - uuid: 1d749a81-6c27-a22e-fa0e574-b023-a6afef040fe5
+    hostname: test-extend
+    admin_pubkey: GmE4JH3bL4NpmzwKCBJemJzRTumJAnbcXLGqce5mREgS
+    node_pubkey: 8ue3VHMnJg2i8pwTQ6mJtvYuS2kd9n1XLLco8GUPfT95
+    exposed_ports:
+      - 46393
+    public_ipv4: ""
+    public_ipv6: ""
+    disk_size_gb: 10
+    vcpus: 1
+    memory_mb: 1000
+    kernel_sha: e765e56166ef321b53399b9638584d1279821dbe3d46191c1f66bbaa075e7919
+    dtrfs_sha: d207644ee60d54009b6ecdfb720e2ec251cde31774dd249fcc7435aca0377990
+    created_at: 2025-04-16T20:37:57.176592933Z
+    updated_at: 2025-04-16T20:37:57.176594069Z
+    price_per_unit: 20000
+    locked_nano: 60000
+    collected_at: 2025-04-20T00:34:15.461240342Z
   - uuid: 5af49a71-4c64-a82e-f50e574-b023-b2a0ef0405ed
     hostname: hallow-hobo
     admin_pubkey: 4qFJJJdRrSB9hCn8rrvYTXHLJg371ab36PJmZ4uxHjGQ