feat: implemented key to use the bech32 content

demeter-run · Sep 16, 2024 · 5b1ffa4 · 5b1ffa4
1 parent 2745c00
commit 5b1ffa4
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 34 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/operator/src/utils.rs b/operator/src/utils.rs
@@ -23,6 +23,7 @@ pub async fn patch_resource_status(
     let patch_params = PatchParams::default();
     api.patch_status(name, &patch_params, &Patch::Merge(status))
         .await?;
+
     Ok(())
 }
 
@@ -38,8 +39,6 @@ pub async fn build_api_key(crd: &CardanoNodePort) -> Result<String, Error> {
     let namespace = crd.namespace().unwrap();
     let name = format!("cardano-node-auth-{}", &crd.name_any());
 
-    let network = &crd.spec.network;
-    let version = &crd.spec.version;
     let password = format!("{}{}", name, namespace).as_bytes().to_vec();
 
     let config = get_config();
@@ -50,7 +49,7 @@ pub async fn build_api_key(crd: &CardanoNodePort) -> Result<String, Error> {
     let argon2 = Argon2::default();
     argon2.hash_password_into(password.as_slice(), salt, &mut output)?;
 
-    let hrp = Hrp::parse(&format!("dmtr_cnode_{version}_{network}_"))?;
+    let hrp = Hrp::parse("cnode")?;
     let with_bech = bech32::encode::<Bech32m>(hrp, output.as_slice())?;
 
     Ok(with_bech)

diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml
@@ -22,3 +22,4 @@ serde = { version = "1.0.197", features = ["derive"] }
 serde_json = "1.0.114"
 toml = "0.8.10"
 notify = "6.1.1"
+bech32 = "0.11.0"
diff --git a/proxy/src/auth.rs b/proxy/src/auth.rs
@@ -50,9 +50,15 @@ impl BackgroundService for AuthBackgroundService {
                 Ok(Some(Event::Restarted(crds))) => {
                     info!("auth: Watcher restarted, reseting consumers");
 
-                    let mut consumers: HashMap<String, Consumer> = Default::default();
+                    let mut consumers: HashMap<Vec<u8>, Consumer> = Default::default();
                     for crd in crds.iter() {
-                        let consumer = self.sync_consumer(crd.into()).await;
+                        let result = Consumer::new(crd);
+                        if let Err(error) = result {
+                            error!(?error, "invalid consumer");
+                            continue;
+                        }
+
+                        let consumer = self.sync_consumer(result.unwrap()).await;
                         consumers.insert(consumer.key.clone(), consumer);
                     }
 
@@ -64,7 +70,13 @@ impl BackgroundService for AuthBackgroundService {
                     Some(_) => {
                         info!("auth: Adding new consumer: {}", crd.name_any());
 
-                        let consumer = self.sync_consumer((&crd).into()).await;
+                        let result = Consumer::new(&crd);
+                        if let Err(error) = result {
+                            error!(?error, "invalid consumer");
+                            continue;
+                        }
+
+                        let consumer = self.sync_consumer(result.unwrap()).await;
 
                         self.state.limiter.write().await.remove(&consumer.key);
                         self.state
@@ -85,7 +97,13 @@ impl BackgroundService for AuthBackgroundService {
                         "auth: Port deleted, removing from state: {}",
                         crd.name_any()
                     );
-                    let consumer = Consumer::from(&crd);
+
+                    let result = Consumer::new(&crd);
+                    if let Err(error) = result {
+                        error!(?error, "invalid consumer");
+                        continue;
+                    }
+                    let consumer = result.unwrap();
                     self.state.consumers.write().await.remove(&consumer.key);
                     self.state.limiter.write().await.remove(&consumer.key);
                 }

diff --git a/proxy/src/main.rs b/proxy/src/main.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, fmt::Display, sync::Arc, time::Duration};
+use std::{collections::HashMap, error::Error, fmt::Display, sync::Arc, time::Duration};
 
 use auth::AuthBackgroundService;
 use dotenv::dotenv;
@@ -80,16 +80,16 @@ fn main() {
 #[derive(Default)]
 pub struct State {
     metrics: Metrics,
-    consumers: RwLock<HashMap<String, Consumer>>,
-    limiter: RwLock<HashMap<String, Vec<Arc<RateLimiter>>>>,
+    consumers: RwLock<HashMap<Vec<u8>, Consumer>>,
+    limiter: RwLock<HashMap<Vec<u8>, Vec<Arc<RateLimiter>>>>,
     tiers: RwLock<HashMap<String, Tier>>,
 }
 impl State {
     pub fn new() -> Self {
         Self::default()
     }
 
-    pub async fn get_consumer(&self, key: &str) -> Option<Consumer> {
+    pub async fn get_consumer(&self, key: &[u8]) -> Option<Consumer> {
         let consumers = self.consumers.read().await.clone();
         consumers.get(key).cloned()
     }
@@ -100,12 +100,32 @@ pub struct Consumer {
     namespace: String,
     port_name: String,
     tier: String,
-    key: String,
+    key: Vec<u8>,
     network: String,
     version: String,
     active_connections: usize,
 }
 impl Consumer {
+    pub fn new(crd: &CardanoNodePort) -> Result<Self, Box<dyn Error>> {
+        let network = crd.spec.network.to_string();
+        let version = crd.spec.version.to_string();
+        let tier = crd.spec.throughput_tier.to_string();
+        let key = crd.status.as_ref().unwrap().auth_token.clone();
+        let namespace = crd.metadata.namespace.as_ref().unwrap().clone();
+        let port_name = crd.name_any();
+
+        let (_hrp, key) = bech32::decode(&key)?;
+
+        Ok(Self {
+            namespace,
+            port_name,
+            tier,
+            key,
+            network,
+            version,
+            active_connections: 0,
+        })
+    }
     pub async fn inc_connections(&self, state: Arc<State>) {
         state
             .consumers
@@ -137,26 +157,6 @@ impl Display for Consumer {
         write!(f, "{}.{}", self.namespace, self.port_name)
     }
 }
-impl From<&CardanoNodePort> for Consumer {
-    fn from(value: &CardanoNodePort) -> Self {
-        let network = value.spec.network.to_string();
-        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();
-
-        Self {
-            namespace,
-            port_name,
-            tier,
-            key,
-            network,
-            version,
-            active_connections: 0,
-        }
-    }
-}
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct Tier {

diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs
@@ -46,7 +46,7 @@ impl ProxyApp {
     pub fn new(config: Arc<Config>, state: Arc<State>) -> Self {
         ProxyApp {
             client_connector: TransportConnector::new(None),
-            host_regex: Regex::new(r"(dmtr_[\w\d-]+)\..+").unwrap(),
+            host_regex: Regex::new(r"([\w\d-]+)\..+").unwrap(),
             config,
             state,
         }
@@ -241,7 +241,13 @@ impl ServerApp for ProxyApp {
 
         let token = captures.get(1)?.as_str().to_string();
 
-        let consumer = self.state.get_consumer(&token).await?;
+        let result = bech32::decode(&token);
+        if let Err(error) = result {
+            error!(?error, "invalid bech32");
+            return None;
+        }
+
+        let consumer = self.state.get_consumer(&result.unwrap().1).await?;
         let instance = format!(
             "node-{}-{}.{}:{}",
             consumer.network, consumer.version, self.config.node_dns, self.config.node_port