fix: Update on ports now work on active connections (#51)

* fix: Update on ports now work on active connections

* Keep consumer on request, but only use key for limiting

Author: gonzalezzfelipe

Diff for: proxy/src/limiter.rs

@@ -1,12 +1,28 @@
 use futures_util::future::join_all;
 use leaky_bucket::RateLimiter;
 use std::sync::Arc;
+use std::{error::Error, fmt::Display};
 
 use crate::{tiers::Tier, Consumer, State};
 
-async fn has_limiter(state: &State, consumer: &Consumer) -> bool {
+#[derive(Debug)]
+pub enum LimiterError {
+    PortDeleted,
+    InvalidTier,
+}
+impl Display for LimiterError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            LimiterError::PortDeleted => f.write_str("Port was deleted"),
+            LimiterError::InvalidTier => f.write_str("Tier is invalid"),
+        }
+    }
+}
+impl Error for LimiterError {}
+
+async fn has_limiter(state: &State, consumer_key: &String) -> bool {
     let rate_limiter_map = state.limiter.read().await;
-    rate_limiter_map.get(&consumer.key).is_some()
+    rate_limiter_map.get(consumer_key).is_some()
 }
 
 async fn add_limiter(state: &State, consumer: &Consumer, tier: &Tier) {
@@ -31,16 +47,24 @@ async fn add_limiter(state: &State, consumer: &Consumer, tier: &Tier) {
         .insert(consumer.key.clone(), rates);
 }
 
-pub async fn limiter(state: Arc<State>, consumer: &Consumer) {
-    let tiers = state.tiers.read().await.clone();
-    let tier = tiers.get(&consumer.tier).unwrap();
-
-    if !has_limiter(&state, consumer).await {
+pub async fn limiter(state: Arc<State>, consumer_key: String) -> Result<(), LimiterError> {
+    if !has_limiter(&state, &consumer_key).await {
+        let consumers = state.consumers.read().await.clone();
+        let consumer = match consumers.get(&consumer_key) {
+            Some(consumer) => consumer,
+            None => return Err(LimiterError::PortDeleted),
+        };
+        let tiers = state.tiers.read().await.clone();
+        let tier = match tiers.get(&consumer.tier) {
+            Some(tier) => tier,
+            None => return Err(LimiterError::InvalidTier),
+        };
         add_limiter(&state, consumer, tier).await;
     }
 
     let rate_limiter_map = state.limiter.read().await.clone();
-    let rates = rate_limiter_map.get(&consumer.key).unwrap();
+    let rates = rate_limiter_map.get(&consumer_key).unwrap();
 
     join_all(rates.iter().map(|r| async { r.acquire_one().await })).await;
+    Ok(())
 }

Diff for: proxy/src/proxy.rs

@@ -199,7 +199,13 @@ async fn handle_websocket(
                     while let Some(result) = client_incoming.next().await {
                         match result {
                             Ok(data) => {
-                                limiter(state.clone(), proxy_req.consumer.as_ref().unwrap()).await;
+                                if let Err(err) =
+                                    limiter(state.clone(), proxy_req.consumer.clone().unwrap().key)
+                                        .await
+                                {
+                                    error!(error = err.to_string(), "Failed to run limiter");
+                                    break;
+                                };
                                 if let Err(err) = instance_outgoing.send(data).await {
                                     error!(
                                         error = err.to_string(),