src/config.rs

@@ -309,7 +309,13 @@ impl Config {
 
     pub fn get_brain_info() -> (String, String) {
         match Self::init_config().network.as_str() {
-            "staging" => ("https://10.254.254.8:31337".to_string(), "staging-brain".to_string()),
+            "staging" => {
+                let url1 = "https://149.22.95.1:47855".to_string(); // staging brain 2
+                let url2 = "https://149.36.48.99:48843".to_string(); // staging brain 3
+
+                let url = if rand::random::<bool>() { url1 } else { url2 };
+                (url, "staging-brain".to_string())
+            }
             "localhost" => ("https://localhost:31337".to_string(), "staging-brain".to_string()),
             _ => ("https://173.234.17.2:39477".to_string(), "testnet-brain".to_string()),
         }

src/snp/grpc.rs

@@ -3,8 +3,9 @@ pub mod proto {
     pub use detee_shared::vm_proto::*;
 }
 
+use crate::call_with_follow_redirect;
 use crate::config::Config;
-use crate::constants::MAX_REDIRECTS;
+use crate::utils::{self, sign_request};
 use lazy_static::lazy_static;
 use log::{debug, info, warn};
 use proto::{
@@ -13,9 +14,7 @@ use proto::{
 };
 use tokio_stream::StreamExt;
 use tonic::metadata::errors::InvalidMetadataValue;
-use tonic::metadata::AsciiMetadataValue;
 use tonic::transport::Channel;
-use tonic::Request;
 
 lazy_static! {
     static ref SECURE_PUBLIC_KEY: String = use_default_string();
@@ -44,6 +43,10 @@ pub enum Error {
     CorruptedBrainUrl,
     #[error("Max redirects exceeded: {0}")]
     MaxRedirectsExceeded(String),
+    #[error("Redirect error: {0}")]
+    RedirectError(String),
+    #[error(transparent)]
+    InternalError(#[from] utils::Error),
 }
 
 impl crate::HumanOutput for VmContract {
@@ -97,20 +100,6 @@ async fn client_from_endpoint(
     Ok(BrainVmCliClient::new(Config::connect_brain_channel(reconnect_endpoint).await?))
 }
 
-fn sign_request<T: std::fmt::Debug>(req: T) -> Result<Request<T>, Error> {
-    let pubkey = Config::get_detee_wallet()?;
-    let timestamp = chrono::Utc::now().to_rfc3339();
-    let signature = Config::try_sign_message(&format!("{timestamp}{req:?}"))?;
-    let timestamp: AsciiMetadataValue = timestamp.parse()?;
-    let pubkey: AsciiMetadataValue = pubkey.parse()?;
-    let signature: AsciiMetadataValue = signature.parse()?;
-    let mut req = Request::new(req);
-    req.metadata_mut().insert("timestamp", timestamp);
-    req.metadata_mut().insert("pubkey", pubkey);
-    req.metadata_mut().insert("request-signature", signature);
-    Ok(req)
-}
-
 pub async fn get_node_list(req: VmNodeFilters) -> Result<Vec<VmNodeListResp>, Error> {
     debug!("Getting nodes from brain...");
     let mut client = client().await?;
@@ -138,30 +127,13 @@ pub async fn get_one_node(req: VmNodeFilters) -> Result<VmNodeListResp, Error> {
 }
 
 pub async fn create_vm(req: NewVmReq) -> Result<NewVmResp, Error> {
-    let mut client = client().await?;
     debug!("Sending NewVmReq to brain: {req:?}");
-    for attempt in 0..MAX_REDIRECTS {
-        match client.new_vm(sign_request(req.clone())?).await {
-            Ok(resp) => return Ok(resp.into_inner()),
-            Err(status)
-                if status.code() == tonic::Code::Unavailable
-                    && status.message() == "moved"
-                    && status.metadata().contains_key("location") =>
-            {
-                let redirect_url = status
-                    .metadata()
-                    .get("location")
-                    .and_then(|v| v.to_str().ok().map(|s| format!("https://{s}")))
-                    .unwrap_or_default();
-                // TODO: change this println to log::info!()
-                println!("{attempt}) server moved to a different URL, trying to reconnect... ({redirect_url})");
-                client = client_from_endpoint(redirect_url).await?;
-                continue;
-            }
-            Err(e) => return Err(e.into()),
-        };
+
+    let client = client().await?;
+    match call_with_follow_redirect!(client, req, new_vm).await {
+        Ok(resp) => Ok(resp.into_inner()),
+        Err(e) => Err(e.into()),
     }
-    Err(Error::MaxRedirectsExceeded(MAX_REDIRECTS.to_string()))
 }
 
 pub async fn list_contracts(req: ListVmContractsReq) -> Result<Vec<VmContract>, Error> {

src/utils.rs

@@ -42,3 +42,49 @@ pub fn shorten_string(my_string: &String) -> String {
         format!("{}", first_part)
     }
 }
+
+#[macro_export]
+macro_rules! call_with_follow_redirect {
+    (
+        $client:expr,
+        $req_data:expr,
+        $method:ident
+    ) => {
+        async {
+            let mut client = $client;
+
+            for attempt in 0..crate::constants::MAX_REDIRECTS {
+                debug!("Attempt #{}: Calling method '{}'...", attempt + 1, stringify!($method));
+
+                let req_data_clone = $req_data.clone();
+                let signed_req = crate::utils::sign_request(req_data_clone)?;
+
+                match client.$method(signed_req).await {
+                    Ok(resp) => return Ok(resp),
+
+                    Err(status)
+                        if status.code() == tonic::Code::Unavailable
+                            && status.message() == "moved" =>
+                    {
+                        let redirect_url = status
+                            .metadata()
+                            .get("location")
+                            .and_then(|v| v.to_str().ok())
+                            .ok_or_else(|| {
+                                Error::RedirectError(
+                                    "Server indicated a move but provided no location".into(),
+                                )
+                            })?;
+
+                        info!("Server moved. Redirecting to {}...", redirect_url);
+
+                        client = client_from_endpoint(format!("https://{}", redirect_url)).await?;
+                        continue;
+                    }
+                    Err(e) => return Err(Error::ResponseStatus(e)),
+                }
+            }
+            Err(Error::MaxRedirectsExceeded(crate::constants::MAX_REDIRECTS.to_string()))
+        }
+    };
+}