cli listing nodes and submiting newvmrequest

cli-mock/Cargo.toml

@@ -4,3 +4,16 @@ version = "0.1.0"
 edition = "2021"
 
 [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"

cli-mock/brain.proto

@@ -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);
+}
+

cli-mock/build.rs

@@ -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));
+}

cli-mock/src/main.rs

@@ -1,3 +1,114 @@
-fn main() {
-    println!("Hello, world!");
+#![allow(dead_code)]
+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(())
 }

daemon-mock/src/main.rs

@@ -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.");
     let mut i = 0;
     while let Some(new_vm) = req.recv().await {
-        let success = i == 3;
         let exposed_ports = match new_vm.public_ipv4 {
             true => Vec::new(),
             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(),
             false => String::new(),
         };
-        if success {
+        if i != 3 {
             let confirmation = NewVmConfirmation {
                 uuid: new_vm.uuid,
                 exposed_ports,