diff --git a/src/constants.rs b/src/constants.rs index 94df60b..ae7b5e5 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -23,6 +23,8 @@ pub static ADMIN_ACCOUNTS: LazyLock> = LazyLock::new(|| { pub const OLD_BRAIN_DATA_PATH: &str = "./saved_data.yaml"; pub const ACCOUNT: &str = "account"; +pub const KICK: &str = "kick"; + pub const VM_NODE: &str = "vm_node"; pub const ACTIVE_VM: &str = "active_vm"; pub const VM_UPDATE_EVENT: &str = "vm_update_event"; diff --git a/src/db/app.rs b/src/db/app.rs index 478d87d..b4dade4 100644 --- a/src/db/app.rs +++ b/src/db/app.rs @@ -12,7 +12,7 @@ use surrealdb::sql::Datetime; use surrealdb::{Notification, RecordId, Surreal}; use tokio_stream::StreamExt; -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct AppNode { pub id: RecordId, pub operator: RecordId, @@ -32,8 +32,9 @@ pub struct AppNode { impl AppNode { pub async fn register(self, db: &Surreal) -> Result { db::Account::get_or_create(db, &self.operator.key().to_string()).await?; - let app_node: Option = db.upsert(self.id.clone()).content(self).await?; - app_node.ok_or(Error::FailedToCreateDBEntry) + let app_node_id = self.id.clone(); + let app_node: Option = db.upsert(app_node_id.clone()).content(self).await?; + app_node.ok_or(Error::FailedToCreateDBEntry(format!("{APP_NODE}:{app_node_id}"))) } } @@ -54,7 +55,7 @@ impl From for AppDaemonMsg { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct NewAppReq { pub id: RecordId, #[serde(rename = "in")] @@ -164,7 +165,7 @@ impl AppNodeWithReports { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct ActiveApp { pub id: RecordId, #[serde(rename = "in")] @@ -210,6 +211,15 @@ impl From for DeletedApp { } impl ActiveApp { + pub fn price_per_minute(&self) -> u64 { + (self.total_units() * self.price_per_unit as f64) as u64 + } + + fn total_units(&self) -> f64 { + // TODO: Optimize this based on price of hardware. + (self.vcpus as f64 * 5f64) + (self.memory_mb as f64 / 200f64) + (self.disk_size_gb as f64) + } + pub async fn activate(db: &Surreal, id: &str) -> Result<(), Error> { let new_app_req = match NewAppReq::get(db, id).await? { Some(r) => r, @@ -293,7 +303,7 @@ impl ActiveApp { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct ActiveAppWithNode { pub id: RecordId, #[serde(rename = "in")] @@ -315,6 +325,29 @@ pub struct ActiveAppWithNode { pub hratls_pubkey: String, } +impl From for ActiveApp { + fn from(val: ActiveAppWithNode) -> Self { + Self { + id: val.id, + admin: val.admin, + app_node: val.app_node.id, + app_name: val.app_name, + mapped_ports: val.mapped_ports, + host_ipv4: val.host_ipv4, + vcpus: val.vcpus, + memory_mb: val.memory_mb, + disk_size_gb: val.disk_size_gb, + created_at: val.created_at, + price_per_unit: val.price_per_unit, + locked_nano: val.locked_nano, + collected_at: val.collected_at, + mr_enclave: val.mr_enclave, + package_url: val.package_url, + hratls_pubkey: val.hratls_pubkey, + } + } +} + impl ActiveAppWithNode { pub async fn get_by_uuid(db: &Surreal, uuid: &str) -> Result, Error> { let contract: Option = diff --git a/src/db/general.rs b/src/db/general.rs index 86c49df..4ffd49c 100644 --- a/src/db/general.rs +++ b/src/db/general.rs @@ -1,7 +1,6 @@ -use crate::constants::{ACCOUNT, MIN_ESCROW, TOKEN_DECIMAL}; -use crate::db::prelude::*; - use super::Error; +use crate::constants::{ACCOUNT, KICK, MIN_ESCROW, TOKEN_DECIMAL}; +use crate::db::prelude::*; use crate::old_brain; use serde::{Deserialize, Serialize}; use surrealdb::engine::remote::ws::Client; @@ -37,7 +36,7 @@ impl Account { Some(account) => Ok(account), None => { let account: Option = db.create(id).await?; - account.ok_or(Error::FailedToCreateDBEntry) + account.ok_or(Error::FailedToCreateDBEntry(ACCOUNT.to_string())) } } } @@ -147,6 +146,25 @@ pub struct Kick { created_at: Datetime, reason: String, contract: RecordId, + node: RecordId, +} + +impl Kick { + pub async fn kicked_in_a_day(db: &Surreal, account: &str) -> Result, Error> { + let yesterday = chrono::Utc::now() - chrono::Duration::days(1); + let mut result = db + .query(format!( + "select * from {KICK} where in = {ACCOUNT}:{account} and created_at > {yesterday};" + )) + .await?; + let kicks: Vec = result.take(0)?; + Ok(kicks) + } + + pub async fn submit(self, db: &Surreal) -> Result<(), Error> { + let _: Vec = db.insert(KICK).relation(self).await?; + Ok(()) + } } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -185,7 +203,7 @@ impl Report { /// This is the operator obtained from the DB, /// however the relation is defined using OperatorRelation -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct Operator { pub account: RecordId, pub app_nodes: u64, @@ -258,3 +276,107 @@ impl Operator { Ok((operator, vm_nodes, app_nodes)) } } + +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 (node_operator, admin, contract_id, node_id, collected_at, price_per_mint, is_vm) = + 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.to_string(), + active_vm.admin.key().to_string(), + active_vm.id, + active_vm.vm_node.id, + active_vm.collected_at, + price_per_minute, + true, + ) + } 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.to_string(), + active_app.admin.key().to_string(), + active_app.id, + active_app.app_node.id, + active_app.collected_at, + price_per_minute, + false, + ) + } else { + return Err(Error::ContractNotFound); + }; + + if node_operator != operator_wallet { + 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; + + if minutes_to_refund > one_week_minute { + minutes_to_refund = one_week_minute; + } + + let mut refund_amount = minutes_to_refund * price_per_mint; + + if !Kick::kicked_in_a_day(db, &admin).await?.is_empty() { + refund_amount = 0; + } + + let mut operator_account = Account::get(db, &node_operator).await?; + let mut admin_account = Account::get(db, &admin).await?; + + if operator_account.escrow < refund_amount { + refund_amount = operator_account.escrow; + } + + log::debug!( + "Removing {refund_amount} escrow from {} and giving it to {}", + node_operator, + admin + ); + + admin_account.balance = admin_account.balance.saturating_add(refund_amount); + operator_account.escrow = operator_account.escrow.saturating_sub(refund_amount); + + let kick = Kick { + id: RecordId::from((KICK, contract_uuid)), + from_account: operator_account.id.clone(), + to_account: admin_account.id.clone(), + created_at: Datetime::default(), + reason: reason.to_string(), + contract: contract_id.clone(), + node: node_id, + }; + + kick.submit(db).await?; + operator_account.save(db).await?; + admin_account.save(db).await?; + + let contract_id = contract_id.to_string(); + + if is_vm && !ActiveVm::delete(db, &contract_id).await? { + return Err(Error::FailedToDeleteContract(format!("vm:{contract_id}"))); + } else if !is_vm && !ActiveApp::delete(db, &contract_id).await? { + return Err(Error::FailedToDeleteContract(format!("app:{contract_id}"))); + } + + Ok(refund_amount) + } +} diff --git a/src/db/mod.rs b/src/db/mod.rs index eabb36c..2129d0e 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -24,8 +24,8 @@ pub enum Error { StdIo(#[from] std::io::Error), #[error(transparent)] TimeOut(#[from] tokio::time::error::Elapsed), - #[error("Failed to create account")] - FailedToCreateDBEntry, + #[error("Failed to create {0}")] + FailedToCreateDBEntry(String), #[error("Unknown Table: {0}")] UnknownTable(String), #[error("Daemon channel got closed: {0}")] @@ -36,6 +36,12 @@ pub enum Error { MinimalEscrow, #[error("Insufficient funds, deposit more tokens")] InsufficientFunds, + #[error("Contract not found")] + ContractNotFound, + #[error("Access denied")] + AccessDenied, + #[error("Failed to delete contract {0}")] + FailedToDeleteContract(String), } pub mod prelude { diff --git a/src/db/vm.rs b/src/db/vm.rs index 82acfc2..ff5e3fe 100644 --- a/src/db/vm.rs +++ b/src/db/vm.rs @@ -728,6 +728,29 @@ pub struct ActiveVmWithNode { pub collected_at: Datetime, } +impl From for ActiveVm { + fn from(val: ActiveVmWithNode) -> Self { + Self { + id: val.id, + admin: val.admin, + vm_node: val.vm_node.id, + hostname: val.hostname, + mapped_ports: val.mapped_ports, + public_ipv4: val.public_ipv4, + public_ipv6: val.public_ipv6, + disk_size_gb: val.disk_size_gb, + vcpus: val.vcpus, + memory_mb: val.memory_mb, + dtrfs_sha: val.dtrfs_sha, + kernel_sha: val.kernel_sha, + created_at: val.created_at, + price_per_unit: val.price_per_unit, + locked_nano: val.locked_nano, + collected_at: val.collected_at, + } + } +} + impl ActiveVmWithNode { pub async fn get_by_uuid(db: &Surreal, uuid: &str) -> Result, Error> { let contract: Option = diff --git a/src/grpc/general.rs b/src/grpc/general.rs index 2976314..35e15b9 100644 --- a/src/grpc/general.rs +++ b/src/grpc/general.rs @@ -125,13 +125,34 @@ impl BrainGeneralCli for GeneralCliServer { } } - async fn kick_contract(&self, _req: Request) -> Result, Status> { - todo!(); - // let req = check_sig_from_req(req)?; - // match self.data.kick_contract(&req.operator_wallet, &req.contract_uuid, &req.reason).await { - // Ok(nano_lp) => Ok(Response::new(KickResp { nano_lp })), - // Err(e) => Err(Status::permission_denied(e.to_string())), - // } + 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 + { + Ok(nano_lp) => Ok(Response::new(KickResp { nano_lp })), + Err(e) + if matches!( + e, + db::Error::ContractNotFound + | db::Error::AccessDenied + | db::Error::FailedToDeleteContract(_) + ) => + { + Err(Status::failed_precondition(e.to_string())) + } + Err(e) => { + log::info!("Failed to kick contract: {e:?}"); + Err(Status::unknown( + "Unknown error. Please try again or contact the DeTEE devs team.", + )) + } + } } async fn ban_user(&self, _req: Request) -> Result, Status> { diff --git a/surql/tables.sql b/surql/tables.sql index 278c7a1..4e9c02e 100644 --- a/surql/tables.sql +++ b/surql/tables.sql @@ -144,6 +144,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 TABLE report TYPE RELATION FROM account TO vm_node|app_node; DEFINE FIELD created_at ON TABLE report TYPE datetime; diff --git a/tests/common/vm_daemon_utils.rs b/tests/common/vm_daemon_utils.rs index 58a9446..2842b85 100644 --- a/tests/common/vm_daemon_utils.rs +++ b/tests/common/vm_daemon_utils.rs @@ -37,7 +37,7 @@ pub async fn register_vm_node( client: &mut BrainVmDaemonClient, key: &Key, operator_wallet: &str, -) -> Result> { +) -> Result> { log::info!("Registering vm_node: {}", key.pubkey); let node_pubkey = key.pubkey.clone();