Upstreams from store

.gitignore

@@ -6,6 +6,7 @@
 /proksi.hcl
 .zed/
 .vscode/
+.idea/
 /tmp
 docker.README.md
 docker-compose.yml

crates/proksi/src/stores/memory_store.rs

@@ -1,45 +1,58 @@
 use async_trait::async_trait;
 use papaya::HashMapRef;
 use std::{error::Error, hash::RandomState};
+use crate::config::RouteUpstream;
 
 use super::certificates::Certificate;
 use super::store_trait::Store;
 
 pub struct MemoryStore {
     /// Map of domain names to certificates (including leaf & chain)
-    inner_certs: papaya::HashMap<String, Certificate>,
+    certs: papaya::HashMap<String, Certificate>,
     /// Map of domain names to challenge tokens and proofs (token, proof)
-    inner_challenges: papaya::HashMap<String, (String, String)>,
+    challenges: papaya::HashMap<String, (String, String)>,
+    /// Map of domain names to routes upstreams
+    upstreams: papaya::HashMap<String, Vec<RouteUpstream>>,
 }
 
 impl MemoryStore {
     pub fn new() -> Self {
         MemoryStore {
-            inner_certs: papaya::HashMap::new(),
-            inner_challenges: papaya::HashMap::new(),
+            certs: papaya::HashMap::new(),
+            challenges: papaya::HashMap::new(),
+            upstreams: papaya::HashMap::new(),
         }
     }
 }
 
 #[async_trait]
 impl Store for MemoryStore {
+    async fn get_upstreams(&self, domain: &str) -> Option<Vec<RouteUpstream>> {
+        self.upstreams.pin().get(domain).cloned()
+    }
+
+    async fn set_upstreams(&self, domain: &str, upstreams: Vec<RouteUpstream>) -> Result<(), Box<dyn Error>> {
+        self.upstreams.pin().insert(domain.to_string(), upstreams);
+        Ok(())
+    }
+
     async fn get_certificate(&self, host: &str) -> Option<Certificate> {
-        self.inner_certs.pin().get(host).cloned()
+        self.certs.pin().get(host).cloned()
     }
 
     async fn set_certificate(&self, host: &str, cert: Certificate) -> Result<(), Box<dyn Error>> {
-        self.inner_certs.pin().insert(host.to_string(), cert);
+        self.certs.pin().insert(host.to_string(), cert);
         Ok(())
     }
 
     async fn get_certificates(
         &self,
     ) -> HashMapRef<'_, String, Certificate, RandomState, seize::LocalGuard<'_>> {
-        self.inner_certs.pin()
+        self.certs.pin()
     }
 
     async fn get_challenge(&self, domain: &str) -> Option<(String, String)> {
-        self.inner_challenges.pin().get(domain).cloned()
+        self.challenges.pin().get(domain).cloned()
     }
 
     async fn set_challenge(
@@ -48,7 +61,7 @@ impl Store for MemoryStore {
         token: String,
         proof: String,
     ) -> Result<(), Box<dyn Error>> {
-        self.inner_challenges
+        self.challenges
             .pin()
             .insert(domain.to_string(), (token, proof));
         Ok(())
@@ -57,6 +70,7 @@ impl Store for MemoryStore {
 
 #[cfg(test)]
 mod tests {
+    use std::borrow::Cow;
     // No need to import super since we're using specific imports
     use crate::stores::certificates::Certificate;
     use crate::stores::store_trait::Store;
@@ -67,6 +81,7 @@ mod tests {
         rsa::Rsa,
         x509::{X509Name, X509},
     };
+    use crate::config::RouteUpstream;
 
     // Helper function to create a test certificate
     fn create_test_certificate(domain: &str) -> Certificate {
@@ -210,4 +225,41 @@ mod tests {
         let challenge = store.get_challenge("nonexistent.com").await;
         assert!(challenge.is_none());
     }
+
+    #[tokio::test]
+    async fn test_upstreams_storage() {
+        let store = MemoryStore::new();
+
+        let domain = "example.com";
+
+        let upstream1 = RouteUpstream {
+            ip: Cow::Borrowed("127.0.0.1"),
+            port: 80,
+            network: None,
+            weight: None,
+            sni: None,
+            headers: None,
+        };
+
+        let upstream2 = RouteUpstream {
+            ip: Cow::Borrowed("127.0.0.1"),
+            port: 443,
+            network: None,
+            weight: None,
+            sni: None,
+            headers: None,
+        };
+
+        let upstreams = vec!(
+            upstream1, upstream2
+        );
+
+        store.set_upstreams(domain, upstreams)
+            .await
+            .unwrap();
+
+        let store_upstreams = store.get_upstreams(domain).await.unwrap();
+
+        assert_eq!(store_upstreams.len(), 2);
+    }
 }

crates/proksi/src/stores/redis_store.rs

@@ -3,6 +3,7 @@ use papaya::HashMapRef;
 use redis::{Client, Commands};
 use serde_json;
 use std::{error::Error, hash::RandomState};
+use crate::config::RouteUpstream;
 
 use super::certificates::{Certificate, SerializableCertificate};
 use super::store_trait::Store;
@@ -11,8 +12,9 @@ const CHALLENGE_TTL_SECONDS: u64 = 300;
 
 pub struct RedisStore {
     pool: r2d2::Pool<redis::Client>,
-    cache: papaya::HashMap<String, Certificate>,
+    certificate_cache: papaya::HashMap<String, Certificate>,
     challenge_cache: papaya::HashMap<String, (String, String)>,
+    upstream_cache: papaya::HashMap<String, Vec<RouteUpstream>>,
 }
 
 impl RedisStore {
@@ -22,11 +24,16 @@ impl RedisStore {
 
         Ok(RedisStore {
             pool,
-            cache: papaya::HashMap::new(),
+            certificate_cache: papaya::HashMap::new(),
             challenge_cache: papaya::HashMap::new(),
+            upstream_cache: papaya::HashMap::new(),
         })
     }
 
+    fn upstreams_key(domain: &str) -> String {
+        format!("proksi:upstreams:{domain}")
+    }
+
     fn certificate_key(domain: &str) -> String {
         format!("proksi:cert:{domain}")
     }
@@ -35,7 +42,7 @@ impl RedisStore {
         format!("proksi:challenge:{domain}")
     }
 
-    fn load_from_redis(&self, domain: &str) -> Option<Certificate> {
+    fn load_certificate_from_redis(&self, domain: &str) -> Option<Certificate> {
         let mut conn = self.pool.get().unwrap();
         let key = Self::certificate_key(domain);
 
@@ -66,16 +73,46 @@ impl RedisStore {
 
 #[async_trait]
 impl Store for RedisStore {
+    async fn get_upstreams(&self, domain: &str) -> Option<Vec<RouteUpstream>> {
+
+        let mut conn = self.pool.get().unwrap();
+        let key = Self::upstreams_key(domain);
+
+        let upstreams_data: Option<String> = conn.get(&key).ok()?;
+
+        if let Some(data) = upstreams_data {
+            if let Ok(upstreams) = serde_json::from_str::<Vec<RouteUpstream>>(&data) {
+                return Some(upstreams);
+            }
+        }
+        None
+    }
+
+    async fn set_upstreams(&self, domain: &str, upstreams: Vec<RouteUpstream>) -> Result<(), Box<dyn Error>> {
+        let mut conn = self.pool.get()?;
+        let key = Self::certificate_key(domain);
+
+        let upstreams_json = serde_json::to_string(&upstreams.clone())?;
+
+        conn.set::<String, String, String>(key, upstreams_json)?;
+
+        // Update cache
+        self.upstream_cache.pin().insert(domain.to_string(), upstreams);
+
+        Ok(())
+    }
+
+
     async fn get_certificate(&self, domain: &str) -> Option<Certificate> {
         // Check cache first
-        if let Some(cert) = self.cache.pin().get(domain) {
+        if let Some(cert) = self.certificate_cache.pin().get(domain) {
             return Some(cert.clone());
         }
 
         // If not in cache, load from Redis
-        if let Some(cert) = self.load_from_redis(domain) {
+        if let Some(cert) = self.load_certificate_from_redis(domain) {
             // Store in cache for future use
-            self.cache.pin().insert(domain.to_string(), cert.clone());
+            self.certificate_cache.pin().insert(domain.to_string(), cert.clone());
             return Some(cert);
         }
 
@@ -93,7 +130,7 @@ impl Store for RedisStore {
         conn.set::<String, String, String>(key, cert_json)?;
 
         // Update cache
-        self.cache.pin().insert(domain.to_string(), cert);
+        self.certificate_cache.pin().insert(domain.to_string(), cert);
 
         Ok(())
     }
@@ -107,7 +144,7 @@ impl Store for RedisStore {
         TEMP_MAP.pin().clear();
 
         // First, copy all cached certificates
-        for (key, value) in &self.cache.pin() {
+        for (key, value) in &self.certificate_cache.pin() {
             TEMP_MAP.pin().insert(key.clone(), value.clone());
         }
 
@@ -130,7 +167,7 @@ impl Store for RedisStore {
                             if let Ok(cert) = Certificate::from_serializable(serializable_cert) {
                                 TEMP_MAP.pin().insert(domain.clone(), cert.clone());
                                 // Update cache with newly found certificate
-                                self.cache.pin().insert(domain, cert);
+                                self.certificate_cache.pin().insert(domain, cert);
                             }
                         }
                     }

crates/proksi/src/stores/store_trait.rs

@@ -1,6 +1,7 @@
 use async_trait::async_trait;
 use papaya::HashMapRef;
 use std::{error::Error, hash::RandomState};
+use crate::config::RouteUpstream;
 
 use super::certificates::Certificate;
 
@@ -9,6 +10,10 @@
     // async fn get_route(&self, host: &str) -> Result<Option<String>, Box<dyn Error>>;
     // async fn remove_route(&self, host: &str) -> Result<(), Box<dyn Error>>;
     // async fn set_route(&self, route: &str) -> Result<(), Box<dyn Error>>;
+
+    async fn get_upstreams(&self, domain: &str) -> Option<Vec<RouteUpstream>>;
+    async fn set_upstreams(&self, domain: &str, upstreams: Vec<RouteUpstream>) -> Result<(), Box<dyn Error>>;
+
     async fn get_certificate(&self, domain: &str) -> Option<Certificate>;
     async fn set_certificate(&self, domain: &str, cert: Certificate) -> Result<(), Box<dyn Error>>;
     async fn get_certificates(