From bff04c7b7ea104099f72545f1093d2ba960e8b3f Mon Sep 17 00:00:00 2001
From: Ramil_Algayev <pro.remred@gmail.com>
Date: Sat, 4 Jan 2025 17:20:16 +0400
Subject: [PATCH] Temporary code to test updates, this is not final at all.

---
 brain.proto |  6 +++---
 src/data.rs | 18 +++++++++---------
 src/grpc.rs |  2 +-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/brain.proto b/brain.proto
index b8ef0ee..7253420 100644
--- a/brain.proto
+++ b/brain.proto
@@ -118,7 +118,7 @@ service BrainDaemonService {
   rpc ListVMContracts (ListVMContractsReq) returns (stream VMContract);
   rpc GetUpdateVMReq (NodePubkey) returns (stream UpdateVMReq);
   rpc SendUpdateVMResp (stream UpdateVMResp) returns (Empty);
-  // rpc SendMeasurementArgs (NodePubkey) returns (MeasurementArgs);
+  //rpc GetMeasurementArgs (ListVMContractsReq) returns (stream MeasurementArgs);
 }
 
 message NodeFilters {
@@ -151,5 +151,5 @@ service BrainCliService {
   rpc GetOneNode (NodeFilters) returns (NodeListResp);
   rpc DeleteVM (DeleteVMReq) returns (Empty);
   rpc UpdateVM (UpdateVMReq) returns (UpdateVMResp);
-  // rpc GetMeasurementArgs (ListVMContractsReq) returns (MeasurementArgs);
-}
+  //rpc GetMeasurementArgs (ListVMContractsReq) returns (MeasurementArgs);
+}
\ No newline at end of file
diff --git a/src/data.rs b/src/data.rs
index 4b0c522..0f09021 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -1,5 +1,5 @@
 #![allow(dead_code)]
-use crate::grpc::brain as grpc;
+use crate::grpc::brain::{self as grpc, MeasurementArgs};
 use chrono::Utc;
 use dashmap::DashMap;
 use log::{debug, info, warn};
@@ -87,6 +87,7 @@ pub struct BrainData {
     daemon_deletevm_tx: DashMap<String, Sender<grpc::DeleteVmReq>>,
     daemon_newvm_tx: DashMap<String, Sender<grpc::NewVmReq>>,
     daemon_updatevm_tx: DashMap<String, Sender<grpc::UpdateVmReq>>,
+    daemon_get_measurement_tx: DashMap<String, Sender<MeasurementArgs>>,
 }
 
 #[derive(Debug)]
@@ -106,6 +107,7 @@ impl BrainData {
             daemon_deletevm_tx: DashMap::new(),
             daemon_newvm_tx: DashMap::new(),
             daemon_updatevm_tx: DashMap::new(),
+            daemon_get_measurement_tx: DashMap::new(),
         }
     }
 
@@ -208,7 +210,7 @@ impl BrainData {
         }
         let mut public_ipv4 = String::new();
         let mut public_ipv6 = String::new();
-        for ip in new_vm_resp.ips {
+        for ip in new_vm_resp.args.clone().unwrap_or(grpc::MeasurementArgs::default()).ips {
             if let Ok(ipv4_addr) = std::net::Ipv4Addr::from_str(&ip.address) {
                 if !ipv4_addr.is_private() && !ipv4_addr.is_link_local() {
                     public_ipv4 = ipv4_addr.to_string();
@@ -221,7 +223,7 @@ impl BrainData {
         }
         let contract = Contract {
             uuid: new_vm_resp.uuid,
-            exposed_ports: new_vm_resp.exposed_ports,
+            exposed_ports: new_vm_resp.args.unwrap_or(grpc::MeasurementArgs::default()).exposed_ports,
             public_ipv4,
             public_ipv6,
             created_at: Utc::now().to_rfc3339(),
@@ -293,9 +295,7 @@ impl BrainData {
                 self.submit_newvm_resp(grpc::NewVmResp {
                     error: "Daemon is offline.".to_string(),
                     uuid: req.uuid,
-                    exposed_ports: Vec::new(),
-                    ovmf_hash: "".to_string(),
-                    ips: Vec::new(),
+                    args: None,
                 })
                 .await;
             }
@@ -319,7 +319,7 @@ impl BrainData {
                 let _ = tx.send(grpc::UpdateVmResp {
                     uuid,
                     error: "Contract does not exist.".to_string(),
-                    ovmf_hash: "".to_string(),
+                    args: None,
                 });
                 return;
             }
@@ -344,7 +344,7 @@ impl BrainData {
                     self.submit_updatevm_resp(grpc::UpdateVmResp {
                         uuid,
                         error: "Daemon is offline.".to_string(),
-                        ovmf_hash: "".to_string(),
+                        args: None,
                     })
                     .await;
                 }
@@ -354,7 +354,7 @@ impl BrainData {
             self.submit_updatevm_resp(grpc::UpdateVmResp {
                 uuid,
                 error: "Daemon is offline.".to_string(),
-                ovmf_hash: "".to_string(),
+                args: None,
             })
             .await;
         }
diff --git a/src/grpc.rs b/src/grpc.rs
index aa82a00..adb280e 100644
--- a/src/grpc.rs
+++ b/src/grpc.rs
@@ -256,7 +256,7 @@ impl BrainDaemonService for BrainDaemonMock {
                     data.submit_updatevm_resp(UpdateVmResp {
                         error: "Daemon not connected.".to_string(),
                         uuid: updatevmreq.uuid,
-                        ovmf_hash: "".to_string(),
+                        args: None,
                     })
                     .await;
                     break;