use crate::old_brain; use serde::{Deserialize, Serialize}; use std::sync::LazyLock; use surrealdb::{ engine::remote::ws::{Client, Ws}, opt::auth::Root, sql::Datetime, RecordId, Surreal, }; static DB: LazyLock> = LazyLock::new(Surreal::init); pub const ACCOUNT: &str = "account"; pub const VM_CONTRACT: &str = "vm_contract"; pub const VM_NODE: &str = "vm_node"; #[derive(thiserror::Error, Debug)] pub enum Error { #[error(transparent)] DataBase(#[from] surrealdb::Error), } pub async fn init() -> surrealdb::Result<()> { DB.connect::("localhost:8000").await?; // Sign in to the server DB.signin(Root { username: "root", password: "root" }).await?; DB.use_ns("brain").use_db("migration").await?; Ok(()) } pub async fn migration0(old_data: &old_brain::BrainData) -> surrealdb::Result<()> { let accounts: Vec = old_data.into(); let vm_nodes: Vec = old_data.into(); let app_nodes: Vec = old_data.into(); let vm_contracts: Vec = old_data.into(); init().await?; println!("Inserting accounts..."); let _: Vec = DB.insert(()).content(accounts).await?; println!("Inserting vm nodes..."); let _: Vec = DB.insert(()).content(vm_nodes).await?; println!("Inserting app nodes..."); let _: Vec = DB.insert(()).content(app_nodes).await?; println!("Inserting vm contracts..."); let _: Vec = DB.insert("vm_contract").relation(vm_contracts).await?; Ok(()) } #[derive(Debug, Serialize, Deserialize)] pub struct Account { pub id: RecordId, pub balance: u64, pub tmp_locked: u64, pub escrow: u64, pub email: String, } impl Account { pub async fn get(address: &str) -> Result { let id = (ACCOUNT, address); let account: Option = DB.select(id).await?; let account = match account { Some(account) => account, None => { Self { id: id.into(), balance: 0, tmp_locked: 0, escrow: 0, email: String::new() } } }; Ok(account) } pub async fn airdrop(account: &str, tokens: u64) -> Result<(), Error> { let tokens = tokens.saturating_mul(1_000_000_000); let _ = DB .query(format!("upsert account:{account} SET balance = (balance || 0) + {tokens};")) .await?; Ok(()) } } #[derive(Debug, Serialize, Deserialize)] pub struct VmNode { pub id: RecordId, pub operator: RecordId, pub country: String, pub region: String, pub city: String, pub ip: String, pub avail_mem_mb: u32, pub avail_vcpus: u32, pub avail_storage_gbs: u32, pub avail_ipv4: u32, pub avail_ipv6: u32, pub avail_ports: u32, pub max_ports_per_vm: u32, pub price: u64, pub offline_minutes: u64, } impl VmNode { pub async fn register(self) -> Result<(), Error> { let _: Option = DB.upsert(self.id.clone()).content(self).await?; Ok(()) } } #[derive(Debug, Serialize, Deserialize)] pub struct VmNodeWithReports { pub id: RecordId, pub operator: RecordId, pub country: String, pub region: String, pub city: String, pub ip: String, pub avail_mem_mb: u32, pub avail_vcpus: u32, pub avail_storage_gbs: u32, pub avail_ipv4: u32, pub avail_ipv6: u32, pub avail_ports: u32, pub max_ports_per_vm: u32, pub price: u64, pub offline_minutes: u64, pub reports: Vec, } #[derive(Debug, Serialize, Deserialize)] pub struct VmContract { pub id: RecordId, #[serde(rename = "in")] pub admin: RecordId, #[serde(rename = "out")] pub vm_node: RecordId, pub state: String, pub hostname: String, pub mapped_ports: Vec<(u32, u32)>, pub public_ipv4: String, pub public_ipv6: String, pub disk_size_gb: u32, pub vcpus: u32, pub memory_mb: u32, pub dtrfs_sha: String, pub kernel_sha: String, pub created_at: Datetime, pub updated_at: Datetime, pub price_per_unit: u64, pub locked_nano: u64, pub collected_at: Datetime, } impl VmContract { /// total hardware units of this VM fn total_units(&self) -> u64 { // TODO: Optimize this based on price of hardware. // I tried, but this can be done better. // Storage cost should also be based on tier (self.vcpus as u64 * 10) + ((self.memory_mb + 256) as u64 / 200) + (self.disk_size_gb as u64 / 10) + (!self.public_ipv4.is_empty() as u64 * 10) } /// Returns price per minute in nanoLP pub fn price_per_minute(&self) -> u64 { self.total_units() * self.price_per_unit } } impl VmContract {} #[derive(Debug, Serialize, Deserialize)] pub struct VmContractWithNode { pub id: RecordId, #[serde(rename = "in")] pub admin: RecordId, #[serde(rename = "out")] pub vm_node: VmNode, pub state: String, pub hostname: String, pub mapped_ports: Vec<(u32, u32)>, pub public_ipv4: String, pub public_ipv6: String, pub disk_size_gb: u32, pub vcpus: u32, pub memory_mb: u32, pub dtrfs_sha: String, pub kernel_sha: String, pub created_at: Datetime, pub updated_at: Datetime, pub price_per_unit: u64, pub locked_nano: u64, pub collected_at: Datetime, } impl VmContractWithNode { pub async fn get_by_uuid(uuid: &str) -> Result, Error> { let contract: Option = DB.query(format!("select * from {VM_CONTRACT}:{uuid} fetch out;")).await?.take(0)?; Ok(contract) } pub async fn list_by_admin(admin: &str) -> Result, Error> { let mut result = DB .query(format!("select * from {VM_CONTRACT} where in = {ACCOUNT}:{admin} fetch out;")) .await?; let contracts: Vec = result.take(0)?; Ok(contracts) } pub async fn list_by_node(admin: &str) -> Result, Error> { let mut result = DB .query(format!("select * from {VM_CONTRACT} where out = {VM_NODE}:{admin} fetch out;")) .await?; let contracts: Vec = result.take(0)?; Ok(contracts) } pub async fn list_by_operator(operator: &str) -> Result, Error> { let mut result = DB .query(format!( "select (select * from ->operator->vm_node<-vm_contract fetch out) as contracts from {ACCOUNT}:{operator};" )) .await?; #[derive(Deserialize)] struct Wrapper { contracts: Vec, } let c: Option = result.take(0)?; match c { Some(c) => Ok(c.contracts), None => Ok(Vec::new()), } } /// total hardware units of this VM fn total_units(&self) -> u64 { // TODO: Optimize this based on price of hardware. // I tried, but this can be done better. // Storage cost should also be based on tier (self.vcpus as u64 * 10) + ((self.memory_mb + 256) as u64 / 200) + (self.disk_size_gb as u64 / 10) + (!self.public_ipv4.is_empty() as u64 * 10) } /// Returns price per minute in nanoLP pub fn price_per_minute(&self) -> u64 { self.total_units() * self.price_per_unit } } #[derive(Debug, Serialize, Deserialize)] pub struct AppNode { pub id: RecordId, pub operator: RecordId, pub country: String, pub region: String, pub city: String, pub ip: String, pub avail_mem_mb: u32, pub avail_vcpus: u32, pub avail_storage_gbs: u32, pub avail_ports: u32, pub max_ports_per_app: u32, pub price: u64, pub offline_minutes: u64, } #[derive(Debug, Serialize, Deserialize)] pub struct AppNodeWithReports { pub id: RecordId, pub operator: RecordId, pub country: String, pub region: String, pub city: String, pub ip: String, pub avail_mem_mb: u32, pub avail_vcpus: u32, pub avail_storage_gbs: u32, pub avail_ports: u32, pub max_ports_per_app: u32, pub price: u64, pub offline_minutes: u64, pub reports: Vec, } #[derive(Debug, Serialize, Deserialize)] pub struct AppContract { id: RecordId, #[serde(rename = "in")] admin: RecordId, #[serde(rename = "out")] app_node: RecordId, state: String, app_name: String, mapped_ports: Vec<(u64, u64)>, host_ipv4: String, vcpus: u64, memory_mb: u64, disk_size_gb: u64, created_at: Datetime, updated_at: Datetime, price_per_unit: u64, locked_nano: u64, collected_at: Datetime, mr_enclave: String, package_url: String, hratls_pubkey: String, } #[derive(Debug, Serialize, Deserialize)] pub struct Ban { id: RecordId, #[serde(rename = "in")] from_account: RecordId, #[serde(rename = "out")] to_account: RecordId, created_at: Datetime, } #[derive(Debug, Serialize, Deserialize)] pub struct Kick { id: RecordId, #[serde(rename = "in")] from_account: RecordId, #[serde(rename = "out")] to_account: RecordId, created_at: Datetime, reason: String, contract: RecordId, } #[derive(Debug, Serialize, Deserialize)] pub struct Report { #[serde(rename = "in")] from_account: RecordId, #[serde(rename = "out")] to_node: RecordId, created_at: Datetime, pub reason: String, } impl Report { // TODO: test this functionality and remove this comment pub async fn create( from_account: RecordId, to_node: RecordId, reason: String, ) -> Result<(), Error> { let _: Vec = DB .insert("report") .relation(Report { from_account, to_node, created_at: Datetime::default(), reason }) .await?; Ok(()) } } /// This is the operator obtained from the DB, /// however the relation is defined using OperatorRelation #[derive(Debug, Serialize, Deserialize)] pub struct Operator { pub account: RecordId, pub app_nodes: u64, pub vm_nodes: u64, pub email: String, pub escrow: u64, pub reports: u64, } impl Operator { pub async fn list() -> Result, Error> { let mut result = DB .query(format!( "array::distinct(array::flatten( [ (select operator from vm_node group by operator).operator, (select operator from app_node group by operator).operator ]));" )) .await?; let operator_accounts: Vec = result.take(0)?; let mut operators: Vec = Vec::new(); for account in operator_accounts.iter() { if let Some(operator) = Self::inspect(&account.key().to_string()).await? { operators.push(operator); } } Ok(operators) } pub async fn inspect(account: &str) -> Result, Error> { let mut result = DB .query(format!( "$vm_nodes = (select id from vm_node where operator = account:{account}).id; $app_nodes = (select id from app_node where operator = account:{account}).id; select *, id as account, email, escrow, $vm_nodes.len() as vm_nodes, $app_nodes.len() as app_nodes, (select id from report where $vm_nodes contains out).len() + (select id from report where $app_nodes contains out).len() as reports from account where id = account:{account};" )) .await?; let operator: Option = result.take(2)?; Ok(operator) } pub async fn inspect_nodes( account: &str, ) -> Result<(Option, Vec, Vec), Error> { let operator = Self::inspect(account).await?; let mut result = DB .query(format!( "select *, operator, <-report.* as reports from vm_node where operator = account:{account};" )) .query(format!( "select *, operator, <-report.* as reports from app_node where operator = account:{account};" )) .await?; let vm_nodes: Vec = result.take(0)?; let app_nodes: Vec = result.take(1)?; Ok((operator, vm_nodes, app_nodes)) } } // TODO: delete all of these From implementation after migration 0 gets executed impl From<&old_brain::BrainData> for Vec { fn from(old_data: &old_brain::BrainData) -> Self { let mut nodes = Vec::new(); for old_node in old_data.vm_nodes.iter() { nodes.push(VmNode { id: RecordId::from((VM_NODE, old_node.public_key.clone())), operator: RecordId::from((ACCOUNT, old_node.operator_wallet.clone())), country: old_node.country.clone(), region: old_node.region.clone(), city: old_node.city.clone(), ip: old_node.ip.clone(), avail_mem_mb: old_node.avail_mem_mb, avail_vcpus: old_node.avail_vcpus, avail_storage_gbs: old_node.avail_storage_gbs, avail_ipv4: old_node.avail_ipv4, avail_ipv6: old_node.avail_ipv6, avail_ports: old_node.avail_ports, max_ports_per_vm: old_node.max_ports_per_vm, price: old_node.price, offline_minutes: old_node.offline_minutes, }); } nodes } } impl From<&old_brain::BrainData> for Vec { fn from(old_data: &old_brain::BrainData) -> Self { let mut contracts = Vec::new(); for old_c in old_data.vm_contracts.iter() { let mut mapped_ports = Vec::new(); for port in old_c.exposed_ports.iter() { mapped_ports.push((*port, 8080 as u32)); } contracts.push(VmContract { id: RecordId::from(( VM_CONTRACT, old_c.node_pubkey.chars().take(20).collect::() + &old_c.uuid.replace("-", "").chars().take(20).collect::(), )), admin: RecordId::from((ACCOUNT, old_c.admin_pubkey.clone())), vm_node: RecordId::from((VM_NODE, old_c.node_pubkey.clone())), state: "active".to_string(), hostname: old_c.hostname.clone(), mapped_ports, public_ipv4: old_c.public_ipv4.clone(), public_ipv6: old_c.public_ipv6.clone(), disk_size_gb: old_c.disk_size_gb, vcpus: old_c.vcpus, memory_mb: old_c.memory_mb, dtrfs_sha: old_c.dtrfs_sha.clone(), kernel_sha: old_c.kernel_sha.clone(), price_per_unit: old_c.price_per_unit, locked_nano: old_c.locked_nano, created_at: old_c.created_at.into(), updated_at: old_c.updated_at.into(), collected_at: old_c.collected_at.into(), }); } contracts } } impl From<&old_brain::BrainData> for Vec { fn from(old_data: &old_brain::BrainData) -> Self { let mut nodes = Vec::new(); for old_node in old_data.app_nodes.iter() { nodes.push(AppNode { id: RecordId::from(("app_node", old_node.node_pubkey.clone())), operator: RecordId::from((ACCOUNT, old_node.operator_wallet.clone())), country: old_node.country.clone(), region: old_node.region.clone(), city: old_node.city.clone(), ip: old_node.ip.clone(), avail_mem_mb: old_node.avail_mem_mb, avail_vcpus: old_node.avail_vcpus, avail_storage_gbs: old_node.avail_storage_mb, avail_ports: old_node.avail_no_of_port, max_ports_per_app: old_node.max_ports_per_app, price: old_node.price, offline_minutes: old_node.offline_minutes, }); } nodes } } impl From<&old_brain::BrainData> for Vec { fn from(old_data: &old_brain::BrainData) -> Self { let mut accounts = Vec::new(); for old_account in old_data.accounts.iter() { let mut a = Account { id: RecordId::from(("account", old_account.key())), balance: old_account.value().balance, tmp_locked: old_account.value().tmp_locked, escrow: 0, email: String::new(), }; if let Some(operator) = old_data.operators.get(old_account.key()) { a.escrow = operator.escrow; a.email = operator.email.clone(); } accounts.push(a); } accounts } }