From d6a8f14124031a3ac8c2bba4661157f29fae5a90 Mon Sep 17 00:00:00 2001
From: Noor <noormohammedb@protonmail.com>
Date: Mon, 16 Jun 2025 10:57:02 +0300
Subject: [PATCH] Improves brain connection reliability

Updates the brain connection logic to randomly select from a list of available URLs for staging and testnet environments.
---
 src/global.rs | 21 ++++++++++++++++++---
 src/grpc.rs   |  6 ++++--
 src/main.rs   |  4 +---
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/src/global.rs b/src/global.rs
index 4a16717..e397a33 100644
--- a/src/global.rs
+++ b/src/global.rs
@@ -4,12 +4,19 @@ use anyhow::Result;
 use ed25519_dalek::SigningKey;
 use lazy_static::lazy_static;
 use log::{info, warn};
+use rand::Rng;
 use sha2::{Digest, Sha256};
-use std::{fs::File, io::Read, io::Write};
+use std::{
+    fs::File,
+    io::{Read, Write},
+    sync::LazyLock,
+};
 
 pub(crate) const DETEE_ROOT_CA: &str = "/etc/detee/root_ca.pem";
-pub(crate) const BRAIN_STAGING: (&str, &str) = ("https://184.107.169.199:49092", "staging-brain");
-pub(crate) const BRAIN_TESTING: (&str, &str) = ("https://164.92.249.180:31337", "testnet-brain");
+pub(crate) const BRAIN_STAGING_URLS: [&str; 3] =
+    ["https://184.107.169.199:49092", "https://149.22.95.1:47855", "https://149.36.48.99:48843"];
+pub(crate) const BRAIN_TESTING_URLS: [&str; 3] =
+    ["https://184.107.169.199:45223", "https://149.22.95.1:44522", "https://149.36.48.99:48638"];
 pub(crate) const VM_BOOT_DIR: &str = "/var/lib/detee/boot/";
 pub(crate) const USED_RESOURCES: &str = "/etc/detee/daemon/used_resources.yaml";
 pub(crate) const VM_CONFIG_DIR: &str = "/etc/detee/daemon/vms/";
@@ -24,6 +31,14 @@ pub(crate) const OVMF_HASH: &str =
 pub(crate) const OVMF_URL: &str =
     "https://drive.google.com/uc?export=download&id=1V-vLkaiLaGmFSjrN84Z6nELQOxKNAoSJ";
 
+pub static BRAIN_STAGING: LazyLock<(&str, &str)> = LazyLock::new(|| {
+    (BRAIN_STAGING_URLS[rand::thread_rng().gen_range(0..BRAIN_STAGING_URLS.len())], "staging-brain")
+});
+
+pub static BRAIN_TESTING: LazyLock<(&str, &str)> = LazyLock::new(|| {
+    (BRAIN_TESTING_URLS[rand::thread_rng().gen_range(0..BRAIN_TESTING_URLS.len())], "testnet-brain")
+});
+
 lazy_static! {
     pub static ref PUBLIC_KEY: String = get_public_key();
     pub static ref IP_INFO: IPInfo = get_ip_info().unwrap();
diff --git a/src/grpc.rs b/src/grpc.rs
index 8ad4f4e..5718cdb 100644
--- a/src/grpc.rs
+++ b/src/grpc.rs
@@ -18,14 +18,16 @@ pub mod snp_proto {
 
 async fn client(network: &str) -> Result<BrainVmDaemonClient<Channel>> {
     let (brain_url, brain_san) = match network {
-        "staging" => BRAIN_STAGING,
-        "testnet" => BRAIN_TESTING,
+        "staging" => *BRAIN_STAGING,
+        "testnet" => *BRAIN_TESTING,
         _ => {
             return Err(anyhow::anyhow!(
                 "The only networks currently supported are staging and testnet."
             ))
         }
     };
+    info!("brain_url: {brain_url}, brain_san: {brain_san}");
+
     let pem = std::fs::read_to_string(DETEE_ROOT_CA)?;
     let ca = Certificate::from_pem(pem);
 
diff --git a/src/main.rs b/src/main.rs
index e2b4d1e..8a243c8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -253,9 +253,7 @@ async fn main() {
 
         info!("Registering with the brain and getting back deleted VMs.");
         match grpc::register_node(&vm_handler.config).await {
-            Ok(deleted_vms) => {
-                vm_handler.clear_deleted_contracts(deleted_vms)
-            }
+            Ok(deleted_vms) => vm_handler.clear_deleted_contracts(deleted_vms),
             Err(e) => log::error!("Could not get contracts from brain: {e:?}"),
         };