From 5ccd432566c6b95d7c764a2ef67b8a793f8775a2 Mon Sep 17 00:00:00 2001 From: Noor Date: Tue, 27 May 2025 21:03:53 +0530 Subject: [PATCH] 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 | 4 ++ src/db/vm.rs | 73 +++++++++++++++++++++++++++++ src/grpc/vm.rs | 29 ++++++++---- tests/common/vm_cli_utils.rs | 35 ++++++++++++++ tests/grpc_vm_cli_test.rs | 89 ++++++++++++++++++++++++++++++++++-- tests/mock_data.yaml | 54 ++++++++++++++++++++++ 6 files changed, 272 insertions(+), 12 deletions(-) diff --git a/src/db/mod.rs b/src/db/mod.rs index c84689a..8cd5ce3 100644 --- a/src/db/mod.rs +++ b/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 { diff --git a/src/db/vm.rs b/src/db/vm.rs index 80d14c0..eeb449c 100644 --- a/src/db/vm.rs +++ b/src/db/vm.rs @@ -196,6 +196,7 @@ impl NewVmReq { error: String, } let _: Option = 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, + 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 ' + $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)] diff --git a/src/grpc/vm.rs b/src/grpc/vm.rs index 11c6ad0..9b5e6aa 100644 --- a/src/grpc/vm.rs +++ b/src/grpc/vm.rs @@ -313,15 +313,26 @@ impl BrainVmCli for VmCliServer { } async fn extend_vm(&self, req: Request) -> Result, Status> { - let _req = check_sig_from_req(req)?; - todo!(); - // match self - // .data - // .extend_vm_contract_time(&req.uuid, &req.admin_pubkey, req.locked_nano) - // { - // Ok(()) => Ok(Response::new(Empty {})), - // Err(e) => Err(Status::unknown(format!("Could not extend contract: {e}"))), - // } + let req = check_sig_from_req(req)?; + match db::ActiveVm::extend_time(&self.db, &req.uuid, &req.admin_pubkey, req.locked_nano) + .await + { + Ok(()) => Ok(Response::new(Empty {})), + Err(e) + if matches!( + 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) -> Result, Status> { diff --git a/tests/common/vm_cli_utils.rs b/tests/common/vm_cli_utils.rs index 4e4b31a..7927410 100644 --- a/tests/common/vm_cli_utils.rs +++ b/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> { + 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) +} diff --git a/tests/grpc_vm_cli_test.rs b/tests/grpc_vm_cli_test.rs index b6d4c71..7338238 100644 --- a/tests/grpc_vm_cli_test.rs +++ b/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); +} diff --git a/tests/mock_data.yaml b/tests/mock_data.yaml index 1bae3f9..08dbdf9 100644 --- a/tests/mock_data.yaml +++ b/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