From ff72ff6c1e45801381c955abd4335a2db429a666 Mon Sep 17 00:00:00 2001
From: Noor <noormohammedb@protonmail.com>
Date: Thu, 12 Jun 2025 19:29:14 +0530
Subject: [PATCH] fixed pub_sub_node field error on re registration

---
 src/constants.rs |  2 ++
 src/db/mod.rs    | 18 +++---------------
 src/db/vm.rs     |  6 ++++--
 src/grpc/vm.rs   |  8 ++++----
 4 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/src/constants.rs b/src/constants.rs
index baf9980..2a29690 100644
--- a/src/constants.rs
+++ b/src/constants.rs
@@ -52,3 +52,5 @@ pub const MIN_ESCROW: u64 = 5000 * TOKEN_DECIMAL;
 
 pub const APP_DAEMON_TIMEOUT: u64 = 20;
 pub const VM_DAEMON_TIMEOUT: u64 = 10;
+
+pub const DEFAULT_ENDPOINT: &str = "127.0.0.1:31337";
diff --git a/src/db/mod.rs b/src/db/mod.rs
index cdecd5a..ef55c45 100644
--- a/src/db/mod.rs
+++ b/src/db/mod.rs
@@ -3,8 +3,8 @@ pub mod general;
 pub mod vm;
 
 use crate::constants::{
-    APP_NODE, DB_SCHEMA_FILES, DELETED_APP, DELETED_VM, MIN_ESCROW, NEW_APP_REQ, NEW_VM_REQ,
-    UPDATE_VM_REQ,
+    APP_NODE, DB_SCHEMA_FILES, DEFAULT_ENDPOINT, DELETED_APP, DELETED_VM, MIN_ESCROW, NEW_APP_REQ,
+    NEW_VM_REQ, UPDATE_VM_REQ,
 };
 use crate::old_brain;
 use prelude::*;
@@ -196,24 +196,12 @@ pub async fn live_appnode_msgs<
     Ok(())
 }
 
-pub async fn set_pubsub_node(db: &Surreal<Client>, id: RecordId) {
-    let pub_endpoint = std::env::var("BRAIN_PUBLIC_ENDPOINT").unwrap_or_default();
-    match db
-        .query(format!("UPDATE $id SET pub_sub_node = '{pub_endpoint}'",))
-        .bind(("id", id))
-        .await
-    {
-        Ok(res) => log::info!("Updated pub_sub_node {:?}", res),
-        Err(e) => log::error!("Could not update pub_sub_node {:?}", e),
-    }
-}
-
 pub async fn check_pubsub_node(
     db: &Surreal<Client>,
     id: RecordId,
 ) -> Result<Option<tonic::Status>, Error> {
     let pub_endpoint =
-        std::env::var("BRAIN_PUBLIC_ENDPOINT").unwrap_or("127.0.0.1:31337".to_string());
+        std::env::var("BRAIN_PUBLIC_ENDPOINT").unwrap_or(DEFAULT_ENDPOINT.to_string());
 
     let mut query_resp = db.query(format!("select pub_sub_node from {id}")).await?;
     let node_endpoint = query_resp
diff --git a/src/db/vm.rs b/src/db/vm.rs
index 5a45677..4f1583a 100644
--- a/src/db/vm.rs
+++ b/src/db/vm.rs
@@ -3,8 +3,8 @@ use std::time::Duration;
 
 use super::Error;
 use crate::constants::{
-    ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, UPDATE_VM_REQ, VM_DAEMON_TIMEOUT, VM_NODE,
-    VM_UPDATE_EVENT,
+    ACCOUNT, ACTIVE_VM, DEFAULT_ENDPOINT, DELETED_VM, NEW_VM_REQ, UPDATE_VM_REQ, VM_DAEMON_TIMEOUT,
+    VM_NODE, VM_UPDATE_EVENT,
 };
 use crate::db::{Account, ErrorFromTable, Report};
 use crate::old_brain;
@@ -19,6 +19,7 @@ use tokio_stream::StreamExt as _;
 pub struct VmNode {
     pub id: RecordId,
     pub operator: RecordId,
+    pub pub_sub_node: String,
     pub country: String,
     pub region: String,
     pub city: String,
@@ -997,6 +998,7 @@ impl From<&old_brain::BrainData> for Vec<VmNode> {
             nodes.push(VmNode {
                 id: RecordId::from((VM_NODE, old_node.public_key.clone())),
                 operator: RecordId::from((ACCOUNT, old_node.operator_wallet.clone())),
+                pub_sub_node: DEFAULT_ENDPOINT.to_string(),
                 country: old_node.country.clone(),
                 region: old_node.region.clone(),
                 city: old_node.city.clone(),
diff --git a/src/grpc/vm.rs b/src/grpc/vm.rs
index 7d62cac..fcc0631 100644
--- a/src/grpc/vm.rs
+++ b/src/grpc/vm.rs
@@ -1,5 +1,5 @@
 #![allow(dead_code)]
-use crate::constants::{ACCOUNT, NEW_VM_REQ, UPDATE_VM_REQ, VM_NODE};
+use crate::constants::{ACCOUNT, DEFAULT_ENDPOINT, NEW_VM_REQ, UPDATE_VM_REQ, VM_NODE};
 use crate::db::prelude as db;
 use crate::grpc::{check_sig_from_parts, check_sig_from_req};
 use detee_shared::common_proto::Empty;
@@ -38,10 +38,11 @@ impl BrainVmDaemon for VmDaemonServer {
     ) -> Result<Response<Self::RegisterVmNodeStream>, Status> {
         let req = check_sig_from_req(req)?;
         info!("Starting registration process for {:?}", req);
-        let id = surrealdb::RecordId::from((VM_NODE, req.node_pubkey.clone()));
         db::VmNode {
-            id: id.clone(),
+            id: surrealdb::RecordId::from((VM_NODE, req.node_pubkey.clone())),
             operator: surrealdb::RecordId::from((ACCOUNT, req.operator_wallet)),
+            pub_sub_node: std::env::var("BRAIN_PUBLIC_ENDPOINT")
+                .unwrap_or(DEFAULT_ENDPOINT.to_string()),
             country: req.country,
             region: req.region,
             city: req.city,
@@ -59,7 +60,6 @@ impl BrainVmDaemon for VmDaemonServer {
         }
         .register(&self.db)
         .await?;
-        db::set_pubsub_node(&self.db, id).await;
 
         info!("Sending deleted contracts to {}", req.node_pubkey);
         let deleted_vms = db::DeletedVm::list_by_node(&self.db, &req.node_pubkey).await?;