testnet/detee-cli
3
0
Fork 0
Code Issues Pull Requests 1 Packages Projects Releases Wiki Activity

Implement macro for grpc redirect

Browse Source 
redirect handling in gRPC client and add macro for follow redirects
random brain url selection for staging environment.
cleanup duplicate sign_request function
This commit is contained in:
Noor 2025-06-13 17:08:52 +05:30
parent 41d9bd104f
commit a25c53d709
Signed by: noormohammedb
GPG Key ID: D83EFB8B3B967146
3 changed files with 64 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/config.rs
View File

@ -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()),
}

50
src/snp/grpc.rs
View File

@ -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;
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(e) => return Err(e.into()),
};
}
Err(Error::MaxRedirectsExceeded(MAX_REDIRECTS.to_string()))
}
pub async fn list_contracts(req: ListVmContractsReq) -> Result<Vec<VmContract>, Error> {

46
src/utils.rs
View File

@ -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()))
}
};
}