refactor: introduce EnvironmentIndex to make the served environment set runtime-mutable

Cargo.toml

@@ -9,6 +9,7 @@ license = "BSD-3-Clause"
 axum = { version = "0.8", features = ["macros"] }
 tokio = { version = "1", features = ["full"] }
 tower = "0.5"
+parking_lot = "0.12"
 tower-http = { version = "0.6", features = ["compression-gzip", "cors", "trace", "normalize-path"] }
 reqwest = { version = "0.13", features = ["json", "gzip"] }
 serde = { version = "1", features = ["derive"] }

src/cache/environment.rs

@@ -20,6 +20,10 @@ pub trait EnvironmentsCache: Send + Sync {
     /// Store environment document and compute context. Returns true if changed.
     async fn put_environment(&self, environment_key: &str, document: Value) -> bool;
 
+    /// Remove everything stored for an environment (document, context,
+    /// identity overrides)
+    async fn remove_environment(&self, environment_key: &str);
+
     /// Get identity override data
     async fn get_identity(&self, environment_api_key: &str, identifier: &str) -> Option<Value>;
 }
@@ -98,10 +102,67 @@ impl EnvironmentsCache for LocalMemEnvironmentsCache {
         changed
     }
 
+    async fn remove_environment(&self, environment_key: &str) {
+        // Same guard order as put_environment so the two can't deadlock
+        let mut environments = self.environments.write().await;
+        let mut contexts = self.contexts.write().await;
+        let mut identity_overrides = self.identity_overrides.write().await;
+
+        environments.remove(environment_key);
+        contexts.remove(environment_key);
+        identity_overrides.remove(environment_key);
+    }
+
     async fn get_identity(&self, environment_api_key: &str, identifier: &str) -> Option<Value> {
         let identity_overrides = self.identity_overrides.read().await;
         identity_overrides
             .get(environment_api_key)
             .and_then(|identities| identities.get(identifier).cloned())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[tokio::test]
+    async fn test_remove_environment_clears_all_stored_state() {
+        // Given
+        let cache = LocalMemEnvironmentsCache::new();
+        let document = json!({
+            "api_key": "client_a",
+            "identity_overrides": [{"identifier": "user_1"}],
+        });
+        cache.put_environment("client_a", document).await;
+        assert!(cache.get_environment("client_a").await.is_some());
+        assert!(cache.get_identity("client_a", "user_1").await.is_some());
+
+        // When
+        cache.remove_environment("client_a").await;
+
+        // Then
+        assert!(cache.get_environment("client_a").await.is_none());
+        assert!(cache.get_context("client_a").await.is_none());
+        assert!(cache.get_identity("client_a", "user_1").await.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_remove_environment_leaves_other_environments_alone() {
+        // Given
+        let cache = LocalMemEnvironmentsCache::new();
+        cache
+            .put_environment("client_a", json!({"api_key": "client_a"}))
+            .await;
+        cache
+            .put_environment("client_b", json!({"api_key": "client_b"}))
+            .await;
+
+        // When
+        cache.remove_environment("client_a").await;
+
+        // Then
+        assert!(cache.get_environment("client_a").await.is_none());
+        assert!(cache.get_environment("client_b").await.is_some());
+    }
+}

src/config/settings.rs

@@ -91,6 +91,8 @@ impl Default for HealthCheckSettings {
 
 #[derive(Debug, Clone, Serialize, Deserialize, Validate)]
 pub struct AppSettings {
+    // Optional so the environment set can also come from runtime discovery
+    #[serde(default)]
     #[validate(nested)]
     pub environment_key_pairs: Vec<EnvironmentKeyPair>,
     #[serde(default = "default_api_url")]
@@ -170,6 +172,16 @@ pub fn get_settings() -> Result<AppSettings> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_config_without_environment_key_pairs_parses_to_empty_valid_set() {
+        // Given a config file omitting environment_key_pairs entirely
+        let settings: AppSettings = serde_json::from_str("{}").unwrap();
+
+        // Then it behaves like an explicitly empty list and validates
+        assert!(settings.environment_key_pairs.is_empty());
+        assert!(settings.validate().is_ok());
+    }
+
     #[test]
     fn test_client_side_key_validation_valid() {
         // Given

src/environments.rs

@@ -0,0 +1,256 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use chrono::{DateTime, Utc};
+use parking_lot::RwLock;
+
+use crate::config::settings::EnvironmentKeyPair;
+
+/// A server-side (`ser.`) key together with the validity metadata the
+/// proxy config endpoint reports. Statically configured keys carry no
+/// metadata and are always valid.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ServerKey {
+    pub key: String,
+    pub active: bool,
+    pub expires_at: Option<DateTime<Utc>>,
+}
+
+impl ServerKey {
+    pub fn is_valid(&self) -> bool {
+        self.active && self.expires_at.is_none_or(|at| at > Utc::now())
+    }
+}
+
+/// The key set of one environment the proxy serves: its client-side key
+/// and every server-side key that can authenticate for it upstream
+/// (multiple during rotation).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct EnvironmentKeys {
+    pub client_key: String,
+    pub server_keys: Vec<ServerKey>,
+}
+
+impl EnvironmentKeys {
+    /// The first server-side key still usable for upstream fetches.
+    pub fn valid_server_key(&self) -> Option<&ServerKey> {
+        self.server_keys.iter().find(|key| key.is_valid())
+    }
+}
+
+/// The runtime-mutable set of environments the proxy serves.
+///
+/// Uses `parking_lot::RwLock`, not tokio's: guards are held only for a map
+/// operation, never across an await, and lookups stay callable from
+/// synchronous code.
+#[derive(Default)]
+pub struct EnvironmentIndex {
+    /// One entry per key an environment owns, client and server alike, all
+    /// pointing at the same record, so a lookup takes whichever key a
+    /// request presents.
+    environment_keys_by_any_key: RwLock<HashMap<String, Arc<EnvironmentKeys>>>,
+}
+
+impl EnvironmentIndex {
+    pub fn from_settings(pairs: &[EnvironmentKeyPair]) -> Self {
+        let index = Self::default();
+        for pair in pairs {
+            index.insert(EnvironmentKeys {
+                client_key: pair.client_side_key.clone(),
+                server_keys: vec![ServerKey {
+                    key: pair.server_side_key.clone(),
+                    active: true,
+                    expires_at: None,
+                }],
+            });
+        }
+        index
+    }
+
+    /// Resolve a presented key — client- or server-side — to its
+    /// environment's keys.
+    pub fn resolve(&self, key: &str) -> Option<Arc<EnvironmentKeys>> {
+        self.environment_keys_by_any_key.read().get(key).cloned()
+    }
+
+    /// Insert or replace an environment's keys. Server keys the
+    /// environment no longer has stop resolving.
+    pub fn insert(&self, keys: EnvironmentKeys) {
+        let keys = Arc::new(keys);
+        let mut environment_keys_by_any_key = self.environment_keys_by_any_key.write();
+
+        if let Some(previous) = environment_keys_by_any_key.get(&keys.client_key).cloned() {
+            for server_key in &previous.server_keys {
+                environment_keys_by_any_key.remove(&server_key.key);
+            }
+        }
+
+        for server_key in &keys.server_keys {
+            environment_keys_by_any_key.insert(server_key.key.clone(), Arc::clone(&keys));
+        }
+        environment_keys_by_any_key.insert(keys.client_key.clone(), keys);
+    }
+
+    /// Remove the environment `key` resolves to (any of its keys works),
+    /// returning its keys so the caller can clear per-key caches.
+    pub fn remove(&self, key: &str) -> Option<Arc<EnvironmentKeys>> {
+        let mut environment_keys_by_any_key = self.environment_keys_by_any_key.write();
+        let keys = environment_keys_by_any_key.get(key).cloned()?;
+
+        environment_keys_by_any_key.remove(&keys.client_key);
+        for server_key in &keys.server_keys {
+            environment_keys_by_any_key.remove(&server_key.key);
+        }
+
+        Some(keys)
+    }
+
+    /// Point-in-time snapshot of every environment's keys, ordered by
+    /// client key so callers iterate deterministically.
+    pub fn snapshot(&self) -> Vec<Arc<EnvironmentKeys>> {
+        let environment_keys_by_any_key = self.environment_keys_by_any_key.read();
+        let mut snapshot: Vec<Arc<EnvironmentKeys>> = environment_keys_by_any_key
+            .iter()
+            .filter(|(key, keys)| key.as_str() == keys.client_key)
+            .map(|(_, keys)| Arc::clone(keys))
+            .collect();
+        snapshot.sort_by(|a, b| a.client_key.cmp(&b.client_key));
+        snapshot
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use chrono::TimeDelta;
+
+    fn pair(client: &str, server: &str) -> EnvironmentKeyPair {
+        EnvironmentKeyPair {
+            client_side_key: client.to_string(),
+            server_side_key: server.to_string(),
+        }
+    }
+
+    fn server_key(key: &str) -> ServerKey {
+        ServerKey {
+            key: key.to_string(),
+            active: true,
+            expires_at: None,
+        }
+    }
+
+    #[test]
+    fn from_settings_resolves_both_keys_to_the_same_environment() {
+        // Given
+        let index = EnvironmentIndex::from_settings(&[pair("client_a", "ser.a")]);
+
+        // When
+        let by_client = index.resolve("client_a").unwrap();
+        let by_server = index.resolve("ser.a").unwrap();
+
+        // Then
+        assert!(Arc::ptr_eq(&by_client, &by_server));
+        assert_eq!(by_client.client_key, "client_a");
+        assert!(by_client.valid_server_key().is_some());
+    }
+
+    #[test]
+    fn resolve_unknown_key_returns_none() {
+        let index = EnvironmentIndex::from_settings(&[pair("client_a", "ser.a")]);
+        assert!(index.resolve("nope").is_none());
+    }
+
+    #[test]
+    fn insert_replaces_keys_and_drops_stale_server_key_index() {
+        // Given
+        let index = EnvironmentIndex::from_settings(&[pair("client_a", "ser.old")]);
+
+        // When the environment's server key is rotated
+        index.insert(EnvironmentKeys {
+            client_key: "client_a".to_string(),
+            server_keys: vec![server_key("ser.new")],
+        });
+
+        // Then only the new server key resolves
+        assert!(index.resolve("ser.old").is_none());
+        assert_eq!(index.resolve("ser.new").unwrap().client_key, "client_a");
+        assert_eq!(index.snapshot().len(), 1);
+    }
+
+    #[test]
+    fn remove_by_any_key_clears_every_index_entry() {
+        // Given an environment with two server keys
+        let index = EnvironmentIndex::default();
+        index.insert(EnvironmentKeys {
+            client_key: "client_a".to_string(),
+            server_keys: vec![server_key("ser.one"), server_key("ser.two")],
+        });
+
+        // When removed via one of its server keys
+        let removed = index.remove("ser.two").unwrap();
+
+        // Then
+        assert_eq!(removed.client_key, "client_a");
+        assert!(index.resolve("client_a").is_none());
+        assert!(index.resolve("ser.one").is_none());
+        assert!(index.resolve("ser.two").is_none());
+        assert!(index.remove("client_a").is_none());
+    }
+
+    #[test]
+    fn snapshot_returns_one_entry_per_environment_sorted_by_client_key() {
+        // Given
+        let index = EnvironmentIndex::from_settings(&[
+            pair("client_b", "ser.b"),
+            pair("client_a", "ser.a"),
+        ]);
+
+        // When
+        let snapshot = index.snapshot();
+
+        // Then
+        let client_keys: Vec<&str> = snapshot.iter().map(|r| r.client_key.as_str()).collect();
+        assert_eq!(client_keys, vec!["client_a", "client_b"]);
+    }
+
+    #[test]
+    fn valid_server_key_skips_inactive_and_expired_keys() {
+        // Given
+        let keys = EnvironmentKeys {
+            client_key: "client_a".to_string(),
+            server_keys: vec![
+                ServerKey {
+                    key: "ser.inactive".to_string(),
+                    active: false,
+                    expires_at: None,
+                },
+                ServerKey {
+                    key: "ser.expired".to_string(),
+                    active: true,
+                    expires_at: Some(Utc::now() - TimeDelta::days(1)),
+                },
+                ServerKey {
+                    key: "ser.valid".to_string(),
+                    active: true,
+                    expires_at: Some(Utc::now() + TimeDelta::days(1)),
+                },
+            ],
+        };
+
+        // When / Then
+        assert_eq!(keys.valid_server_key().unwrap().key, "ser.valid");
+    }
+
+    #[test]
+    fn valid_server_key_returns_none_when_no_key_is_usable() {
+        let keys = EnvironmentKeys {
+            client_key: "client_a".to_string(),
+            server_keys: vec![ServerKey {
+                key: "ser.inactive".to_string(),
+                active: false,
+                expires_at: None,
+            }],
+        };
+        assert!(keys.valid_server_key().is_none());
+    }
+}

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod cache;
 pub mod config;
+pub mod environments;
 pub mod error;
 pub mod models;
 pub mod routes;