cli listing nodes and submiting newvmrequest

This commit is contained in:
ghe0 2024-12-22 20:00:12 +02:00
parent 47eb811a6e
commit 6ee235a2ec
Signed by: ghe0
GPG Key ID: 451028EE56A0FBB4
6 changed files with 1473 additions and 4 deletions

1231
cli-mock/Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

@ -4,3 +4,16 @@ version = "0.1.0"
edition = "2021" edition = "2021"
[dependencies] [dependencies]
anyhow = "1.0.94"
env_logger = "0.11.6"
lazy_static = "1.5.0"
log = "0.4.22"
prost = "0.13.4"
prost-types = "0.13.4"
rand = "0.8.5"
tokio = { version = "1.42.0", features = ["macros", "rt-multi-thread"] }
tokio-stream = "0.1.17"
tonic = "0.12"
[build-dependencies]
tonic-build = "0.12"

109
cli-mock/brain.proto Normal file

@ -0,0 +1,109 @@
syntax = "proto3";
package brain;
message Empty {
}
message NodePubkey {
string node_pubkey = 1;
}
message RegisterNodeRequest {
string node_pubkey = 1;
string owner_pubkey = 2;
string ip = 3;
string country = 4;
string city = 5;
uint32 avail_ports = 6;
uint32 avail_ipv4 = 7;
uint32 avail_ipv6 = 8;
uint32 avail_vcpus = 9;
uint32 avail_memory_mb = 10;
uint32 avail_storage_gb = 11;
uint32 max_ports_per_vm = 12;
}
message NewVMRequest {
string uuid = 1; // UUID is empty when CLI sends request; brain sets UUID
string hostname = 2;
string admin_pubkey = 3;
string node_pubkey = 4;
repeated uint32 extra_ports = 5;
bool public_ipv4 = 6;
bool public_ipv6 = 7;
uint32 disk_size_gb = 8;
uint32 vcpus = 9;
uint32 memory_mb = 10;
string kernel_url = 11;
string kernel_sha = 12;
string dtrfs_url = 13;
string dtrfs_sha = 14;
}
message VMContract {
string uuid = 1;
string hostname = 2;
string admin_pubkey = 3;
string node_pubkey = 4;
repeated uint32 exposed_ports = 5;
string public_ipv4 = 6;
string public_ipv6 = 7;
uint32 disk_size_gb = 8;
uint32 vcpus = 9;
uint32 memory_mb = 10;
string kernel_sha = 11;
string dtrfs_sha = 12;
string created_at = 13;
}
message ListVMContractsReq {
string admin_pubkey = 1;
string node_pubkey = 2;
}
message NewVMConfirmation {
string uuid = 1;
repeated uint32 exposed_ports = 2;
string public_ipv4 = 3;
string public_ipv6 = 4;
string error = 5;
}
message DeletedVMUpdate {
string uuid = 1;
}
service BrainDaemonService {
rpc RegisterNode (RegisterNodeRequest) returns (Empty);
rpc GetNewVMReqs (NodePubkey) returns (stream NewVMRequest);
rpc SendVMConfirmations (stream NewVMConfirmation) returns (Empty);
rpc DeletedVMUpdates (NodePubkey) returns (stream DeletedVMUpdate);
rpc ListVMContracts (ListVMContractsReq) returns (stream VMContract);
}
message NodeFilters {
uint32 free_ports = 1;
bool offers_ipv4 = 2;
bool offers_ipv6 = 3;
uint32 vcpus = 4;
uint32 memory_mb = 5;
uint32 storage_gb = 6;
string country = 7;
}
message NodeListResp {
string node_pubkey = 1;
string country = 2;
string city = 3;
string ip = 4; // required for latency test
uint32 server_rating = 5;
uint32 provider_rating = 6;
}
service BrainCliService {
rpc CreateVMContract (NewVMRequest) returns (NewVMConfirmation);
rpc ListVMContracts (ListVMContractsReq) returns (stream VMContract);
rpc ListNodes (NodeFilters) returns (stream NodeListResp);
rpc DeleteVM (DeletedVMUpdate) returns (Empty);
}

6
cli-mock/build.rs Normal file

@ -0,0 +1,6 @@
fn main() {
tonic_build::configure()
.build_server(true)
.compile_protos(&["brain.proto"], &["proto"])
.unwrap_or_else(|e| panic!("Failed to compile protos {:?}", e));
}

@ -1,3 +1,114 @@
fn main() { #![allow(dead_code)]
println!("Hello, world!"); pub mod brain {
tonic::include_proto!("brain");
}
use anyhow::Result;
use brain::{
brain_cli_service_client::BrainCliServiceClient, NewVmRequest, NodeFilters, NodeListResp,
};
use lazy_static::lazy_static;
use log::{debug, info, warn};
use rand::{distributions::Alphanumeric, thread_rng, Rng};
use tokio_stream::StreamExt;
use tonic::transport::Channel;
lazy_static! {
static ref SECURE_PUBLIC_KEY: String = generate_random_string();
}
fn generate_random_string() -> String {
let rng = thread_rng();
rng.sample_iter(&Alphanumeric)
.take(16) // Adjust the length as needed
.map(char::from)
.collect()
}
async fn get_node_list(mut client: BrainCliServiceClient<Channel>) -> Result<Vec<NodeListResp>> {
info!("Getting nodes from brain...");
let mut nodes = Vec::new();
let mut grpc_stream = client
.list_nodes(NodeFilters {
free_ports: 0,
offers_ipv4: true,
offers_ipv6: true,
vcpus: 0,
memory_mb: 0,
storage_gb: 0,
country: String::new(),
})
.await?
.into_inner();
while let Some(stream_update) = grpc_stream.next().await {
match stream_update {
Ok(node) => {
debug!("Received node from brain: {node:?}");
nodes.push(node);
}
Err(e) => {
warn!("Received error instead of node list: {e:?}");
}
}
}
info!("Brain terminated list_nodes stream.");
Ok(nodes)
}
async fn submit_vm_request(
mut client: BrainCliServiceClient<Channel>,
node_pubkey: &str,
) -> Result<()> {
let req = NewVmRequest {
uuid: String::new(),
admin_pubkey: SECURE_PUBLIC_KEY.clone(),
node_pubkey: node_pubkey.to_string(),
hostname: "cool-server".to_string(),
vcpus: 4,
memory_mb: 4000,
disk_size_gb: 20,
public_ipv4: true,
public_ipv6: true,
extra_ports: Vec::new(),
kernel_url: "thisIsMyURL".to_string(),
kernel_sha: "thisIsMySha".to_string(),
dtrfs_url: "thisIsMyURL".to_string(),
dtrfs_sha: "thisIsMySha".to_string(),
};
info!("Creating VM {req:?}");
let result = client.create_vm_contract(req).await;
match result {
Ok(confirmation) => {
let confirmation = confirmation.into_inner();
if confirmation.error.is_empty() {
info!("Got VM confirmation: {confirmation:?}");
} else {
warn!("Got VM confirmation error: {}", confirmation.error);
};
}
Err(e) => log::error!("Could not create vm: {e:?}"),
};
Ok(())
}
#[tokio::main]
async fn main() -> Result<()> {
env_logger::builder()
.filter_level(log::LevelFilter::Debug)
.init();
info!("Hello! My name is {}", SECURE_PUBLIC_KEY.clone());
let client = BrainCliServiceClient::connect("http://[::1]:31337").await?;
let nodes = get_node_list(client.clone()).await?;
for node in nodes {
if let Err(e) = submit_vm_request(client.clone(), &node.node_pubkey).await {
log::error!(
"Received error when creating VM on node {}: {e:?}",
&node.node_pubkey
);
}
}
Ok(())
} }

@ -119,7 +119,6 @@ async fn handle_vm_requests(mut req: Receiver<NewVmRequest>, resp: Sender<NewVmC
info!("Started to handle vm requests. 1 out of 5 requests will return error."); info!("Started to handle vm requests. 1 out of 5 requests will return error.");
let mut i = 0; let mut i = 0;
while let Some(new_vm) = req.recv().await { while let Some(new_vm) = req.recv().await {
let success = i == 3;
let exposed_ports = match new_vm.public_ipv4 { let exposed_ports = match new_vm.public_ipv4 {
true => Vec::new(), true => Vec::new(),
false => vec![20321, 20415, 25912], false => vec![20321, 20415, 25912],
@ -132,7 +131,7 @@ async fn handle_vm_requests(mut req: Receiver<NewVmRequest>, resp: Sender<NewVmC
true => " 2a02:2f2d:d301:3100:afe8:a85e:54a0:dd28".to_string(), true => " 2a02:2f2d:d301:3100:afe8:a85e:54a0:dd28".to_string(),
false => String::new(), false => String::new(),
}; };
if success { if i != 3 {
let confirmation = NewVmConfirmation { let confirmation = NewVmConfirmation {
uuid: new_vm.uuid, uuid: new_vm.uuid,
exposed_ports, exposed_ports,