From ab4361ad5e470589f7e0e1f0db33e3dd1d77134e Mon Sep 17 00:00:00 2001
From: Ramil_Algayev <pro.remred@gmail.com>
Date: Mon, 6 Jan 2025 23:41:39 +0400
Subject: [PATCH] now update_vm_resp.args.dtrfs_api_endpoint does get the node
 ip from the brain if it didn't have a public ip

---
 src/data.rs | 41 ++++++++++++++++++++++++++++++++---------
 1 file changed, 32 insertions(+), 9 deletions(-)

diff --git a/src/data.rs b/src/data.rs
index 3aa8fc2..25c26bd 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -184,7 +184,7 @@ impl BrainData {
         self.daemon_newvm_tx.insert(node_pubkey.to_string(), tx);
     }
 
-    pub async fn submit_newvm_resp(&self, new_vm_resp: grpc::NewVmResp) {
+    pub async fn submit_newvm_resp(&self, mut new_vm_resp: grpc::NewVmResp) {
         let newvmreq = match self.tmp_newvm_reqs.remove(&new_vm_resp.uuid) {
             Some((_, r)) => r,
             None => {
@@ -192,9 +192,18 @@ impl BrainData {
                     "Received confirmation for ghost NewVMReq {}",
                     new_vm_resp.uuid
                 );
+                new_vm_resp.error = "Received confirmation for ghost NewVMReq.".to_string();
                 return;
             }
         };
+        let args = match new_vm_resp.args.clone() {
+            Some(args) => args,
+            None => {
+            log::error!("NewVmResp does not contain MeasurementArgs for {}", new_vm_resp.uuid);
+            new_vm_resp.error = "Daemon did not return measurement args.".to_string();
+            return;
+            }
+        };
         if let Err(e) = newvmreq.1.send(new_vm_resp.clone()) {
             log::error!(
                 "CLI RX for {} dropped before receiving confirmation {:?}. Error is: {:?}",
@@ -208,13 +217,6 @@ impl BrainData {
         }
         let mut public_ipv4 = String::new();
         let mut public_ipv6 = String::new();
-        let args = match new_vm_resp.args {
-            Some(args) => args,
-            None => {
-            log::error!("NewVmResp does not contain MeasurementArgs for {}", new_vm_resp.uuid);
-            return;
-            }
-        };
 
         for ip in args.ips {
             if let Ok(ipv4_addr) = std::net::Ipv4Addr::from_str(&ip.address) {
@@ -247,7 +249,7 @@ impl BrainData {
         self.contracts.write().unwrap().push(contract);
     }
 
-    pub async fn submit_updatevm_resp(&self, resp: grpc::UpdateVmResp) {
+    pub async fn submit_updatevm_resp(&self, mut resp: grpc::UpdateVmResp) {
         let updatevmreq = match self.tmp_updatevm_reqs.remove(&resp.uuid) {
             Some((_, r)) => r,
             None => {
@@ -255,9 +257,15 @@ impl BrainData {
                     "Received confirmation for ghost UpdateVMRequest {}",
                     resp.uuid
                 );
+                resp.error = "Received confirmation for ghost UpdateVMRequest.".to_string();
                 return;
             }
         };
+        if let None = resp.args {
+            log::error!("NewVmResp does not contain MeasurementArgs for {}", resp.uuid);
+            resp.error = "Daemon did not return measurement args.".to_string();
+            return;
+        };
         if let Err(e) = updatevmreq.1.send(resp.clone()) {
             log::error!("CLI RX dropped before receiving UpdateVMResp {resp:?}. Error is: {e:?}");
         }
@@ -266,6 +274,21 @@ impl BrainData {
         }
         let mut contracts = self.contracts.write().unwrap();
         if let Some(contract) = contracts.iter_mut().find(|c| c.uuid == resp.uuid) {
+            let args =  resp.args.as_mut().unwrap();
+            if args.dtrfs_api_endpoint.starts_with(':') {
+                if let Some(node) = self.find_nodes_by_pubkey(&contract.node_pubkey) {
+                    args.dtrfs_api_endpoint = format!("{}{}", node.ip, args.dtrfs_api_endpoint);
+                } else {
+                    // This should never happen.
+                    log::error!(
+                        "Node {} not found for contract {}. Cannot update contract.",
+                        contract.node_pubkey,
+                        contract.uuid
+                    );
+                    return;
+                }
+            }
+
             contract.disk_size_gb = updatevmreq.0.disk_size_gb;
             contract.vcpus = updatevmreq.0.vcpus;
             contract.memory_mb = updatevmreq.0.memory_mb;