PERP-4586 | Refactor CosmosInterceptor and parse_cosmos_grpc

Author: anakinzhed

packages/cosmos/src/client.rs

@@ -5,7 +5,6 @@ pub(crate) mod query;
 
 use std::{
     collections::HashMap,
-    str::FromStr,
     sync::{Arc, Weak},
 };
 
@@ -32,7 +31,7 @@ use cosmos_sdk_proto::{
 use parking_lot::{Mutex, RwLock};
 use tokio::{sync::mpsc::Receiver, task::JoinSet, time::Instant};
 use tonic::{
-    metadata::{MetadataKey, MetadataValue},
+    metadata::{Ascii, MetadataKey, MetadataValue},
     service::Interceptor,
     Status,
 };
@@ -743,19 +742,13 @@ impl Cosmos {
 }
 
 #[derive(Clone)]
-pub struct CosmosInterceptor(pub HashMap<String, String>);
+pub struct CosmosInterceptor(Arc<[(MetadataKey<Ascii>, MetadataValue<Ascii>)]>);
 
 impl Interceptor for CosmosInterceptor {
     fn call(&mut self, mut request: tonic::Request<()>) -> Result<tonic::Request<()>, Status> {
-        let req = request.metadata_mut();
-        for (key, value) in &self.0 {
-            let key = MetadataKey::from_bytes(key.as_bytes())
-                .map_err(|_| Status::internal(format!("Invalid header key: '{}'", key)))?;
-
-            let value = MetadataValue::from_str(value)
-                .map_err(|_| Status::internal(format!("Invalid header value for '{}'", key)))?;
-
-            req.insert(key, value);
+        let meta = request.metadata_mut();
+        for (key, value) in self.0.iter() {
+            meta.insert(key, value.clone());
         }
         Ok(request)
     }

packages/cosmos/src/client/node.rs

@@ -1,6 +1,6 @@
 use std::{
-    collections::HashMap,
     ops::Deref,
+    str::FromStr,
     sync::Arc,
     time::{Duration, Instant},
 };
@@ -9,7 +9,9 @@ use chrono::{DateTime, Utc};
 use parking_lot::RwLock;
 use tonic::{
     codegen::InterceptedService,
+    metadata::{Ascii, MetadataKey, MetadataValue},
     transport::{Channel, ClientTlsConfig, Endpoint, Uri},
+    Status,
 };
 
 use crate::{
@@ -91,23 +93,47 @@ impl LastError {
     }
 }
 
-fn parse_cosmos_grpc(value: &str) -> (String, HashMap<String, String>) {
-    let parts: Vec<&str> = value.splitn(2, '#').collect();
-
-    let endpoint = parts[0].trim().to_string();
-    let headers = if parts.len() == 2 {
-        parts[1]
-            .split(';')
-            .filter_map(|pair| {
-                let mut kv = pair.splitn(2, '=');
-                Some((kv.next()?.to_string(), kv.next()?.to_string()))
-            })
-            .collect()
-    } else {
-        HashMap::new()
+type ParseCosmosGrpcResult =
+    Result<(String, Arc<[(MetadataKey<Ascii>, MetadataValue<Ascii>)]>), Status>;
+
+pub fn parse_cosmos_grpc(value: &str) -> ParseCosmosGrpcResult {
+    let (endpoint, raw_headers) = match value.split_once('#') {
+        Some((endpoint, headers)) => (endpoint.trim().to_string(), Some(headers)),
+        None => (value.trim().to_string(), None),
+    };
+
+    let headers = {
+        let mut parsed = Vec::new();
+        if let Some(hdrs) = raw_headers {
+            for pair in hdrs.split(';').filter(|s| !s.trim().is_empty()) {
+                let (key, val) = pair.split_once('=').ok_or_else(|| {
+                    Status::invalid_argument(format!("Malformed header: '{}'", pair))
+                })?;
+
+                let key = MetadataKey::from_bytes(key.trim().as_bytes()).map_err(|_| {
+                    Status::invalid_argument(format!("Invalid header key '{}'", key))
+                })?;
+
+                let val_str = val.trim();
+
+                if val_str.is_empty() {
+                    return Err(Status::invalid_argument(format!(
+                        "Header '{}' has empty value",
+                        key
+                    )));
+                }
+
+                let val = MetadataValue::from_str(val_str).map_err(|_| {
+                    Status::invalid_argument(format!("Invalid header value for '{}'", key))
+                })?;
+
+                parsed.push((key, val));
+            }
+        }
+        Arc::from(parsed.into_boxed_slice())
     };
 
-    (endpoint, headers)
+    Ok((endpoint, headers))
 }
 
 impl CosmosBuilder {
@@ -116,7 +142,11 @@ impl CosmosBuilder {
         grpc_url: &Arc<String>,
         is_fallback: bool,
     ) -> Result<Node, BuilderError> {
-        let (url, headers) = parse_cosmos_grpc(grpc_url.as_str());
+        let (url, headers) =
+            parse_cosmos_grpc(grpc_url.as_str()).map_err(|e| BuilderError::InvalidGrpcHeaders {
+                grpc_url: grpc_url.clone(),
+                source: e,
+            })?;
         let grpc_url = Arc::<String>::new(url);
 
         let grpc_endpoint =

packages/cosmos/src/error.rs

@@ -102,6 +102,11 @@ pub enum BuilderError {
     SanityQueryFailed { source: QueryError },
     #[error("Could not find Sei gas config for chain ID {chain_id} at {url}")]
     SeiGasConfigNotFound { chain_id: String, url: String },
+    #[error("Invalid gRPC headers in {grpc_url}: {source}")]
+    InvalidGrpcHeaders {
+        grpc_url: Arc<String>,
+        source: tonic::Status,
+    },
 }
 
 /// Parse errors while interacting with chain data.