chore: Handle forwarding with consumer instead of regex

Author: gonzalezzfelipe

proxy/src/auth.rs

@@ -33,7 +33,7 @@ pub fn start(state: Arc<State>) {
                         .iter()
                         .map(|crd| {
                             let consumer = Consumer::from(crd);
-                            (consumer.hash_key.clone(), consumer)
+                            (consumer.key.clone(), consumer)
                         })
                         .collect();
                     *state.consumers.write().await = consumers;
@@ -52,7 +52,7 @@ pub fn start(state: Arc<State>) {
                             .consumers
                             .write()
                             .await
-                            .insert(consumer.hash_key.clone(), consumer);
+                            .insert(consumer.key.clone(), consumer);
                     }
                     None => {
                         // New ports are created without status. When the status is added, a new
@@ -67,7 +67,7 @@ pub fn start(state: Arc<State>) {
                         crd.name_any()
                     );
                     let consumer = Consumer::from(&crd);
-                    state.consumers.write().await.remove(&consumer.hash_key);
+                    state.consumers.write().await.remove(&consumer.key);
                     state.limiter.write().await.remove(&consumer.key);
                 }
                 // Empty response from stream. Should never happen.

proxy/src/limiter.rs

@@ -50,7 +50,7 @@ async fn add_limiter(state: &State, consumer: &Consumer, tier: &Tier) {
 pub async fn limiter(state: Arc<State>, consumer: &Consumer) -> Result<(), LimiterError> {
     if !has_limiter(&state, consumer).await {
         let consumers = state.consumers.read().await.clone();
-        let refreshed_consumer = match consumers.get(&consumer.hash_key) {
+        let refreshed_consumer = match consumers.get(&consumer.key) {
             Some(consumer) => consumer,
             None => return Err(LimiterError::PortDeleted),
         };

proxy/src/main.rs

@@ -52,7 +52,7 @@ impl State {
     pub fn try_new() -> Result<Self, Box<dyn Error>> {
         let config = Config::new();
         let metrics = Metrics::try_new(Registry::default())?;
-        let host_regex = Regex::new(r"(dmtr_[\w\d-]+)?\.?([\w-]+)-v([\d]).+")?;
+        let host_regex = Regex::new(r"(dmtr_[\w\d-]+)?\.?.+")?;
         let consumers = Default::default();
         let tiers = Default::default();
         let limiter = Default::default();
@@ -67,10 +67,8 @@ impl State {
         })
     }
 
-    pub async fn get_consumer(&self, network: &str, version: &str, key: &str) -> Option<Consumer> {
-        let consumers = self.consumers.read().await.clone();
-        let hash_key = format!("{}.{}.{}", network, version, key);
-        consumers.get(&hash_key).cloned()
+    pub async fn get_consumer(&self, key: &str) -> Option<Consumer> {
+        self.consumers.read().await.clone().get(key).cloned()
     }
 }
 
@@ -80,7 +78,16 @@ pub struct Consumer {
     port_name: String,
     tier: String,
     key: String,
-    hash_key: String,
+    network: String,
+    version: String,
+}
+impl Consumer {
+    pub fn instance(&self, state: State) -> String {
+        format!(
+            "ogmios-{}-{}.{}:{}",
+            self.network, self.version, state.config.ogmios_dns, state.config.ogmios_port
+        )
+    }
 }
 impl Display for Consumer {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -90,19 +97,19 @@ impl Display for Consumer {
 impl From<&OgmiosPort> for Consumer {
     fn from(value: &OgmiosPort) -> Self {
         let network = value.spec.network.to_string();
-        let version = value.spec.version;
+        let version = value.spec.version.to_string();
         let tier = value.spec.throughput_tier.to_string();
         let key = value.status.as_ref().unwrap().auth_token.clone();
         let namespace = value.metadata.namespace.as_ref().unwrap().clone();
         let port_name = value.name_any();
 
-        let hash_key = format!("{}.{}.{}", network, version, key);
         Self {
             namespace,
             port_name,
             tier,
             key,
-            hash_key,
+            network,
+            version,
         }
     }
 }

proxy/src/metrics.rs

@@ -11,7 +11,7 @@ use tracing::{error, info, instrument};
 
 use crate::proxy::ProxyRequest;
 use crate::utils::{full, ProxyResponse};
-use crate::{Consumer, State};
+use crate::State;
 
 #[derive(Debug, Clone)]
 pub struct Metrics {
@@ -69,75 +69,51 @@ impl Metrics {
     }
 
     pub fn count_ws_total_frame(&self, proxy_req: &ProxyRequest) {
-        let consumer = proxy_req
-            .consumer
-            .as_ref()
-            .unwrap_or(&Consumer::default())
-            .clone();
-
         self.ws_total_frame
             .with_label_values(&[
                 &proxy_req.namespace,
                 &proxy_req.instance,
                 &proxy_req.host,
-                &consumer.to_string(),
-                &consumer.tier,
+                &proxy_req.consumer.to_string(),
+                &proxy_req.consumer.tier,
             ])
             .inc()
     }
 
     pub fn inc_ws_total_connection(&self, proxy_req: &ProxyRequest) {
-        let consumer = proxy_req
-            .consumer
-            .as_ref()
-            .unwrap_or(&Consumer::default())
-            .clone();
-
         self.ws_total_connection
             .with_label_values(&[
                 &proxy_req.namespace,
                 &proxy_req.instance,
                 &proxy_req.host,
-                &consumer.to_string(),
-                &consumer.tier,
+                &proxy_req.consumer.to_string(),
+                &proxy_req.consumer.tier,
             ])
             .inc()
     }
 
     pub fn dec_ws_total_connection(&self, proxy_req: &ProxyRequest) {
-        let consumer = proxy_req
-            .consumer
-            .as_ref()
-            .unwrap_or(&Consumer::default())
-            .clone();
-
         self.ws_total_connection
             .with_label_values(&[
                 &proxy_req.namespace,
                 &proxy_req.instance,
                 &proxy_req.host,
-                &consumer.to_string(),
-                &consumer.tier,
+                &proxy_req.consumer.to_string(),
+                &proxy_req.consumer.tier,
             ])
             .dec()
     }
 
     pub fn count_http_total_request(&self, proxy_req: &ProxyRequest, status_code: StatusCode) {
-        let consumer = proxy_req
-            .consumer
-            .as_ref()
-            .unwrap_or(&Consumer::default())
-            .clone();
-
         self.http_total_request
             .with_label_values(&[
                 &proxy_req.namespace,
                 &proxy_req.instance,
                 &proxy_req.host,
                 &status_code.as_u16().to_string(),
                 &proxy_req.protocol.to_string(),
-                &consumer.to_string(),
-                &consumer.tier,
+                &proxy_req.consumer.to_string(),
+                &proxy_req.consumer.tier,
             ])
             .inc()
     }

proxy/src/proxy.rs

@@ -101,20 +101,13 @@ async fn handle(
         _ => {
             let proxy_req_result = ProxyRequest::new(&mut hyper_req, &state).await;
             if proxy_req_result.is_none() {
-                return Ok(Response::builder()
-                    .status(StatusCode::BAD_GATEWAY)
-                    .body(full("Invalid hostname"))
-                    .unwrap());
-            }
-
-            let proxy_req = proxy_req_result.unwrap();
-            if proxy_req.consumer.is_none() {
                 return Ok(Response::builder()
                     .status(StatusCode::UNAUTHORIZED)
                     .body(full("Unauthorized"))
                     .unwrap());
             }
 
+            let proxy_req = proxy_req_result.unwrap();
             let response_result = match proxy_req.protocol {
                 Protocol::Http => handle_http(hyper_req, &proxy_req).await,
                 Protocol::Websocket => handle_websocket(hyper_req, &proxy_req, state.clone()).await,
@@ -199,11 +192,7 @@ async fn handle_websocket(
                     while let Some(result) = client_incoming.next().await {
                         match result {
                             Ok(data) => {
-                                if let Err(err) = limiter(
-                                    state.clone(),
-                                    proxy_req.consumer.clone().as_ref().unwrap(),
-                                )
-                                .await
+                                if let Err(err) = limiter(state.clone(), &proxy_req.consumer).await
                                 {
                                     error!(error = err.to_string(), "Failed to run limiter");
                                     break;
@@ -276,7 +265,7 @@ pub struct ProxyRequest {
     pub namespace: String,
     pub host: String,
     pub instance: String,
-    pub consumer: Option<Consumer>,
+    pub consumer: Consumer,
     pub protocol: Protocol,
 }
 impl ProxyRequest {
@@ -288,11 +277,6 @@ impl ProxyRequest {
         let network = captures.get(2)?.as_str().to_string();
         let version = captures.get(3)?.as_str().to_string();
 
-        let instance = format!(
-            "ogmios-{network}-{version}.{}:{}",
-            state.config.ogmios_dns, state.config.ogmios_port
-        );
-
         let namespace = state.config.proxy_namespace.clone();
 
         let protocol = get_header(hyper_req, UPGRADE.as_str())
@@ -314,7 +298,11 @@ impl ProxyRequest {
         }
 
         let token = get_header(hyper_req, DMTR_API_KEY).unwrap_or_default();
-        let consumer = state.get_consumer(&network, &version, &token).await;
+        let consumer = state.get_consumer(&token).await?;
+        let instance = format!(
+            "ogmios-{network}-{version}.{}:{}",
+            state.config.ogmios_dns, state.config.ogmios_port
+        );
 
         Some(Self {
             namespace,