From ba872917c244c7e5ff34352f6a9faedf86fcc4fc Mon Sep 17 00:00:00 2001 From: Noor Date: Fri, 21 Mar 2025 18:11:04 +0530 Subject: [PATCH] refactor: reorganize super-cli gRPC method to general module change grpc type import as proto in snp grpc module and updated its related imports --- src/bin/super-detee-cli.rs | 8 ++-- src/config.rs | 3 +- src/general/grpc.rs | 68 +++++++++++++++++++++++++++++++- src/snp/deploy.rs | 10 ++--- src/snp/grpc.rs | 80 ++++---------------------------------- src/snp/injector.rs | 4 +- src/snp/mod.rs | 40 +++++++++---------- src/snp/update.rs | 6 +-- 8 files changed, 109 insertions(+), 110 deletions(-) diff --git a/src/bin/super-detee-cli.rs b/src/bin/super-detee-cli.rs index 67d7baa..bf25a9a 100644 --- a/src/bin/super-detee-cli.rs +++ b/src/bin/super-detee-cli.rs @@ -106,7 +106,7 @@ fn handle_completion(matches: &ArgMatches, mut cmd: Command) { fn handle_account(matches: &ArgMatches) { match matches.subcommand() { Some(("show", _)) => cli_print(Ok(config::Config::get_account_data())), - Some(("list", _)) => match block_on(snp::grpc::admin_list_accounts()) { + Some(("list", _)) => match block_on(general::grpc::admin_list_accounts()) { Ok(accounts) => { for a in accounts { println!("{} {} {}", a.pubkey, a.balance, a.tmp_locked); @@ -123,7 +123,7 @@ fn handle_account(matches: &ArgMatches) { } fn handle_contracts(_matches: &ArgMatches) { - match block_on(snp::grpc::admin_list_contracts()) { + match block_on(general::grpc::admin_list_contracts()) { Ok(contracts) => { for c in contracts { println!( @@ -139,7 +139,7 @@ fn handle_contracts(_matches: &ArgMatches) { fn handle_airdrop(matches: &ArgMatches) { let wallet: String = matches.get_one::("wallet").unwrap().clone(); let tokens = *matches.get_one::("amount").unwrap(); - match block_on(snp::grpc::admin_airdrop(wallet, tokens)) { + match block_on(general::grpc::admin_airdrop(wallet, tokens)) { Ok(()) => println!("Success."), Err(e) => println!("Could not give airdrop due to error: {e:?}"), } @@ -148,7 +148,7 @@ fn handle_airdrop(matches: &ArgMatches) { fn handle_slash(matches: &ArgMatches) { let wallet: String = matches.get_one::("wallet").unwrap().clone(); let tokens = *matches.get_one::("amount").unwrap(); - match block_on(snp::grpc::admin_slash(wallet, tokens)) { + match block_on(general::grpc::admin_slash(wallet, tokens)) { Ok(()) => println!("Success."), Err(e) => println!("Could not slash wallet due to error: {e:?}"), } diff --git a/src/config.rs b/src/config.rs index fb722b7..a0ab8a4 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,3 +1,4 @@ +use crate::{general, utils::block_on}; use ed25519_dalek::SigningKey; use log::{debug, info, warn}; use openssl::bn::BigNum; @@ -339,7 +340,7 @@ impl Config { match Self::get_detee_wallet() { Ok(key) => { account_data.wallet_address = key.to_string(); - match crate::utils::block_on(crate::snp::grpc::get_balance(&key)) { + match block_on(general::grpc::get_balance(&key)) { Ok(account) => { account_data.account_balance = account.balance as f64 / 1_000_000_000.0; account_data.locked_funds = account.tmp_locked as f64 / 1_000_000_000.0; diff --git a/src/general/grpc.rs b/src/general/grpc.rs index 5a27f26..4a8653d 100644 --- a/src/general/grpc.rs +++ b/src/general/grpc.rs @@ -1,6 +1,7 @@ use crate::config::Config; +use crate::snp::grpc::proto::VmContract; use crate::utils::sign_request; -use log::{debug, warn}; +use log::{debug, info, warn}; use tokio_stream::StreamExt; // use tonic::metadata::errors::InvalidMetadataValue; @@ -11,7 +12,8 @@ pub mod proto { use proto::brain_general_cli_client::BrainGeneralCliClient; use proto::{ - BanUserReq, Empty, InspectOperatorResp, KickReq, ListOperatorsResp, Pubkey, RegOperatorReq, + Account, AccountBalance, AirdropReq, BanUserReq, Empty, InspectOperatorResp, KickReq, + ListOperatorsResp, Pubkey, RegOperatorReq, SlashReq, }; #[derive(thiserror::Error, Debug)] @@ -33,6 +35,14 @@ pub enum Error { ConfigError(#[from] crate::config::Error), } +pub async fn get_balance(account: &str) -> Result { + let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; + let response = + client.get_balance(sign_request(Pubkey { pubkey: account.to_string() })?).await?; + log::info!("Received account from brain: {response:?}"); + Ok(response.into_inner()) +} + pub async fn register_operator(escrow: u64, email: String) -> Result<(), Error> { debug!("Connecting to brain to register operator..."); let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; @@ -96,3 +106,57 @@ pub async fn ban_user(user_wallet: String) -> Result<(), Error> { .await?; Ok(()) } + +// super admin + +pub async fn admin_list_accounts() -> Result, Error> { + let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; + let mut accounts = Vec::new(); + let mut grpc_stream = client.list_accounts(sign_request(Empty {})?).await?.into_inner(); + while let Some(stream_update) = grpc_stream.next().await { + match stream_update { + Ok(account) => { + info!("Received account from brain: {account:?}"); + accounts.push(account); + } + Err(e) => { + warn!("Received error instead of contracts: {e:?}"); + } + } + } + debug!("Brain terminated list_contracts stream."); + Ok(accounts) +} + +pub async fn admin_list_contracts() -> Result, Error> { + let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; + let mut contracts = Vec::new(); + let mut grpc_stream = client.list_all_vm_contracts(sign_request(Empty {})?).await?.into_inner(); + while let Some(stream_update) = grpc_stream.next().await { + match stream_update { + Ok(contract) => { + info!("Received contract from brain: {contract:?}"); + contracts.push(contract); + } + Err(e) => { + warn!("Received error instead of contracts: {e:?}"); + } + } + } + debug!("Brain terminated list_contracts stream."); + Ok(contracts) +} + +pub async fn admin_airdrop(pubkey: String, tokens: u64) -> Result<(), Error> { + let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; + let req = sign_request(AirdropReq { pubkey, tokens })?; + let _ = client.airdrop(req).await?; + Ok(()) +} + +pub async fn admin_slash(pubkey: String, tokens: u64) -> Result<(), Error> { + let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; + let req = sign_request(SlashReq { pubkey, tokens })?; + let _ = client.slash(req).await?; + Ok(()) +} diff --git a/src/snp/deploy.rs b/src/snp/deploy.rs index 1f234bb..a9b909b 100644 --- a/src/snp/deploy.rs +++ b/src/snp/deploy.rs @@ -1,5 +1,5 @@ use super::{ - grpc::{self, brain}, + grpc::{self, proto}, injector, Distro, Dtrfs, Error, VmSshArgs, DEFAULT_ARCHLINUX, DEFAULT_DTRFS, }; use crate::config::Config; @@ -106,7 +106,7 @@ impl Request { } // returns node IP and data regarding the new VM - fn send_vm_request(&self) -> Result<(String, brain::NewVmResp), Error> { + fn send_vm_request(&self) -> Result<(String, proto::NewVmResp), Error> { let admin_pubkey = Config::get_detee_wallet()?; let node = self.get_node()?; let (extra_ports, public_ipv4): (Vec, bool) = match &self.ipv4 { @@ -130,7 +130,7 @@ impl Request { self.hours, self.price, ); - let brain_req = brain::NewVmReq { + let brain_req = proto::NewVmReq { uuid: String::new(), hostname: self.hostname.clone(), admin_pubkey, @@ -155,12 +155,12 @@ impl Request { Ok((node.ip, new_vm_resp)) } - pub fn get_node(&self) -> Result { + pub fn get_node(&self) -> Result { let (free_ports, offers_ipv4) = match &self.ipv4 { IPv4Config::PublishPorts(vec) => (vec.len() as u32, false), IPv4Config::PublicIPv4 => (0, true), }; - let filters = brain::VmNodeFilters { + let filters = proto::VmNodeFilters { free_ports, offers_ipv4, offers_ipv6: self.public_ipv6, diff --git a/src/snp/grpc.rs b/src/snp/grpc.rs index f40f677..5d0ebf3 100644 --- a/src/snp/grpc.rs +++ b/src/snp/grpc.rs @@ -1,21 +1,16 @@ -pub mod brain { - // tonic::include_proto!("vm_proto"); - pub use detee_shared::common_proto::*; +pub mod proto { pub use detee_shared::general_proto::*; pub use detee_shared::vm_proto::*; } use crate::config::Config; -use crate::snp::brain::Account; -use crate::snp::brain::AccountBalance; -use crate::snp::brain::Pubkey; -use brain::{ - brain_vm_cli_client::BrainVmCliClient, DeleteVmReq, ExtendVmReq, ListVmContractsReq, NewVmReq, - NewVmResp, ReportNodeReq, UpdateVmReq, UpdateVmResp, VmContract, VmNodeFilters, VmNodeListResp, -}; use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient; use lazy_static::lazy_static; use log::{debug, info, warn}; +use proto::{ + brain_vm_cli_client::BrainVmCliClient, DeleteVmReq, ExtendVmReq, ListVmContractsReq, NewVmReq, + NewVmResp, ReportNodeReq, UpdateVmReq, UpdateVmResp, VmContract, VmNodeFilters, VmNodeListResp, +}; use tokio_stream::StreamExt; use tonic::metadata::errors::InvalidMetadataValue; use tonic::metadata::AsciiMetadataValue; @@ -118,28 +113,6 @@ pub async fn get_node_list(req: VmNodeFilters) -> Result, Er Ok(nodes) } -pub async fn get_balance(account: &str) -> Result { - let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; - let response = - client.get_balance(sign_request(Pubkey { pubkey: account.to_string() })?).await?; - log::info!("Received account from brain: {response:?}"); - Ok(response.into_inner()) -} - -pub async fn admin_airdrop(pubkey: String, tokens: u64) -> Result<(), Error> { - let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; - let req = sign_request(brain::AirdropReq { pubkey, tokens })?; - let _ = client.airdrop(req).await?; - Ok(()) -} - -pub async fn admin_slash(pubkey: String, tokens: u64) -> Result<(), Error> { - let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; - let req = sign_request(brain::SlashReq { pubkey, tokens })?; - let _ = client.slash(req).await?; - Ok(()) -} - pub async fn get_one_node(req: VmNodeFilters) -> Result { let mut client = BrainVmCliClient::connect(Config::get_brain_url()).await?; let response = client.get_one_vm_node(sign_request(req)?).await?; @@ -193,45 +166,6 @@ pub async fn list_contracts(req: ListVmContractsReq) -> Result, Ok(contracts) } -pub async fn admin_list_contracts() -> Result, Error> { - let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; - let mut contracts = Vec::new(); - let mut grpc_stream = - client.list_all_vm_contracts(sign_request(brain::Empty {})?).await?.into_inner(); - while let Some(stream_update) = grpc_stream.next().await { - match stream_update { - Ok(contract) => { - info!("Received contract from brain: {contract:?}"); - contracts.push(contract); - } - Err(e) => { - warn!("Received error instead of contracts: {e:?}"); - } - } - } - debug!("Brain terminated list_contracts stream."); - Ok(contracts) -} - -pub async fn admin_list_accounts() -> Result, Error> { - let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?; - let mut accounts = Vec::new(); - let mut grpc_stream = client.list_accounts(sign_request(brain::Empty {})?).await?.into_inner(); - while let Some(stream_update) = grpc_stream.next().await { - match stream_update { - Ok(account) => { - info!("Received account from brain: {account:?}"); - accounts.push(account); - } - Err(e) => { - warn!("Received error instead of contracts: {e:?}"); - } - } - } - debug!("Brain terminated list_contracts stream."); - Ok(accounts) -} - pub async fn delete_vm(uuid: &str) -> Result<(), Error> { let mut client = BrainVmCliClient::connect(Config::get_brain_url()).await?; let req = DeleteVmReq { uuid: uuid.to_string(), admin_pubkey: Config::get_detee_wallet()? }; @@ -291,8 +225,8 @@ pub async fn update_vm(req: UpdateVmReq) -> Result { } } -pub async fn get_contract_by_uuid(uuid: &str) -> Result { - let req = brain::ListVmContractsReq { +pub async fn get_contract_by_uuid(uuid: &str) -> Result { + let req = ListVmContractsReq { wallet: Config::get_detee_wallet()?, uuid: uuid.to_string(), ..Default::default() diff --git a/src/snp/injector.rs b/src/snp/injector.rs index 122918f..58b2867 100644 --- a/src/snp/injector.rs +++ b/src/snp/injector.rs @@ -1,4 +1,4 @@ -use crate::{config::Config, snp::grpc::brain}; +use crate::{config::Config, snp::grpc::proto}; use log::debug; use std::net::IpAddr; @@ -9,7 +9,7 @@ pub struct Args { pub vcpus: u32, pub kernel: String, pub initrd: String, - pub args: brain::MeasurementArgs, + pub args: proto::MeasurementArgs, } #[derive(thiserror::Error, Debug)] diff --git a/src/snp/mod.rs b/src/snp/mod.rs index 0e146d2..1ed51d5 100644 --- a/src/snp/mod.rs +++ b/src/snp/mod.rs @@ -9,7 +9,7 @@ use crate::{ config::{self, Config}, snp, }; -use grpc::brain; +use grpc::proto; use lazy_static::lazy_static; use serde::{Deserialize, Serialize}; use tabled::Tabled; @@ -53,10 +53,10 @@ impl crate::HumanOutput for VmSshArgs { } } -impl TryFrom for VmSshArgs { +impl TryFrom for VmSshArgs { type Error = Error; - fn try_from(contract: grpc::brain::VmContract) -> Result { + fn try_from(contract: grpc::proto::VmContract) -> Result { let mut args = VmSshArgs { ..Default::default() }; args.uuid = contract.uuid; args.hostname = contract.hostname; @@ -72,7 +72,7 @@ impl TryFrom for VmSshArgs { "This VM does not have a public IP. Getting node IP for node {}", contract.node_pubkey ); - let node = block_on(snp::grpc::get_one_node(brain::VmNodeFilters { + let node = block_on(snp::grpc::get_one_node(proto::VmNodeFilters { node_pubkey: contract.node_pubkey.clone(), ..Default::default() }))?; @@ -200,10 +200,10 @@ fn display_mins(minutes: &u64) -> String { format!("{hours}h {mins}m") } -impl From for VmContract { - fn from(brain_contract: brain::VmContract) -> Self { +impl From for VmContract { + fn from(brain_contract: proto::VmContract) -> Self { let node_pubkey = brain_contract.node_pubkey.clone(); - let location = match block_on(snp::grpc::get_one_node(brain::VmNodeFilters { + let location = match block_on(snp::grpc::get_one_node(proto::VmNodeFilters { node_pubkey: node_pubkey.clone(), ..Default::default() })) { @@ -241,8 +241,8 @@ pub struct TabledVmNode { pub reports: usize, } -impl From for TabledVmNode { - fn from(brain_node: brain::VmNodeListResp) -> Self { +impl From for TabledVmNode { + fn from(brain_node: proto::VmNodeListResp) -> Self { Self { operator: brain_node.operator, location: brain_node.city + ", " + &brain_node.region + ", " + &brain_node.country, @@ -255,7 +255,7 @@ impl From for TabledVmNode { pub fn ssh(uuid: &str, just_print: bool) -> Result { log::info!("Getting VM information about {uuid} from brain..."); - let req = brain::ListVmContractsReq { + let req = proto::ListVmContractsReq { wallet: Config::get_detee_wallet()?, uuid: uuid.to_string(), ..Default::default() @@ -288,9 +288,9 @@ pub fn ssh(uuid: &str, just_print: bool) -> Result { std::process::exit(2); } -pub fn get_one_contract(uuid: &str) -> Result { +pub fn get_one_contract(uuid: &str) -> Result { log::debug!("Getting contract {uuid} from brain..."); - let req = brain::ListVmContractsReq { + let req = proto::ListVmContractsReq { wallet: Config::get_detee_wallet()?, uuid: uuid.to_string(), ..Default::default() @@ -302,9 +302,9 @@ pub fn get_one_contract(uuid: &str) -> Result Result { +pub fn get_node_by_contract(uuid: &str) -> Result { let contract = get_one_contract(uuid)?; - Ok(block_on(snp::grpc::get_one_node(brain::VmNodeFilters { + Ok(block_on(snp::grpc::get_one_node(proto::VmNodeFilters { node_pubkey: contract.node_pubkey, ..Default::default() }))?) @@ -315,7 +315,7 @@ pub fn delete_contract(uuid: &str) -> Result<(), Error> { } pub fn list_contracts(as_operator: bool) -> Result, Error> { - let req = brain::ListVmContractsReq { + let req = proto::ListVmContractsReq { wallet: Config::get_detee_wallet()?, as_operator, ..Default::default() @@ -346,7 +346,7 @@ pub fn append_uuid_list(uuid: &str, hostname: &str) -> Result<(), Error> { Ok(()) } -impl super::HumanOutput for Vec { +impl super::HumanOutput for Vec { fn human_cli_print(&self) { let nodes: Vec = self.iter().map(|n| n.clone().into()).collect(); let style = tabled::settings::Style::rounded(); @@ -356,14 +356,14 @@ impl super::HumanOutput for Vec { } } -pub fn print_nodes() -> Result, Error> { +pub fn print_nodes() -> Result, Error> { log::debug!("This will support flags in the future, but we have only one node atm."); - let req = brain::VmNodeFilters { ..Default::default() }; + let req = proto::VmNodeFilters { ..Default::default() }; Ok(block_on(grpc::get_node_list(req))?) } -pub fn inspect_node(ip: String) -> Result { - let req = brain::VmNodeFilters { ip, ..Default::default() }; +pub fn inspect_node(ip: String) -> Result { + let req = proto::VmNodeFilters { ip, ..Default::default() }; Ok(block_on(grpc::get_one_node(req))?) } diff --git a/src/snp/update.rs b/src/snp/update.rs index 879d672..4a53e1a 100644 --- a/src/snp/update.rs +++ b/src/snp/update.rs @@ -1,5 +1,5 @@ use super::{ - grpc::{self, brain}, + grpc::{self, proto}, injector, Dtrfs, Error, }; use crate::config::Config; @@ -71,12 +71,12 @@ impl Request { } // returns node IP and data regarding the new VM - fn send_update_vm_request(&self, uuid: &str) -> Result { + fn send_update_vm_request(&self, uuid: &str) -> Result { let (kernel_url, kernel_sha, dtrfs_url, dtrfs_sha) = match self.dtrfs.clone() { Some(dtrfs) => (dtrfs.kernel_url, dtrfs.kernel_sha, dtrfs.dtrfs_url, dtrfs.dtrfs_sha), None => (String::new(), String::new(), String::new(), String::new()), }; - Ok(block_on(grpc::update_vm(brain::UpdateVmReq { + Ok(block_on(grpc::update_vm(proto::UpdateVmReq { uuid: uuid.to_string(), admin_pubkey: Config::get_detee_wallet()?, disk_size_gb: self.disk_size_gb,