diff --git a/src/config.rs b/src/config.rs index 283ece6..0495f5f 100644 --- a/src/config.rs +++ b/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::() { 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()), } diff --git a/src/snp/grpc.rs b/src/snp/grpc.rs index 27b9ee4..9be8310 100644 --- a/src/snp/grpc.rs +++ b/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(req: T) -> Result, 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, Error> { debug!("Getting nodes from brain..."); let mut client = client().await?; @@ -138,30 +127,13 @@ pub async fn get_one_node(req: VmNodeFilters) -> Result { } pub async fn create_vm(req: NewVmReq) -> Result { - 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, Error> { diff --git a/src/utils.rs b/src/utils.rs index 022da20..5d58dcc 100644 --- a/src/utils.rs +++ b/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())) + } + }; +}