diff --git a/src/config/settings.rs b/src/config/settings.rs index a812d96..929ce32 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -137,6 +137,10 @@ pub struct AppSettings { pub api_poll_frequency_seconds: u64, #[serde(default = "default_api_poll_timeout")] pub api_poll_timeout_seconds: u64, + // A zero interval would panic the flush task silently. + #[serde(default = "default_usage_flush_interval")] + #[validate(range(min = 1))] + pub usage_flush_interval_seconds: u64, #[serde(default)] pub server: ServerSettings, #[serde(default)] @@ -159,6 +163,10 @@ fn default_api_poll_timeout() -> u64 { 5 } +fn default_usage_flush_interval() -> u64 { + 60 +} + impl Default for AppSettings { fn default() -> Self { Self { @@ -167,6 +175,7 @@ impl Default for AppSettings { api_url: default_api_url(), api_poll_frequency_seconds: default_api_poll_frequency(), api_poll_timeout_seconds: default_api_poll_timeout(), + usage_flush_interval_seconds: default_usage_flush_interval(), server: ServerSettings::default(), logging: LoggingSettings::default(), endpoint_caches: EndpointCachesSettings::default(), @@ -221,6 +230,16 @@ mod tests { assert!(settings.validate().is_err()); } + #[test] + fn test_config_with_zero_usage_flush_interval_is_invalid() { + // Given an interval that would panic the flush task + let settings: AppSettings = + serde_json::from_str(r#"{"usage_flush_interval_seconds": 0}"#).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 diff --git a/src/environments.rs b/src/environments.rs index cb2d61e..383378e 100644 --- a/src/environments.rs +++ b/src/environments.rs @@ -90,6 +90,11 @@ impl EnvironmentIndex { index } + /// Whether the key belongs to a statically configured environment. + pub fn is_static(&self, key: &str) -> bool { + self.protected.contains(key) + } + /// Resolve a presented key — client- or server-side — to its /// environment's keys. A server-side key resolves only while it is /// valid, so a deactivation delivered by the proxy config and an diff --git a/src/lib.rs b/src/lib.rs index 2bfc72a..14cb2b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,3 +6,4 @@ pub mod models; pub mod routes; pub mod services; pub mod state; +pub mod usage; diff --git a/src/main.rs b/src/main.rs index 405c5fd..38b1930 100644 --- a/src/main.rs +++ b/src/main.rs @@ -44,6 +44,11 @@ async fn main() -> anyhow::Result<()> { polling_service.poll_environments().await; }); + let usage_service = environment_service.clone(); + tokio::spawn(async move { + usage_service.flush_usage_periodically().await; + }); + let addr = SocketAddr::from(( settings .server diff --git a/src/services/environment.rs b/src/services/environment.rs index 8d53ebe..d92884c 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -6,6 +6,7 @@ use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, }; use crate::services::feature_utils::filter_out_server_key_only_flag_results; +use crate::usage::{Resource, UsageCounts, UsageRow}; use chrono::{DateTime, Utc}; use flagsmith_flag_engine::engine::get_evaluation_result; use flagsmith_flag_engine::engine_eval::{FlagResult, add_identity_to_context}; @@ -24,6 +25,7 @@ pub struct EnvironmentService { pub settings: AppSettings, pub last_updated_at: Arc>>>, environments: EnvironmentIndex, + usage: UsageCounts, } impl EnvironmentService { @@ -52,6 +54,7 @@ impl EnvironmentService { settings, last_updated_at: Arc::new(RwLock::new(None)), environments, + usage: UsageCounts::default(), } } @@ -205,10 +208,26 @@ impl EnvironmentService { } } - fn resolve_key(&self, environment_key: &str) -> Result> { - self.environments + /// Resolve a presented key, counting the request for usage reporting. + /// Every SDK entry point resolves through here, so a served request + /// cannot be missed. + fn resolve_key( + &self, + environment_key: &str, + resource: Resource, + ) -> Result> { + let keys = self + .environments .resolve(environment_key) - .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string())) + .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; + self.track_usage(&keys.client_key, resource); + Ok(keys) + } + + fn track_usage(&self, client_key: &str, resource: Resource) { + if self.settings.proxy_key.is_some() && !self.environments.is_static(client_key) { + self.usage.increment(client_key, resource); + } } async fn fetch_environment(&self, keys: &EnvironmentKeys) -> Result { @@ -267,6 +286,14 @@ impl EnvironmentService { .client .get(&next_url) .header("X-Environment-Key", server_side_key); + // Core excludes marked fetches from API usage — the proxy + // reports served requests instead. Static environments stay + // unmarked and keep their old billing. + if !self.environments.is_static(server_side_key) { + if let Some(proxy_key) = &self.settings.proxy_key { + request = request.header("X-Proxy-Key", proxy_key); + } + } // If-Modified-Since is meaningful only on the first request; the // upstream pagination cursor (page_id) drives subsequent fetches. if document.is_none() { @@ -320,7 +347,11 @@ impl EnvironmentService { } pub async fn get_environment(&self, environment_key: &str) -> Result> { - let keys = self.resolve_key(environment_key)?; + // Lookup, not an SDK entry point: callers count via resolve_key. + let keys = self + .environments + .resolve(environment_key) + .ok_or_else(|| EdgeProxyError::FlagsmithUnknownKey(environment_key.to_string()))?; // Documents are cached under the client key, whichever key was presented self.cache @@ -333,7 +364,7 @@ impl EnvironmentService { 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)?; + self.resolve_key(environment_key, Resource::EnvironmentDocument)?; if self.endpoint_cache.is_environment_document_cache_enabled() { let cache_key = CacheKey::new( @@ -392,7 +423,7 @@ impl EnvironmentService { // 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)?; + self.resolve_key(environment_key, Resource::Flags)?; if self.endpoint_cache.is_flags_cache_enabled() { let cache_key = CacheKey::new( @@ -468,7 +499,7 @@ impl EnvironmentService { // 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)?; + self.resolve_key(environment_key, Resource::Identities)?; if self.endpoint_cache.is_identities_cache_enabled() { // Create cache key from identity data @@ -557,6 +588,78 @@ impl EnvironmentService { self.refresh_environment_caches().await; } } + + /// The usage endpoint's batch cap — MAX_USAGE_ROWS in the edge_proxy + /// app. Flushes are chunked to it so a large environment set can + /// never be rejected outright. + const MAX_ROWS_PER_FLUSH: usize = 1000; + + /// Report the counts accumulated since the last flush to the usage + /// endpoint, in chunks the server accepts. A rejected (4xx) chunk is + /// dropped — retrying cannot heal a rejection, and losing one window + /// beats resending a poisoned batch forever. Any other failure keeps + /// the chunk for the next flush. Returns false when any chunk was + /// not accepted. + pub async fn flush_usage(&self) -> bool { + let Some(proxy_key) = &self.settings.proxy_key else { + return true; + }; + let mut rows = self.usage.drain(); + let url = format!("{}/proxy/usage/", self.settings.api_url); + let mut all_success = true; + + while !rows.is_empty() { + let chunk: Vec = rows + .drain(..rows.len().min(Self::MAX_ROWS_PER_FLUSH)) + .collect(); + let result = self + .client + .post(&url) + .header("X-Proxy-Key", proxy_key) + .json(&chunk) + .send() + .await; + match result { + Ok(response) if response.status().is_success() => {} + Ok(response) if response.status().is_client_error() => { + error!( + "Usage report rejected with {}: dropping {} rows", + response.status(), + chunk.len() + ); + all_success = false; + } + Ok(response) => { + error!("Failed to report usage: {}", response.status()); + self.usage.merge(chunk); + all_success = false; + } + Err(e) => { + error!("Failed to report usage: {}", e); + self.usage.merge(chunk); + all_success = false; + } + } + } + + all_success + } + + pub async fn flush_usage_periodically(self: Arc) { + if self.settings.proxy_key.is_none() { + return; + } + let mut interval = tokio::time::interval(Duration::from_secs( + self.settings.usage_flush_interval_seconds, + )); + // The first tick completes immediately, before anything is counted. + interval.tick().await; + + loop { + interval.tick().await; + self.flush_usage().await; + } + } } /// Format the cached document's `updated_at` as an RFC 2822 `If-Modified-Since` diff --git a/src/usage.rs b/src/usage.rs new file mode 100644 index 0000000..c5c365e --- /dev/null +++ b/src/usage.rs @@ -0,0 +1,150 @@ +use serde::Serialize; +use std::collections::HashMap; +use std::sync::Mutex; + +/// An SDK endpoint as the usage endpoint names it. +#[derive(Serialize, Clone, Copy, PartialEq, Eq, Hash, Debug)] +#[serde(rename_all = "kebab-case")] +pub enum Resource { + Flags, + Identities, + EnvironmentDocument, +} + +/// One row of the `POST /proxy/usage/` body. +#[derive(Serialize, Debug, PartialEq)] +pub struct UsageRow { + pub client_side_key: String, + pub resource: Resource, + pub count: u64, +} + +/// Requests served since the last flush, aggregated per environment and +/// resource. Keyed by resolved client keys only, so the map is bounded by +/// the served environment set. +/// +/// Uses `std::sync::Mutex`: every operation writes, and guards are held +/// only for a map operation, never across an await. +#[derive(Default)] +pub struct UsageCounts { + by_key: Mutex>, +} + +impl UsageCounts { + pub fn increment(&self, client_key: &str, resource: Resource) { + let mut by_key = self.by_key.lock().expect("usage counts lock poisoned"); + *by_key + .entry((client_key.to_string(), resource)) + .or_default() += 1; + } + + /// Take everything counted so far, leaving the map empty. + pub fn drain(&self) -> Vec { + let by_key = std::mem::take(&mut *self.by_key.lock().expect("usage counts lock poisoned")); + by_key + .into_iter() + .map(|((client_side_key, resource), count)| UsageRow { + client_side_key, + resource, + count, + }) + .collect() + } + + /// Add drained counts back, on top of anything counted since — for + /// when a flush fails and the rows must survive until the next one. + pub fn merge(&self, counts: Vec) { + let mut by_key = self.by_key.lock().expect("usage counts lock poisoned"); + for row in counts { + let entry = by_key + .entry((row.client_side_key, row.resource)) + .or_default(); + *entry = entry.saturating_add(row.count); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn increment_aggregates_by_key_and_resource() { + // Given + let counts = UsageCounts::default(); + + // When + counts.increment("client", Resource::Flags); + counts.increment("client", Resource::Flags); + counts.increment("client", Resource::Identities); + + // Then + let mut rows = counts.drain(); + rows.sort_by_key(|row| format!("{:?}", row.resource)); + assert_eq!( + rows, + vec![ + UsageRow { + client_side_key: "client".to_string(), + resource: Resource::Flags, + count: 2, + }, + UsageRow { + client_side_key: "client".to_string(), + resource: Resource::Identities, + count: 1, + }, + ] + ); + } + + #[test] + fn drain_leaves_nothing_behind() { + // Given + let counts = UsageCounts::default(); + counts.increment("client", Resource::Flags); + + // When + counts.drain(); + + // Then + assert!(counts.drain().is_empty()); + } + + #[test] + fn merge_adds_drained_counts_back() { + // Given a drained batch and a request counted since + let counts = UsageCounts::default(); + counts.increment("client", Resource::Flags); + let drained = counts.drain(); + counts.increment("client", Resource::Flags); + + // When + counts.merge(drained); + + // Then + let rows = counts.drain(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].count, 2); + } + + #[test] + fn usage_row_serializes_to_the_contract_shape() { + // Given + let row = UsageRow { + client_side_key: "client".to_string(), + resource: Resource::EnvironmentDocument, + count: 3, + }; + + // Then + assert_eq!( + serde_json::to_value(&row).unwrap(), + serde_json::json!({ + "client_side_key": "client", + "resource": "environment-document", + "count": 3, + }) + ); + } +} diff --git a/tests/test_usage_tracking.rs b/tests/test_usage_tracking.rs new file mode 100644 index 0000000..61a39d1 --- /dev/null +++ b/tests/test_usage_tracking.rs @@ -0,0 +1,423 @@ +use axum_test::TestServer; +use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; +use edge_proxy::routes::create_router; +use edge_proxy::services::EnvironmentService; +use serde_json::{Value, json}; +use wiremock::matchers::{header, method, path}; +use wiremock::{Mock, MockServer, Request, 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, proxy_key: Option<&str>, pairs: Vec) -> AppSettings { + AppSettings { + environment_key_pairs: pairs, + proxy_key: proxy_key.map(str::to_string), + api_url: api_url.to_string(), + ..AppSettings::default() + } +} + +fn config_body() -> Value { + 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, + }]) +} + +fn document_body() -> 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) { + Mock::given(method("GET")) + .and(path("/proxy/config/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(200).set_body_json(config_body())) + .mount(mock_server) + .await; +} + +async fn mount_document(mock_server: &MockServer) { + Mock::given(method("GET")) + .and(path("/environment-document/")) + .respond_with(ResponseTemplate::new(200).set_body_json(document_body())) + .mount(mock_server) + .await; +} + +async fn mount_usage(mock_server: &MockServer, status: u16, up_to: Option) { + let mut mock = Mock::given(method("POST")) + .and(path("/proxy/usage/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(status)); + if let Some(n) = up_to { + mock = mock.up_to_n_times(n); + } + mock.mount(mock_server).await; +} + +async fn requests_to(mock_server: &MockServer, url_path: &str) -> Vec { + mock_server + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() == url_path) + .collect() +} + +/// The usage rows of the request body, sorted by resource for +/// order-independent assertions. +fn usage_rows(request: &Request) -> Vec { + let mut rows: Vec = serde_json::from_slice(&request.body).unwrap(); + rows.sort_by_key(|row| row["resource"].as_str().unwrap().to_string()); + rows +} + +#[tokio::test] +async fn test_served_requests_flush_aggregated_usage() { + // Given a discovered environment served through the full router + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 204, None).await; + let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + let server = TestServer::new(app).unwrap(); + + // When SDK traffic arrives under both keys, then usage is flushed + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .get("/api/v1/identities") + .add_query_param("identifier", "user_1") + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .post("/api/v1/identities") + .json(&json!({"identifier": "user_2"})) + .add_header("X-Environment-Key", CLIENT_KEY) + .await + .assert_status_ok(); + server + .get("/api/v1/environment-document") + .add_header("X-Environment-Key", SERVER_KEY) + .await + .assert_status_ok(); + let flushed = service.flush_usage().await; + + // Then one POST reports everything, keyed by the client key even for + // requests that presented the server key + assert!(flushed); + let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!(posts.len(), 1); + assert_eq!( + usage_rows(&posts[0]), + vec![ + json!({"client_side_key": CLIENT_KEY, "resource": "environment-document", "count": 1}), + json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 2}), + json!({"client_side_key": CLIENT_KEY, "resource": "identities", "count": 2}), + ] + ); +} + +#[tokio::test] +async fn test_unresolved_keys_are_never_counted() { + // Given a proxy serving one environment + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + let (app, service) = create_router(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + let server = TestServer::new(app).unwrap(); + + // When requests present an unknown key or none at all + server + .get("/api/v1/flags") + .add_header("X-Environment-Key", "unknown_key") + .await + .assert_status_unauthorized(); + server + .get("/api/v1/flags") + .await + .assert_status_unauthorized(); + + // Then there is nothing to flush and no request is made + assert!(service.flush_usage().await); + assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); +} + +#[tokio::test] +async fn test_failed_flush_merges_counts_into_the_next() { + // Given a served request and a usage endpoint that fails once + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 500, Some(1)).await; + mount_usage(&mock_server, 204, None).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // When the first flush fails and another request is served + assert!(!service.flush_usage().await); + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // Then the next flush carries both counts — nothing lost, nothing + // double-counted + assert!(service.flush_usage().await); + let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!(posts.len(), 2); + assert_eq!( + usage_rows(&posts[1]), + vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 2})] + ); +} + +#[tokio::test] +async fn test_flush_without_proxy_key_is_inert() { + // Given a statically configured proxy with no proxy key + let mock_server = MockServer::start().await; + mount_document(&mock_server).await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + None, + vec![EnvironmentKeyPair { + client_side_key: CLIENT_KEY.to_string(), + server_side_key: SERVER_KEY.to_string(), + }], + )); + service.refresh_environment_caches().await; + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // When / Then: flushing succeeds without reporting anything + assert!(service.flush_usage().await); + assert!(requests_to(&mock_server, "/proxy/usage/").await.is_empty()); +} + +#[tokio::test] +async fn test_rejected_flush_drops_the_batch_instead_of_retrying_it() { + // Given a served request and a usage endpoint that rejects the batch + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 400, None).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + + // When the flush is rejected + assert!(!service.flush_usage().await); + + // Then the rows are dropped, not resent forever: the next flush has + // nothing to send + assert!(service.flush_usage().await); + assert_eq!(requests_to(&mock_server, "/proxy/usage/").await.len(), 1); +} + +#[tokio::test] +async fn test_flush_chunks_batches_to_the_server_cap() { + // Given served requests for more environments than one batch may hold + let environments: Vec = (0..1001) + .map(|n| { + json!({ + "id": n, + "name": format!("env {n}"), + "client_side_key": format!("client_{n}"), + "server_side_keys": [ + {"key": format!("ser.key_{n}"), "active": true, "expires_at": null} + ], + "updated_at": "2026-08-15T08:57:43.311081Z", + "project_id": 35, + "organisation_id": 82, + }) + }) + .collect(); + let mock_server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/proxy/config/")) + .respond_with(ResponseTemplate::new(200).set_body_json(Value::Array(environments))) + .mount(&mock_server) + .await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 204, None).await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + service.refresh_environment_caches().await; + for n in 0..1001 { + service + .get_flags_response_data(&format!("client_{n}"), None) + .await + .unwrap(); + } + + // When + assert!(service.flush_usage().await); + + // Then the rows arrive split across two accepted requests + let posts = requests_to(&mock_server, "/proxy/usage/").await; + let row_counts: Vec = posts.iter().map(|post| usage_rows(post).len()).collect(); + assert_eq!(row_counts, vec![1000, 1]); +} + +#[tokio::test] +async fn test_static_environment_usage_is_neither_counted_nor_marked() { + // Given a proxy serving a static environment alongside a discovered one + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + mount_document(&mock_server).await; + mount_usage(&mock_server, 204, None).await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + Some(PROXY_KEY), + vec![EnvironmentKeyPair { + client_side_key: "static_client".to_string(), + server_side_key: "ser.static_key".to_string(), + }], + )); + service.refresh_environment_caches().await; + + // When both environments serve a request, and usage is flushed + service + .get_flags_response_data("static_client", None) + .await + .unwrap(); + service + .get_flags_response_data(CLIENT_KEY, None) + .await + .unwrap(); + assert!(service.flush_usage().await); + + // Then the static environment keeps its old billing: its document + // fetch is not marked as the proxy's own, and it is not reported + let fetches = requests_to(&mock_server, "/environment-document/").await; + let marked: Vec = fetches + .iter() + .map(|request| request.headers.contains_key("X-Proxy-Key")) + .collect(); + let static_fetches: Vec<&Request> = fetches + .iter() + .filter(|request| request.headers["X-Environment-Key"] == "ser.static_key") + .collect(); + assert!(!static_fetches.is_empty()); + assert!( + static_fetches + .iter() + .all(|request| !request.headers.contains_key("X-Proxy-Key")) + ); + assert!(marked.contains(&true)); + let posts = requests_to(&mock_server, "/proxy/usage/").await; + assert_eq!( + usage_rows(&posts[0]), + vec![json!({"client_side_key": CLIENT_KEY, "resource": "flags", "count": 1})] + ); +} + +#[tokio::test] +async fn test_document_fetch_carries_the_proxy_key() { + // Given a document endpoint that only answers requests marked as the + // proxy's own + let mock_server = MockServer::start().await; + mount_config(&mock_server).await; + Mock::given(method("GET")) + .and(path("/environment-document/")) + .and(header("X-Proxy-Key", PROXY_KEY)) + .respond_with(ResponseTemplate::new(200).set_body_json(document_body())) + .mount(&mock_server) + .await; + let service = EnvironmentService::new(settings(&mock_server.uri(), Some(PROXY_KEY), vec![])); + + // When / Then + assert!(service.refresh_environment_caches().await); + assert!(service.get_environment(CLIENT_KEY).await.is_ok()); +} + +#[tokio::test] +async fn test_static_document_fetch_omits_the_proxy_key() { + // Given a statically configured proxy with no proxy key + let mock_server = MockServer::start().await; + mount_document(&mock_server).await; + let service = EnvironmentService::new(settings( + &mock_server.uri(), + None, + vec![EnvironmentKeyPair { + client_side_key: CLIENT_KEY.to_string(), + server_side_key: SERVER_KEY.to_string(), + }], + )); + + // When + assert!(service.refresh_environment_caches().await); + + // Then the document fetch is byte-identical to today's + let fetches = requests_to(&mock_server, "/environment-document/").await; + assert!(!fetches.is_empty()); + assert!( + fetches + .iter() + .all(|request| !request.headers.contains_key("X-Proxy-Key")) + ); +}