From 76c88810a6aa5dc2a351a0c5e4024aae1c754fc0 Mon Sep 17 00:00:00 2001 From: ghe0 Date: Sun, 9 Feb 2025 00:30:01 +0200 Subject: [PATCH] rename structs so that they say "VM" --- build.rs | 2 +- src/grpc.rs | 50 +++++++++++++++++++++---------------------- src/main.rs | 18 ++++++++-------- snp.proto => vm.proto | 50 +++++++++++++++++++++++++++---------------- 4 files changed, 67 insertions(+), 53 deletions(-) rename snp.proto => vm.proto (79%) diff --git a/build.rs b/build.rs index d1b9140..469b7d1 100644 --- a/build.rs +++ b/build.rs @@ -1,6 +1,6 @@ fn main() { tonic_build::configure() .build_server(true) - .compile_protos(&["snp.proto"], &["proto"]) + .compile_protos(&["vm.proto"], &["proto"]) .unwrap_or_else(|e| panic!("Failed to compile protos {:?}", e)); } diff --git a/src/grpc.rs b/src/grpc.rs index ccd1880..71cd4fa 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -1,8 +1,8 @@ use crate::global::*; -use crate::snp_proto::DaemonMessage; +use crate::snp_proto::VmDaemonMessage; use anyhow::Result; use log::{debug, info, warn}; -use snp_proto::{brain_daemon_client::BrainDaemonClient, BrainMessage, Contract, RegisterNodeReq}; +use snp_proto::{brain_vm_daemon_client::BrainVmDaemonClient, BrainVmMessage, VmContract, RegisterVmNodeReq}; use tokio::{ sync::mpsc::{Receiver, Sender}, task::JoinSet, @@ -11,34 +11,34 @@ use tokio_stream::{wrappers::ReceiverStream, StreamExt}; use tonic::transport::Channel; pub mod snp_proto { - tonic::include_proto!("snp_proto"); + tonic::include_proto!("vm_proto"); } -impl From for snp_proto::DaemonMessage { +impl From for snp_proto::VmDaemonMessage { fn from(value: snp_proto::NewVmResp) -> Self { - snp_proto::DaemonMessage { msg: Some(snp_proto::daemon_message::Msg::NewVmResp(value)) } + snp_proto::VmDaemonMessage { msg: Some(snp_proto::vm_daemon_message::Msg::NewVmResp(value)) } } } -impl From for snp_proto::DaemonMessage { +impl From for snp_proto::VmDaemonMessage { fn from(value: snp_proto::UpdateVmResp) -> Self { - snp_proto::DaemonMessage { msg: Some(snp_proto::daemon_message::Msg::UpdateVmResp(value)) } + snp_proto::VmDaemonMessage { msg: Some(snp_proto::vm_daemon_message::Msg::UpdateVmResp(value)) } } } -impl From for snp_proto::DaemonMessage { - fn from(value: snp_proto::NodeResources) -> Self { - snp_proto::DaemonMessage { msg: Some(snp_proto::daemon_message::Msg::NodeResources(value)) } +impl From for snp_proto::VmDaemonMessage { + fn from(value: snp_proto::VmNodeResources) -> Self { + snp_proto::VmDaemonMessage { msg: Some(snp_proto::vm_daemon_message::Msg::VmNodeResources(value)) } } } -pub async fn register_node(config: &crate::config::Config) -> Result> { +pub async fn register_node(config: &crate::config::Config) -> Result> { use tonic::metadata::AsciiMetadataValue; use tonic::Request; - let mut client = BrainDaemonClient::connect(config.brain_url.clone()).await?; + let mut client = BrainVmDaemonClient::connect(config.brain_url.clone()).await?; debug!("Starting node registration..."); let ip_info = IP_INFO.clone(); - let req = RegisterNodeReq { + let req = RegisterVmNodeReq { node_pubkey: PUBLIC_KEY.clone(), owner_pubkey: config.owner_wallet.clone(), main_ip: ip_info.ip, @@ -60,7 +60,7 @@ pub async fn register_node(config: &crate::config::Config) -> Result { @@ -85,9 +85,9 @@ fn sign_stream_auth(contracts: Vec) -> Result, + mut client: BrainVmDaemonClient, contracts: Vec, - tx: Sender, + tx: Sender, ) -> Result<()> { debug!("starting to listen for messages from brain"); let mut grpc_stream = client.brain_messages(sign_stream_auth(contracts)?).await?.into_inner(); @@ -107,15 +107,15 @@ async fn receive_messages( } async fn send_messages( - mut client: BrainDaemonClient, + mut client: BrainVmDaemonClient, contracts: Vec, - rx: Receiver, - tx: Sender, + rx: Receiver, + tx: Sender, ) -> Result<()> { debug!("starting daemon message stream to brain"); let rx_stream = ReceiverStream::new(rx); - tx.send(DaemonMessage { - msg: Some(snp_proto::daemon_message::Msg::Auth(sign_stream_auth(contracts)?)), + tx.send(VmDaemonMessage { + msg: Some(snp_proto::vm_daemon_message::Msg::Auth(sign_stream_auth(contracts)?)), }) .await?; client.daemon_messages(rx_stream).await?; @@ -126,13 +126,13 @@ async fn send_messages( pub struct ConnectionData { pub contracts: Vec, pub brain_url: String, - pub brain_msg_tx: Sender, - pub daemon_msg_rx: Receiver, - pub daemon_msg_tx: Sender, + pub brain_msg_tx: Sender, + pub daemon_msg_rx: Receiver, + pub daemon_msg_tx: Sender, } pub async fn connect_and_run(cd: ConnectionData) -> Result<()> { - let client = BrainDaemonClient::connect(cd.brain_url).await?; + let client = BrainVmDaemonClient::connect(cd.brain_url).await?; let mut streaming_tasks = JoinSet::new(); streaming_tasks.spawn(receive_messages(client.clone(), cd.contracts.clone(), cd.brain_msg_tx)); diff --git a/src/main.rs b/src/main.rs index 5f468f4..35da7aa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,8 +14,8 @@ use tokio::{ #[allow(dead_code)] struct VMHandler { - receiver: Receiver, - sender: Sender, + receiver: Receiver, + sender: Sender, config: Config, res: state::Resources, } @@ -23,8 +23,8 @@ struct VMHandler { #[allow(dead_code)] impl VMHandler { fn new( - receiver: Receiver, - sender: Sender, + receiver: Receiver, + sender: Sender, ) -> Self { let config = match Config::load_from_disk(DAEMON_CONFIG_PATH) { Ok(config) => config, @@ -68,7 +68,7 @@ impl VMHandler { } } let avail_storage_gb = avail_storage_gb as u32; - let res = snp_proto::NodeResources { + let res = snp_proto::VmNodeResources { node_pubkey: PUBLIC_KEY.clone(), avail_ports: (self.config.public_port_range.len() - self.res.reserved_ports.len()) as u32, @@ -180,15 +180,15 @@ impl VMHandler { self.send_node_resources().await; while let Some(brain_msg) = self.receiver.recv().await { match brain_msg.msg { - Some(snp_proto::brain_message::Msg::NewVmReq(new_vm_req)) => { + Some(snp_proto::brain_vm_message::Msg::NewVmReq(new_vm_req)) => { self.handle_new_vm_req(new_vm_req).await; } - Some(snp_proto::brain_message::Msg::UpdateVmReq(update_vm_req)) => { + Some(snp_proto::brain_vm_message::Msg::UpdateVmReq(update_vm_req)) => { if let Err(e) = self.handle_update_vm_req(update_vm_req).await { log::error!("Could not update vm: {e:?}"); } } - Some(snp_proto::brain_message::Msg::DeleteVm(delete_vm_req)) => { + Some(snp_proto::brain_vm_message::Msg::DeleteVm(delete_vm_req)) => { let uuid = delete_vm_req.uuid.clone(); if let Err(e) = self.handle_delete_vm(delete_vm_req) { log::error!("Could not delete vm {uuid}: {e:?}"); @@ -201,7 +201,7 @@ impl VMHandler { } } - fn clear_deleted_contracts(&mut self, contracts: Vec) { + fn clear_deleted_contracts(&mut self, contracts: Vec) { for uuid in self.res.existing_vms.clone() { if contracts.iter().find(|c| c.uuid == uuid).is_none() { info!("VM {uuid} exists locally but not found in brain. Deleting..."); diff --git a/snp.proto b/vm.proto similarity index 79% rename from snp.proto rename to vm.proto index fa9654e..e755772 100644 --- a/snp.proto +++ b/vm.proto @@ -1,5 +1,5 @@ syntax = "proto3"; -package snp_proto; +package vm_proto; message Empty { } @@ -13,7 +13,7 @@ message AccountBalance { uint64 tmp_locked = 2; } -message Contract { +message VmContract { string uuid = 1; string hostname = 2; string admin_pubkey = 3; @@ -53,7 +53,7 @@ message MeasurementIP { } // This should also include a block hash or similar, for auth -message RegisterNodeReq { +message RegisterVmNodeReq { string node_pubkey = 1; string owner_pubkey = 2; string main_ip = 3; @@ -64,7 +64,7 @@ message RegisterNodeReq { uint64 price = 7; } -message NodeResources { +message VmNodeResources { string node_pubkey = 1; uint32 avail_ports = 2; uint32 avail_ipv4 = 3; @@ -123,7 +123,7 @@ message DeleteVmReq { string admin_pubkey = 2; } -message BrainMessage { +message BrainVmMessage { oneof Msg { NewVmReq new_vm_req = 1; UpdateVmReq update_vm_req = 2; @@ -138,28 +138,28 @@ message DaemonStreamAuth { string signature = 4; } -message DaemonMessage { +message VmDaemonMessage { oneof Msg { DaemonStreamAuth auth = 1; NewVmResp new_vm_resp = 2; UpdateVmResp update_vm_resp = 3; - NodeResources node_resources = 4; + VmNodeResources vm_node_resources = 4; } } -service BrainDaemon { - rpc RegisterNode (RegisterNodeReq) returns (stream Contract); - rpc BrainMessages (DaemonStreamAuth) returns (stream BrainMessage); - rpc DaemonMessages (stream DaemonMessage) returns (Empty); +service BrainVmDaemon { + rpc RegisterVmNode (RegisterVmNodeReq) returns (stream VmContract); + rpc BrainMessages (DaemonStreamAuth) returns (stream BrainVmMessage); + rpc DaemonMessages (stream VmDaemonMessage) returns (Empty); } -message ListContractsReq { +message ListVmContractsReq { string admin_pubkey = 1; string node_pubkey = 2; string uuid = 3; } -message NodeFilters { +message VmNodeFilters { uint32 free_ports = 1; bool offers_ipv4 = 2; bool offers_ipv6 = 3; @@ -173,7 +173,7 @@ message NodeFilters { string node_pubkey = 11; } -message NodeListResp { +message VmNodeListResp { string node_pubkey = 1; string country = 2; string region = 3; @@ -191,14 +191,28 @@ message ExtendVmReq { uint64 locked_nano = 3; } +message AirdropReq { + string pubkey = 1; + uint64 tokens = 2; +} + +message Account { + string pubkey = 1; + uint64 balance = 2; + uint64 tmp_locked = 3; +} + service BrainCli { - rpc GetAirdrop (Pubkey) returns (Empty); rpc GetBalance (Pubkey) returns (AccountBalance); rpc NewVm (NewVmReq) returns (NewVmResp); - rpc ListContracts (ListContractsReq) returns (stream Contract); - rpc ListNodes (NodeFilters) returns (stream NodeListResp); - rpc GetOneNode (NodeFilters) returns (NodeListResp); + rpc ListVmContracts (ListVmContractsReq) returns (stream VmContract); + rpc ListVmNodes (VmNodeFilters) returns (stream VmNodeListResp); + rpc GetOneVmNode (VmNodeFilters) returns (VmNodeListResp); rpc DeleteVm (DeleteVmReq) returns (Empty); rpc UpdateVm (UpdateVmReq) returns (UpdateVmResp); rpc ExtendVm (ExtendVmReq) returns (Empty); + // admin commands + rpc Airdrop (AirdropReq) returns (Empty); + rpc ListAllVmContracts (Empty) returns (stream VmContract); + rpc ListAccounts (Empty) returns (stream Account); }