diff --git a/src/config/settings.rs b/src/config/settings.rs index 64bca8f..a812d96 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -126,6 +126,11 @@ pub struct AppSettings { #[serde(default)] #[validate(nested)] pub environment_key_pairs: Vec, + /// Credential for the proxy config endpoint; when set, the served + /// environments are kept in sync with it at every poll. + #[serde(default)] + #[validate(length(min = 1))] + pub proxy_key: Option, #[serde(default = "default_api_url")] pub api_url: String, #[serde(default = "default_api_poll_frequency")] @@ -158,6 +163,7 @@ impl Default for AppSettings { fn default() -> Self { Self { environment_key_pairs: vec![], + proxy_key: None, api_url: default_api_url(), api_poll_frequency_seconds: default_api_poll_frequency(), api_poll_timeout_seconds: default_api_poll_timeout(), @@ -206,6 +212,26 @@ pub fn get_settings() -> Result { mod tests { use super::*; + #[test] + fn test_config_with_empty_proxy_key_is_invalid() { + // Given an empty proxy_key, which would silently fail every sync + let settings: AppSettings = serde_json::from_str(r#"{"proxy_key": ""}"#).unwrap(); + + // Then + assert!(settings.validate().is_err()); + } + + #[test] + fn test_config_with_only_a_proxy_key_is_valid() { + // Given a config file that relies entirely on the proxy config + let settings: AppSettings = serde_json::from_str(r#"{"proxy_key": "pk.secret"}"#).unwrap(); + + // Then + assert_eq!(settings.proxy_key.as_deref(), Some("pk.secret")); + assert!(settings.environment_key_pairs.is_empty()); + assert!(settings.validate().is_ok()); + } + #[test] fn test_config_without_environment_key_pairs_parses_to_empty_valid_set() { // Given a config file omitting environment_key_pairs entirely diff --git a/src/environments.rs b/src/environments.rs index cfdb34f..cb2d61e 100644 --- a/src/environments.rs +++ b/src/environments.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, RwLock}; use chrono::{DateTime, Utc}; @@ -37,6 +37,17 @@ impl EnvironmentKeys { } } +/// What a `sync_to` call did, for logging and cache invalidation. +#[derive(Debug, Default)] +pub struct SyncResult { + /// Environments inserted or updated. + pub changed: usize, + /// Previous versions replaced by an update. + pub replaced: Vec>, + /// Environments no longer in the config. + pub removed: Vec>, +} + /// The runtime-mutable set of environments the proxy serves. /// /// Every environment is indexed under its client key *and* each of its @@ -53,11 +64,19 @@ impl EnvironmentKeys { #[derive(Default)] pub struct EnvironmentIndex { by_key: RwLock>>, + /// Every key of the statically configured environments. Immutable + /// after construction; `sync_to` never overrides or removes an + /// environment whose keys appear here. + protected: HashSet, } impl EnvironmentIndex { pub fn from_settings(pairs: &[EnvironmentKeyPair]) -> Self { - let index = Self::default(); + let mut index = Self::default(); + for pair in pairs { + index.protected.insert(pair.client_side_key.clone()); + index.protected.insert(pair.server_side_key.clone()); + } for pair in pairs { index.insert(EnvironmentKeys { client_key: pair.client_side_key.clone(), @@ -72,13 +91,28 @@ impl EnvironmentIndex { } /// Resolve a presented key — client- or server-side — to its - /// environment's keys. + /// environment's keys. A server-side key resolves only while it is + /// valid, so a deactivation delivered by the proxy config and an + /// expiry passing between polls both take effect on the next request. pub fn resolve(&self, key: &str) -> Option> { - self.by_key + let keys = self + .by_key .read() .expect("environment index lock poisoned") .get(key) - .cloned() + .cloned()?; + + if key != keys.client_key { + let presented = keys + .server_keys + .iter() + .find(|server_key| server_key.key == key)?; + if !presented.is_valid() { + return None; + } + } + + Some(keys) } /// Insert or replace an environment's keys, dropping index entries @@ -124,6 +158,62 @@ impl EnvironmentIndex { Some(keys) } + /// Bring the index in line with `desired` (what the proxy config + /// reports). Statically configured environments are never overridden + /// or removed, and a desired environment whose keys collide with a + /// static environment's keys is skipped entirely. Safe against + /// concurrent readers; assumes a single writer — the poll task is the + /// sole caller — so per-operation locking suffices. + /// + /// The caller owns cache invalidation for everything returned in + /// `replaced` and `removed` (request caches are keyed by presented + /// key and consulted before the key gate). + pub fn sync_to(&self, desired: Vec) -> SyncResult { + let mut result = SyncResult::default(); + + let desired_clients: HashSet = + desired.iter().map(|keys| keys.client_key.clone()).collect(); + + for keys in desired { + if self.is_protected(&keys) { + continue; + } + let unchanged = self + .resolve(&keys.client_key) + .is_some_and(|current| *current == keys); + if unchanged { + continue; + } + result.changed += 1; + if let Some(previous) = self.insert(keys) { + result.replaced.push(previous); + } + } + + for current in self.snapshot() { + if desired_clients.contains(¤t.client_key) + || self.protected.contains(¤t.client_key) + { + continue; + } + if let Some(removed) = self.remove(¤t.client_key) { + result.removed.push(removed); + } + } + + result + } + + /// True when the environment is statically configured, or any of its + /// keys collides with a static environment's key namespace. + fn is_protected(&self, keys: &EnvironmentKeys) -> bool { + self.protected.contains(&keys.client_key) + || keys + .server_keys + .iter() + .any(|server_key| self.protected.contains(&server_key.key)) + } + /// Point-in-time snapshot of every environment's keys, ordered by /// client key so callers iterate deterministically. pub fn snapshot(&self) -> Vec> { @@ -243,6 +333,125 @@ mod tests { assert_eq!(client_keys, vec!["client_a", "client_b"]); } + fn environment(client: &str, server: &str) -> EnvironmentKeys { + EnvironmentKeys { + client_key: client.to_string(), + server_keys: vec![server_key(server)], + } + } + + #[test] + fn sync_to_inserts_new_and_removes_absent_environments() { + // Given an index serving env_a while the config now says env_b + let index = EnvironmentIndex::default(); + index.insert(environment("client_a", "ser.a")); + + // When + let result = index.sync_to(vec![environment("client_b", "ser.b")]); + + // Then + assert_eq!(result.changed, 1); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].client_key, "client_a"); + assert!(index.resolve("client_a").is_none()); + assert!(index.resolve("ser.a").is_none()); + assert!(index.resolve("client_b").is_some()); + assert!(index.resolve("ser.b").is_some()); + } + + #[test] + fn sync_to_never_touches_protected_environments() { + // Given a statically configured environment the config omits — + // and also claims with different keys + let index = EnvironmentIndex::from_settings(&[pair("client_a", "ser.a")]); + + // When the config omits it entirely + let result = index.sync_to(vec![]); + + // Then it survives + assert_eq!(result.removed.len(), 0); + assert!(index.resolve("ser.a").is_some()); + + // When the config claims it with a different server key + let result = index.sync_to(vec![environment("client_a", "ser.other")]); + + // Then the static pairing wins + assert_eq!(result.changed, 0); + assert!(index.resolve("ser.a").is_some()); + assert!(index.resolve("ser.other").is_none()); + + // When the config claims a different environment reusing the + // static environment's server key + let result = index.sync_to(vec![EnvironmentKeys { + client_key: "client_b".to_string(), + server_keys: vec![server_key("ser.a")], + }]); + + // Then it is skipped entirely rather than hijacking the key + assert_eq!(result.changed, 0); + assert!(index.resolve("client_b").is_none()); + assert_eq!(index.resolve("ser.a").unwrap().client_key, "client_a"); + } + + #[test] + fn sync_to_rotation_returns_the_replaced_version() { + // Given + let index = EnvironmentIndex::default(); + index.insert(environment("client_a", "ser.old")); + + // When the config rotates the server key + let result = index.sync_to(vec![environment("client_a", "ser.new")]); + + // Then the replaced version comes back for cache invalidation + assert_eq!(result.changed, 1); + assert_eq!(result.replaced[0].server_keys[0].key, "ser.old"); + assert!(index.resolve("ser.old").is_none()); + assert!(index.resolve("ser.new").is_some()); + } + + #[test] + fn sync_to_leaves_unchanged_environments_alone() { + // Given + let index = EnvironmentIndex::default(); + index.insert(environment("client_a", "ser.a")); + let before = index.resolve("client_a").unwrap(); + + // When the config reports the same keys + let result = index.sync_to(vec![environment("client_a", "ser.a")]); + + // Then nothing changed — not even the Arc identity + assert_eq!(result.changed, 0); + assert!(result.replaced.is_empty()); + assert!(result.removed.is_empty()); + assert!(Arc::ptr_eq(&before, &index.resolve("client_a").unwrap())); + } + + #[test] + fn resolve_rejects_invalid_server_keys_but_keeps_the_client_key() { + // Given an environment whose server keys are deactivated or expired + let index = EnvironmentIndex::default(); + index.insert(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)), + }, + ], + }); + + // Then the invalid keys stop authenticating, the client key doesn't + assert!(index.resolve("client_a").is_some()); + assert!(index.resolve("ser.inactive").is_none()); + assert!(index.resolve("ser.expired").is_none()); + } + #[test] fn valid_server_key_skips_inactive_and_expired_keys() { // Given diff --git a/src/main.rs b/src/main.rs index 0a3a060..405c5fd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,26 +19,31 @@ async fn main() -> anyhow::Result<()> { settings.api_poll_frequency_seconds ); - // serde defaults environment_key_pairs, so a typo'd field name parses - // as an empty set and the proxy would report healthy while rejecting - // every request — make that state loud. - if settings.environment_key_pairs.is_empty() { + // serde defaults these fields, so a typo'd field name parses as an + // empty set and the proxy would report healthy while rejecting every + // request — make that state loud. + if settings.environment_key_pairs.is_empty() && settings.proxy_key.is_none() { warn!( - "No environments configured: environment_key_pairs is empty or \ - missing, so every request will be rejected with 401" + "No environments configured: environment_key_pairs and proxy_key \ + are both empty or missing, so every request will be rejected \ + with 401" ); } let (app, environment_service) = create_router(settings.clone()); + // Refreshes must never overlap: a delayed older poll finishing after a + // newer one could restore removed environments or rotated keys. The + // poll loop is serial, so it just has to start after the initial + // refresh completes. + info!("Loading initial environment data..."); + environment_service.refresh_environment_caches().await; + let polling_service = environment_service.clone(); tokio::spawn(async move { polling_service.poll_environments().await; }); - info!("Loading initial environment data..."); - environment_service.refresh_environment_caches().await; - let addr = SocketAddr::from(( settings .server diff --git a/src/models/mod.rs b/src/models/mod.rs index 1da49b5..0c90610 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,6 +1,8 @@ pub mod engine; +pub mod proxy_config; pub mod request; pub mod response; +pub use proxy_config::ProxyConfigEnvironment; pub use request::{IdentityWithTraits, TraitModel}; pub use response::{APIFeature, APIFeatureState, IdentityResponse}; diff --git a/src/models/proxy_config.rs b/src/models/proxy_config.rs new file mode 100644 index 0000000..36e6df3 --- /dev/null +++ b/src/models/proxy_config.rs @@ -0,0 +1,43 @@ +use chrono::{DateTime, Utc}; +use serde::Deserialize; + +use crate::environments::{EnvironmentKeys, ServerKey}; + +/// An environment as the proxy config endpoint reports it. The response +/// carries more fields (id, name, project/organisation ids, updated_at); +/// only what the proxy acts on is declared, serde ignores the rest. +#[derive(Debug, Clone, Deserialize)] +pub struct ProxyConfigEnvironment { + pub client_side_key: String, + #[serde(default)] + pub server_side_keys: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct ProxyConfigServerKey { + pub key: String, + pub active: bool, + pub expires_at: Option>, +} + +impl From for EnvironmentKeys { + fn from(environment: ProxyConfigEnvironment) -> Self { + let mut server_keys: Vec = environment + .server_side_keys + .into_iter() + .map(|server_key| ServerKey { + key: server_key.key, + active: server_key.active, + expires_at: server_key.expires_at, + }) + .collect(); + // The endpoint does not guarantee key order; sort so a reordered + // response is not mistaken for a changed environment. + server_keys.sort_by(|a, b| a.key.cmp(&b.key)); + + Self { + client_key: environment.client_side_key, + server_keys, + } + } +} diff --git a/src/services/environment.rs b/src/services/environment.rs index 0ed6415..8d53ebe 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -2,7 +2,9 @@ use crate::cache::{CacheKey, EndpointCache, EnvironmentsCache, LocalMemEnvironme use crate::config::settings::AppSettings; use crate::environments::{EnvironmentIndex, EnvironmentKeys}; use crate::error::{EdgeProxyError, Result}; -use crate::models::{APIFeatureState, IdentityResponse, IdentityWithTraits}; +use crate::models::{ + APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, +}; use crate::services::feature_utils::filter_out_server_key_only_flag_results; use chrono::{DateTime, Utc}; use flagsmith_flag_engine::engine::get_evaluation_result; @@ -60,7 +62,9 @@ impl EnvironmentService { } pub async fn refresh_environment_caches(&self) -> bool { - let mut all_success = true; + // Sync first so an environment added to the proxy config gets its + // document fetched in the same pass. + let mut all_success = self.sync_proxy_config().await; for keys in self.environments.snapshot() { match self.fetch_environment(&keys).await { @@ -103,18 +107,90 @@ impl EnvironmentService { self.endpoint_cache.clear_environment(environment_key).await; return; }; - self.cache.remove_environment(&keys.client_key).await; - self.clear_endpoint_caches(&keys).await; + self.purge_environment_caches(&keys).await; info!("Environment removed for key: {}", keys.client_key); } + /// Stopping service of an environment means both cache layers must + /// forget it: the document/context cache under its client key, and + /// the endpoint caches under every key it can be presented by. + async fn purge_environment_caches(&self, keys: &EnvironmentKeys) { + self.cache.remove_environment(&keys.client_key).await; + self.clear_endpoint_caches(keys).await; + } + + /// Bring the served environments in line with the proxy config, when + /// one is configured. Returns false when the fetch fails, keeping the + /// current set untouched — an outage or a rejected proxy key can never + /// wipe the proxy. An environment is removed only when a successful + /// fetch no longer lists it. + async fn sync_proxy_config(&self) -> bool { + let Some(proxy_key) = &self.settings.proxy_key else { + return true; + }; + + let config = match self.fetch_proxy_config(proxy_key).await { + Ok(config) => config, + Err(e) => { + error!("Failed to fetch proxy config: {}", e); + return false; + } + }; + + let desired: Vec = config + .into_iter() + .map(EnvironmentKeys::from) + .filter(|keys| { + let usable = keys.valid_server_key().is_some(); + if !usable { + debug!( + "Skipping proxy config environment {}: no usable server-side key", + keys.client_key + ); + } + usable + }) + .collect(); + let result = self.environments.sync_to(desired); + + for keys in &result.removed { + self.purge_environment_caches(keys).await; + info!("Environment removed from proxy config: {}", keys.client_key); + } + // A replaced version's keys may no longer resolve (rotation); + // clear everything cached under them — current keys repopulate. + for keys in &result.replaced { + self.clear_endpoint_caches(keys).await; + } + if result.changed > 0 || !result.removed.is_empty() { + info!( + "Proxy config applied: {} changed, {} removed", + result.changed, + result.removed.len() + ); + } + + true + } + + async fn fetch_proxy_config(&self, proxy_key: &str) -> Result> { + let url = format!("{}/proxy/config/", self.settings.api_url); + let response = self + .client + .get(&url) + .header("X-Proxy-Key", proxy_key) + .send() + .await?; + response.error_for_status_ref()?; + Ok(response.json().await?) + } + /// The poll-loop counterpart: the loop iterates a snapshot, so a /// removal completing mid-fetch would have its environment's data /// re-inserted by put_environment and pinned until restart. async fn discard_environment_if_removed(&self, keys: &EnvironmentKeys) { if self.environments.resolve(&keys.client_key).is_none() { - self.cache.remove_environment(&keys.client_key).await; - self.clear_endpoint_caches(keys).await; + self.purge_environment_caches(keys).await; } } @@ -136,9 +212,11 @@ impl EnvironmentService { } async fn fetch_environment(&self, keys: &EnvironmentKeys) -> Result { - // An environment stuck here fails every poll, keeping /health red - // until a config change — or, once reconciliation exists, until it - // removes or re-keys the environment. + // The proxy config sync skips environments with no usable key, so + // only a statically configured environment can be stuck here (or, + // for at most one tick, one whose last key expired mid-pass). It + // fails every poll, keeping /health red until the config file is + // fixed — deliberately: that is the misconfiguration signal. let server_key = keys.valid_server_key().ok_or_else(|| { EdgeProxyError::ServiceUnavailable(format!( "no active server-side key for environment {}", @@ -253,6 +331,10 @@ impl EnvironmentService { /// Get pre-serialized environment document bytes with endpoint caching pub async fn get_environment_bytes(&self, environment_key: &str) -> Result> { + // Gate before the cache lookup: a server key that expired since the + // last poll must not keep reading cached responses. + self.resolve_key(environment_key)?; + if self.endpoint_cache.is_environment_document_cache_enabled() { let cache_key = CacheKey::new( environment_key.to_string(), @@ -305,6 +387,13 @@ impl EnvironmentService { environment_key: &str, feature_name: Option<&str>, ) -> Result> { + // Validation only — lookups below still use the raw presented key, so + // a server-side key 503s on this endpoint (contexts are stored under + // the client key). Must run before the cache lookup: a server key + // that expired since the last poll must not keep reading cached + // responses. + self.resolve_key(environment_key)?; + if self.endpoint_cache.is_flags_cache_enabled() { let cache_key = CacheKey::new( environment_key.to_string(), @@ -319,11 +408,6 @@ impl EnvironmentService { } } - // Validation only — lookups below still use the raw presented key, so - // a server-side key 503s on this endpoint (contexts are stored under - // the client key) - self.resolve_key(environment_key)?; - let context = self .cache .get_context(environment_key) @@ -381,6 +465,11 @@ impl EnvironmentService { identity: &IdentityWithTraits, environment_key: &str, ) -> Result { + // Validation only — same server-side-key caveat as + // get_flags_response_data, and before the cache lookup for the same + // expired-key reason. + self.resolve_key(environment_key)?; + if self.endpoint_cache.is_identities_cache_enabled() { // Create cache key from identity data let cache_params = @@ -398,9 +487,6 @@ impl EnvironmentService { } } - // Validation only — same server-side-key caveat as get_flags_response_data - self.resolve_key(environment_key)?; - // Get pre-computed context from cache let context = self .cache @@ -461,6 +547,9 @@ impl EnvironmentService { let mut interval = tokio::time::interval(Duration::from_secs( self.settings.api_poll_frequency_seconds, )); + // The first tick completes immediately; the caller has already done + // the initial refresh, so consume it to poll one full period later. + interval.tick().await; loop { interval.tick().await; @@ -615,6 +704,83 @@ mod tests { assert!(service.cache.get_environment("client").await.is_some()); } + #[tokio::test] + async fn expired_server_key_cannot_read_primed_endpoint_caches() { + // Given every endpoint cache primed under a server key that has + // since expired without any config change (so nothing cleared them) + let settings = AppSettings { + endpoint_caches: EndpointCachesSettings { + flags: EndpointCacheSettings { + use_cache: true, + cache_max_size: 10, + }, + identities: EndpointCacheSettings { + use_cache: true, + cache_max_size: 10, + }, + environment_document: EndpointCacheSettings { + use_cache: true, + cache_max_size: 10, + }, + }, + ..AppSettings::default() + }; + let service = EnvironmentService::new(settings); + service.environments.insert(EnvironmentKeys { + client_key: "client".to_string(), + server_keys: vec![crate::environments::ServerKey { + key: "ser.expired".to_string(), + active: true, + expires_at: Some(Utc::now() - chrono::TimeDelta::seconds(1)), + }], + }); + + let identity = crate::models::IdentityWithTraits::new("id".to_string()); + let identity_params = serde_json::to_string(&identity).unwrap(); + service + .endpoint_cache + .put_flags( + CacheKey::new("ser.expired".into(), "flags".into(), "".into()), + serde_json::json!([]), + ) + .await; + service + .endpoint_cache + .put_identity( + CacheKey::new("ser.expired".into(), "identities".into(), identity_params), + serde_json::json!({"flags": [], "traits": []}), + ) + .await; + service + .endpoint_cache + .put_environment_document( + CacheKey::new( + "ser.expired".into(), + "environment_document".into(), + "".into(), + ), + b"{}".to_vec().into(), + ) + .await; + + // When / Then: every endpoint rejects the expired key instead of + // serving the cached response + assert!(matches!( + service.get_flags_response_data("ser.expired", None).await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); + assert!(matches!( + service + .get_identity_response_data(&identity, "ser.expired") + .await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); + assert!(matches!( + service.get_environment_bytes("ser.expired").await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); + } + fn link(value: &str) -> HeaderMap { let mut h = HeaderMap::new(); h.insert(LINK, HeaderValue::from_str(value).unwrap()); diff --git a/tests/test_proxy_config_discovery.rs b/tests/test_proxy_config_discovery.rs new file mode 100644 index 0000000..b2df11c --- /dev/null +++ b/tests/test_proxy_config_discovery.rs @@ -0,0 +1,344 @@ +use edge_proxy::config::settings::{ + AppSettings, EndpointCacheSettings, EndpointCachesSettings, EnvironmentKeyPair, +}; +use edge_proxy::error::EdgeProxyError; +use edge_proxy::services::EnvironmentService; +use serde_json::{Value, json}; +use wiremock::matchers::{header, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const PROXY_KEY: &str = "pk.test_proxy_key"; +const CLIENT_KEY: &str = "config_client_key"; +const SERVER_KEY: &str = "ser.config_key"; + +fn settings(api_url: &str, pairs: Vec) -> AppSettings { + AppSettings { + environment_key_pairs: pairs, + proxy_key: Some(PROXY_KEY.to_string()), + api_url: api_url.to_string(), + endpoint_caches: EndpointCachesSettings { + flags: EndpointCacheSettings { + use_cache: true, + cache_max_size: 10, + }, + environment_document: EndpointCacheSettings { + use_cache: true, + cache_max_size: 10, + }, + ..Default::default() + }, + ..AppSettings::default() + } +} + +/// The frozen contract shape, extra fields included, so the tests prove +/// serde tolerates everything the endpoint actually sends. +fn config_body(environments: &[(&str, &str)]) -> Value { + Value::Array( + environments + .iter() + .map(|(client_key, server_key)| { + json!({ + "id": 30, + "name": "Test Environment", + "client_side_key": client_key, + "server_side_keys": [ + {"key": server_key, "active": true, "expires_at": null} + ], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }) + }) + .collect(), + ) +} + +fn document_body(client_key: &str) -> Value { + json!({ + "id": 1, + "api_key": client_key, + "name": "Test", + "updated_at": "2026-08-22T00:00:00Z", + "allow_client_traits": true, + "hide_sensitive_data": false, + "hide_disabled_flags": null, + "use_identity_composite_key_for_hashing": true, + "use_identity_overrides_in_local_eval": true, + "project": { + "id": 1, + "name": "project-1", + "hide_disabled_flags": false, + "segments": [], + "server_key_only_feature_ids": [], + "organisation": { + "id": 1, + "name": "org-1", + "feature_analytics": false, + "persist_trait_data": true, + "stop_serving_flags": false, + }, + }, + "feature_states": [ + { + "multivariate_feature_state_values": [], + "feature_state_value": "config_value", + "feature": {"id": 1, "name": "config_flag", "type": "STANDARD"}, + "enabled": true, + "featurestate_uuid": "fs-uuid-1", + } + ], + "identity_overrides": [], + }) +} + +async fn mount_config(mock_server: &MockServer, body: Value, up_to: Option) { + let mut mock = Mock::given(method("GET")) + .and(path("/proxy/config/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(200).set_body_json(body)); + if let Some(n) = up_to { + mock = mock.up_to_n_times(n); + } + mock.mount(mock_server).await; +} + +async fn mount_document(mock_server: &MockServer, server_key: &str, client_key: &str) { + Mock::given(method("GET")) + .and(path("/environment-document/")) + .and(header("X-Environment-Key", server_key)) + .respond_with(ResponseTemplate::new(200).set_body_json(document_body(client_key))) + .mount(mock_server) + .await; +} + +#[tokio::test] +async fn test_environment_from_proxy_config_is_served_after_one_refresh() { + // Given a proxy configured with only a proxy key + let mock_server = MockServer::start().await; + mount_config(&mock_server, config_body(&[(CLIENT_KEY, SERVER_KEY)]), None).await; + mount_document(&mock_server, SERVER_KEY, CLIENT_KEY).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), vec![])); + + // When + let all_success = service.refresh_environment_caches().await; + + // Then both keys serve, document and flags included + assert!(all_success); + assert!(service.get_environment(CLIENT_KEY).await.is_ok()); + assert!(service.get_environment(SERVER_KEY).await.is_ok()); + let flags = service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + assert_eq!(flags.len(), 1); +} + +#[tokio::test] +async fn test_environment_dropped_from_proxy_config_is_removed_and_uncached() { + // Given a served environment with a primed flags endpoint cache + let mock_server = MockServer::start().await; + mount_config( + &mock_server, + config_body(&[(CLIENT_KEY, SERVER_KEY)]), + Some(1), + ) + .await; + mount_config(&mock_server, config_body(&[]), None).await; + mount_document(&mock_server, SERVER_KEY, CLIENT_KEY).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), vec![])); + service.refresh_environment_caches().await; + assert!( + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .is_ok() + ); + + // When the next config response omits the environment + let all_success = service.refresh_environment_caches().await; + + // Then it is gone — including the primed cache entry, which is + // consulted before the key gate + assert!(all_success); + assert!(matches!( + service.get_flags_response_data(CLIENT_KEY, None).await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); + assert!(matches!( + service.get_environment(SERVER_KEY).await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); +} + +#[tokio::test] +async fn test_static_environment_absent_from_proxy_config_is_never_removed() { + // Given a statically configured environment the config never mentions + let mock_server = MockServer::start().await; + mount_config(&mock_server, config_body(&[]), None).await; + mount_document(&mock_server, "ser.static_key", "static_client").await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + vec![EnvironmentKeyPair { + client_side_key: "static_client".to_string(), + server_side_key: "ser.static_key".to_string(), + }], + )); + + // When + let all_success = service.refresh_environment_caches().await; + + // Then the static environment keeps being served + assert!(all_success); + assert!(service.get_environment("static_client").await.is_ok()); + assert!(service.get_environment("ser.static_key").await.is_ok()); +} + +#[tokio::test] +async fn test_proxy_config_fetch_failure_removes_nothing() { + // Given a served environment and a proxy config endpoint that starts + // failing + let mock_server = MockServer::start().await; + mount_config( + &mock_server, + config_body(&[(CLIENT_KEY, SERVER_KEY)]), + Some(1), + ) + .await; + Mock::given(method("GET")) + .and(path("/proxy/config/")) + .respond_with(ResponseTemplate::new(500)) + .mount(&mock_server) + .await; + mount_document(&mock_server, SERVER_KEY, CLIENT_KEY).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), vec![])); + assert!(service.refresh_environment_caches().await); + + // When the next sync fails + let all_success = service.refresh_environment_caches().await; + + // Then the failure is reported but nothing is removed + assert!(!all_success); + assert!(service.get_environment(CLIENT_KEY).await.is_ok()); + assert!(service.get_environment(SERVER_KEY).await.is_ok()); +} + +#[tokio::test] +async fn test_server_key_rotation_stops_serving_the_old_key() { + // Given a served environment whose document bytes are cached under + // the old server key + let mock_server = MockServer::start().await; + mount_config( + &mock_server, + config_body(&[(CLIENT_KEY, "ser.old_key")]), + Some(1), + ) + .await; + mount_config( + &mock_server, + config_body(&[(CLIENT_KEY, "ser.new_key")]), + None, + ) + .await; + mount_document(&mock_server, "ser.old_key", CLIENT_KEY).await; + mount_document(&mock_server, "ser.new_key", CLIENT_KEY).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), vec![])); + service.refresh_environment_caches().await; + assert!(service.get_environment_bytes("ser.old_key").await.is_ok()); + + // When the config rotates the server key + let all_success = service.refresh_environment_caches().await; + + // Then the old key stops serving — cached bytes included — and the + // new key works + assert!(all_success); + assert!(matches!( + service.get_environment_bytes("ser.old_key").await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); + assert!(service.get_environment_bytes("ser.new_key").await.is_ok()); + assert!(service.get_environment(CLIENT_KEY).await.is_ok()); +} + +#[tokio::test] +async fn test_environment_without_usable_keys_is_skipped_not_failed() { + // Given the config reports a healthy environment and a brand-new one + // with no server-side keys yet + let mock_server = MockServer::start().await; + let body = json!([ + { + "id": 30, + "name": "healthy", + "client_side_key": CLIENT_KEY, + "server_side_keys": [{"key": SERVER_KEY, "active": true, "expires_at": null}], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }, + { + "id": 31, + "name": "no keys yet", + "client_side_key": "keyless_client", + "server_side_keys": [], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }, + ]); + mount_config(&mock_server, body, None).await; + mount_document(&mock_server, SERVER_KEY, CLIENT_KEY).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), vec![])); + + // When + let all_success = service.refresh_environment_caches().await; + + // Then the poll succeeds — /health stays green — the healthy + // environment serves, and the key-less one is simply not indexed + assert!(all_success); + assert!(service.get_environment(CLIENT_KEY).await.is_ok()); + assert!(matches!( + service.get_environment("keyless_client").await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); +} + +#[tokio::test] +async fn test_environment_whose_only_key_is_deactivated_is_dropped() { + // Given a served environment whose only key the config then reports + // as inactive + let mock_server = MockServer::start().await; + mount_config( + &mock_server, + config_body(&[(CLIENT_KEY, SERVER_KEY)]), + Some(1), + ) + .await; + let deactivated = json!([{ + "id": 30, + "name": "Test", + "client_side_key": CLIENT_KEY, + "server_side_keys": [{"key": SERVER_KEY, "active": false, "expires_at": null}], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }]); + mount_config(&mock_server, deactivated, None).await; + mount_document(&mock_server, SERVER_KEY, CLIENT_KEY).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), vec![])); + service.refresh_environment_caches().await; + assert!(service.get_environment(SERVER_KEY).await.is_ok()); + + // When + let all_success = service.refresh_environment_caches().await; + + // Then the environment is dropped entirely and the poll stays healthy + assert!(all_success); + assert!(matches!( + service.get_environment(SERVER_KEY).await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); + assert!(matches!( + service.get_environment(CLIENT_KEY).await, + Err(EdgeProxyError::FlagsmithUnknownKey(_)) + )); +}